Initial commit

pkg/proto/proto.go

@@ -0,0 +1,260 @@
+package proto
+
+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 == "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(field.Index()))
+	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
+
+}

pkg/proto/proto_test.go

@@ -0,0 +1,18 @@
+package proto
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDefinitions(t *testing.T) {
+	files, err := os.ReadDir("fixtures")
+	assert.NoError(t, err)
+
+	for _, file := range files {
+		t.Run(file.Name(), func(t *testing.T) {
+		})
+	}
+}

pkg/scan/scan.go

@@ -0,0 +1,58 @@
+package scan
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+
+	"google.golang.org/protobuf/encoding/protowire"
+)
+
+func ScanFile(path string) ([][]byte, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("Couldn't open file: %w", err)
+	}
+	return Scan(data), nil
+}
+
+func consumeBytes(data []byte, position int) int {
+	start := position
+	for {
+		_, _, len := protowire.ConsumeField(data[position:])
+		if len < 0 {
+			_ = protowire.ParseError(len)
+			return position - start
+
+		}
+		position += len
+	}
+}
+
+func Scan(data []byte) [][]byte {
+	results := make([][]byte, 0)
+
+	for {
+		index := bytes.Index(data, []byte(".proto"))
+		if index == -1 {
+			break
+		}
+
+		// Look for a 0xa byte: 0 0001 010
+		//                      ^ no more bytes for varint
+		//                           ^ field 1 (file name)
+		//                             ^ type 2 (LEN)
+		start := bytes.LastIndexByte(data[:index], 0xa)
+		if start == -1 {
+			// fmt.Printf("Warn: couldn't find start to proto definition at offset:0x%s\n", strconv.FormatInt(int64(index), 16))
+			data = data[index+1:]
+			continue
+		}
+
+		length := consumeBytes(data, start)
+		results = append(results, data[start:start+length])
+		data = data[start+length:]
+	}
+
+	return results
+}

pkg/scan/scan_test.go

@@ -0,0 +1 @@
+package scan