protodump/pkg/protodump/proto.go

package protodump

import (
	"fmt"
	"strconv"
	"strings"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
	"google.golang.org/protobuf/types/descriptorpb"
)

type ProtoDefinition struct {
	builder     strings.Builder
	indendation int
	pb          descriptorpb.FileDescriptorProto
	descriptor  protoreflect.FileDescriptor
	filename    string
}

// TODO add proto2 support

func (pd *ProtoDefinition) indent() {
	pd.indendation += 1
}

func (pd *ProtoDefinition) dedent() {
	pd.indendation -= 1
}

func (pd *ProtoDefinition) writeIndented(s string) {
	pd.builder.WriteString(strings.Repeat("    ", pd.indendation))
	pd.write(s)
}

func (pd *ProtoDefinition) write(s string) {
	pd.builder.WriteString(s)
}

func (pd *ProtoDefinition) String() string {
	return pd.builder.String()
}

func (pd *ProtoDefinition) Filename() string {
	return pd.filename
}

func (pd *ProtoDefinition) writeMethod(method protoreflect.MethodDescriptor) {
	// TODO need to handle method options
	pd.writeIndented("rpc ")
	pd.write(string(method.Name()))
	pd.write("(")
	if method.IsStreamingClient() {
		pd.write("streaming ")
	}
	pd.write(string(method.Input().Name()))
	pd.write(") returns (")
	if method.IsStreamingServer() {
		pd.write("streaming ")
	}
	pd.write(string(method.Output().Name()))
	pd.write(") {}\n")
}

func (pd *ProtoDefinition) writeService(service protoreflect.ServiceDescriptor) {
	// TODO need to handle service options
	pd.write("service ")
	pd.write(string(service.Name()))
	pd.write(" {\n")
	pd.indent()
	for i := 0; i < service.Methods().Len(); i++ {
		pd.writeMethod(service.Methods().Get(i))
	}
	pd.dedent()
	pd.writeIndented("}\n\n")
}

func (pd *ProtoDefinition) writeType(field protoreflect.FieldDescriptor) {
	kind := field.Kind().String()
	if kind == "message" {
		pd.write(string(field.Message().Name()))
	} else if kind == "enum" {
		pd.write(string(field.Enum().Name()))
	} else if kind == "map" {
		pd.write("map<")
		pd.writeType(field.MapKey())
		pd.write(", ")
		pd.writeType(field.MapValue())
		pd.write(">")
	} else {
		pd.write(kind)
	}
}

func (pd *ProtoDefinition) writeOneof(oneof protoreflect.OneofDescriptor) {
	// TODO need to handle oneof options
	pd.writeIndented("")
	pd.write("oneof ")
	pd.write(string(oneof.Name()))
	pd.write(" {\n")
	pd.indent()
	for i := 0; i < oneof.Fields().Len(); i++ {
		pd.writeField(oneof.Fields().Get(i))
	}
	pd.dedent()
	pd.writeIndented("}\n")
}

func (pd *ProtoDefinition) writeField(field protoreflect.FieldDescriptor) {
	// TODO need to handle options
	pd.writeIndented("")
	if field.HasOptionalKeyword() {
		pd.write("optional ")
	} else if field.Cardinality().String() == "repeated" {
		pd.write("repeated ")
	}
	pd.writeType(field)
	pd.write(" ")
	pd.write(string(field.Name()))
	pd.write(" = ")
	pd.write(strconv.Itoa(int(field.Number())))
	pd.write(";\n")
}

func (pd *ProtoDefinition) writeEnum(enum protoreflect.EnumDescriptor) {
	pd.write("enum ")
	pd.write(string(enum.Name()))
	pd.write(" {\n")
	// TODO need to handle enum options (allow_alias)
	pd.indent()
	for i := 0; i < enum.Values().Len(); i++ {
		value := enum.Values().Get(i)
		pd.writeIndented(string(value.Name()))
		pd.write(" = ")
		pd.write(fmt.Sprintf("%d", value.Number()))
		pd.write(";\n")
	}
	pd.dedent()
	pd.writeIndented("}\n\n")
}

