descriptors.go

// Copyright 2018 Twitch Interactive, Inc.  All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may not
// use this file except in compliance with the License. A copy of the License is
// located at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// or in the "license" file accompanying this file. This file is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

// package descriptors provides tools for manipulating and inspecting protobuf
// descriptors.
package descriptors

import (
	"bytes"
	"compress/gzip"
	"io/ioutil"

	"github.com/golang/protobuf/proto"
	protobuf "github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/pkg/errors"
)

// UnpackFile reads gz as a gzipped, protobuf-encoded FileDescriptorProto. This
// is the format used to store descriptors in protoc-gen-go and
// protoc-gen-twirp.
func UnpackFile(gz []byte) (*protobuf.FileDescriptorProto, error) {
	r, err := gzip.NewReader(bytes.NewReader(gz))
	if err != nil {
		return nil, errors.Wrap(err, "failed to open gzip reader")
	}
	defer r.Close()

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "failed to uncompress descriptor")
	}

	fd := new(protobuf.FileDescriptorProto)
	if err := proto.Unmarshal(b, fd); err != nil {
		return nil, errors.Wrap(err, "malformed FileDescriptorProto")
	}

	return fd, nil
}

// A DescribableMessage provides a gzipped, protobuf-encoded
// FileDescriptorProto, and a series of ints which index into the File to
// provide the address of a DescriptorProto describing a message.
//
// This interface should be fulfilled by any message structs generated by
// protoc-gen-go.
type DescribableMessage interface {
	Descriptor() ([]byte, []int)
}

// MessageDescriptor returns the DescriptorProto describing a message value, and
// the FileDescriptorProto in which the message is defined.
func MessageDescriptor(msg DescribableMessage) (*protobuf.FileDescriptorProto, *protobuf.DescriptorProto, error) {
	gz, path := msg.Descriptor()
	fd, err := UnpackFile(gz)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to unpack gzipped descriptor")
	}

	d, err := MessageInFile(fd, path)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to find message")
	}
	return fd, d, nil
}

// MessageInFile finds a message in fd using the given path as an address.
func MessageInFile(fd *protobuf.FileDescriptorProto, path []int) (*protobuf.DescriptorProto, error) {
	if path[0] > len(fd.MessageType) {
		return nil, errors.Errorf("message index %d out of bounds on file", path[0])
	}
	d := fd.MessageType[path[0]]
	for _, i := range path[1:] {
		if i < 0 || i > len(d.NestedType) {
			return nil, errors.Errorf("nested message index %d out of bounds on type %q", i, d.GetName())
		}
		d = d.NestedType[i]
	}
	return d, nil

}

// A DescribableEnum provides a gzipped, protobuf-encoded FileDescriptorProto,
// and a series of ints which index into the File to provide the address of an
// EnumDescriptorProto describing an enum.
//
// This interface should be fulfilled by any enums generated by protoc-gen-go.
type DescribableEnum interface {
	EnumDescriptor() ([]byte, []int)
}

// EnumDescriptor returns the EnumDescriptorProto describing an enum value, and
// the FileDescriptorProto in which the enum is defined.
func EnumDescriptor(enum DescribableEnum) (*protobuf.FileDescriptorProto, *protobuf.EnumDescriptorProto, error) {
	gz, path := enum.EnumDescriptor()
	fd, err := UnpackFile(gz)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to unpack gzipped descriptor")
	}
	ed, err := EnumInFile(fd, path)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to find enum")
	}
	return fd, ed, nil
}

// EnumInFile uses the given path to find an enum in the given file.
func EnumInFile(fd *protobuf.FileDescriptorProto, path []int) (*protobuf.EnumDescriptorProto, error) {
	if len(path) == 1 {
		// This is an enum declared at the top level of a file.
		if path[0] < 0 || path[0] > len(fd.EnumType) {
			return nil, errors.Errorf("enum index %d out of bounds on file", path[0])
		}
		return fd.EnumType[path[0]], nil
	}

	// This is an enum declared inside a message. We need to find the message, and
	// then the enum within it.
	//
	// The last element of the path will be the index of the enum inside a message
	// descriptor's EnumType slice. Everything before that indexes us through
	// messages.
	msgPath := path[0 : len(path)-1]
	md, err := MessageInFile(fd, msgPath)
	if err != nil {
		return nil, errors.Wrap(err, "unable to find enum inside message")
	}
	enumIdx := path[len(path)-1]
	if enumIdx < 0 || enumIdx > len(md.EnumType) {
		return nil, errors.Errorf("enum index %d out of bounds on message type %q", enumIdx, md.GetName())
	}
	return md.EnumType[enumIdx], nil
}

// A DescribableService provides a gzipped, protobuf-encoded
// FileDescriptorProto, and an int which indexes into the File to provide the
// address of a ServiceDescriptorProto describing a service.
//
// This interface should be fulfilled by any servers generated by
// protoc-gen-twirp.
type DescribableService interface {
	ServiceDescriptor() ([]byte, int)
}

// ServiceDescriptor returns the ServiceDescriptorProto describing a service
// value, and the FileDescriptorProto in which the service is defined.
func ServiceDescriptor(svc DescribableService) (*protobuf.FileDescriptorProto, *protobuf.ServiceDescriptorProto, error) {
	gz, idx := svc.ServiceDescriptor()
	fd, err := UnpackFile(gz)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to unpack gzipped descriptor")
	}

	sd, err := ServiceInFile(fd, idx)
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to find service")
	}

	return fd, sd, nil
}

// ServiceInFile uses the given index to find the service within the given
// FileDescriptorProto.
func ServiceInFile(fd *protobuf.FileDescriptorProto, index int) (*protobuf.ServiceDescriptorProto, error) {
	if index > len(fd.Service) {
		return nil, errors.Errorf("service index %d out of bounds on file", index)
	}
	return fd.Service[index], nil
}