func (pd *ProtoDefinition) writeMessage(message protoreflect.MessageDescriptor) {
	// TODO need to handle message options
	pd.write("message ")
	pd.write(string(message.Name()))
	pd.write(" {\n")
	pd.indent()

	for i := 0; i < message.ReservedNames().Len(); i++ {
		name := message.ReservedNames().Get(i)
		pd.writeIndented("reserved ")
		pd.write(string(name))
		pd.write(";\n")
	}

	for i := 0; i < message.ReservedRanges().Len(); i++ {
		reservedRange := message.ReservedRanges().Get(i)
		pd.writeIndented("reserved ")
		if reservedRange[0] == reservedRange[1] {
			pd.write(string(reservedRange[0]))
		} else {
			if reservedRange[0] > reservedRange[1] {
				reservedRange[1], reservedRange[0] = reservedRange[0], reservedRange[1]
			}
			pd.write(fmt.Sprintf("%d", reservedRange[0]))
			pd.write(" to ")
			pd.write(fmt.Sprintf("%d", reservedRange[1]))
		}
		pd.write(";\n")
	}

	for i := 0; i < message.Fields().Len(); i++ {
		field := message.Fields().Get(i)
		if field.ContainingOneof() == nil {
			pd.writeField(field)
		}
	}

	for i := 0; i < message.Oneofs().Len(); i++ {
		pd.writeOneof(message.Oneofs().Get(i))
	}
	pd.dedent()
	pd.writeIndented("}\n\n")
}

func (pd *ProtoDefinition) writeImport(fileImport protoreflect.FileImport) {
	pd.write("import ")
	if fileImport.IsPublic {
		pd.write("public ")
	}
	pd.write("\"")
	pd.write(fileImport.Path())
	pd.write("\";\n")
}

func (pd *ProtoDefinition) writeFileDescriptor() {
	pd.write("syntax = \"")
	pd.write(pd.descriptor.Syntax().String())
	pd.write("\";\n\n")

	pd.write("package ")
	pd.write(string(pd.descriptor.Package().Name()))
	pd.write(";\n\n")

	// TODO need to handle FileOptions
	filepb := protodesc.ToFileDescriptorProto(pd.descriptor)
	options := filepb.Options
	if options != nil {
	}

	for i := 0; i < pd.descriptor.Imports().Len(); i++ {
		pd.writeImport(pd.descriptor.Imports().Get(i))
	}

	if pd.descriptor.Imports().Len() > 0 {
		pd.write("\n")
	}

	for i := 0; i < pd.descriptor.Services().Len(); i++ {
		pd.writeService(pd.descriptor.Services().Get(i))
	}

	for i := 0; i < pd.descriptor.Messages().Len(); i++ {
		pd.writeMessage(pd.descriptor.Messages().Get(i))
	}

	for i := 0; i < pd.descriptor.Enums().Len(); i++ {
		pd.writeEnum(pd.descriptor.Enums().Get(i))
	}
}

func NewFromBytes(payload []byte) (*ProtoDefinition, error) {
	var pb descriptorpb.FileDescriptorProto
	err := proto.Unmarshal(payload, &pb)
	if err != nil {
		return nil, fmt.Errorf("Couldn't unmarshal proto: %w", err)
	}

	return NewFromDescriptor(pb)
}

func NewFromDescriptor(pb descriptorpb.FileDescriptorProto) (*ProtoDefinition, error) {
	fileOptions := protodesc.FileOptions{AllowUnresolvable: true}
	descriptor, err := fileOptions.New(&pb, &protoregistry.Files{})

	if err != nil {
		return nil, fmt.Errorf("Couldn't create FileDescriptor: %w", err)
	}

	pd := ProtoDefinition{
		pb:         pb,
		descriptor: descriptor,
		filename:   descriptor.Path(),
	}

	pd.writeFileDescriptor()

	return &pd, nil

}