diff --git a/pkg/protodump/fixtures/test.proto b/pkg/protodump/fixtures/test.proto index ee508f7..552db03 100644 --- a/pkg/protodump/fixtures/test.proto +++ b/pkg/protodump/fixtures/test.proto @@ -1,6 +1,6 @@ syntax = "proto3"; -package helloworld; +package hello.world; option go_package = "./;helloworld"; @@ -30,14 +30,14 @@ message HelloRequest { } string name = 1; - helloworld.HelloRequest.Test blah = 2; - helloworld.HelloRequest.NestedEnum blahblah = 3; - helloworld.HelloReply blaha = 4; + hello.world.HelloRequest.Test blah = 2; + hello.world.HelloRequest.NestedEnum blahblah = 3; + hello.world.HelloReply blaha = 4; } message HelloReply { - helloworld.HelloRequest.Test blah = 2; - helloworld.HelloRequest.NestedEnum blahblah = 3; + hello.world.HelloRequest.Test blah = 2; + hello.world.HelloRequest.NestedEnum blahblah = 3; string message = 1; } diff --git a/pkg/protodump/proto.go b/pkg/protodump/proto.go index 6db37d0..b0fc313 100644 --- a/pkg/protodump/proto.go +++ b/pkg/protodump/proto.go @@ -17,13 +17,11 @@ import ( type ProtoDefinition struct { builder strings.Builder indendation int - pb descriptorpb.FileDescriptorProto + pb *descriptorpb.FileDescriptorProto descriptor protoreflect.FileDescriptor filename string } -// TODO add proto2 support - func (pd *ProtoDefinition) indent() { pd.indendation += 1 } @@ -118,6 +116,8 @@ func (pd *ProtoDefinition) writeField(field protoreflect.FieldDescriptor) { pd.write("optional ") } else if field.Cardinality().String() == "repeated" { pd.write("repeated ") + } else if field.Cardinality().String() == "required" && pd.descriptor.Syntax().String() == "proto2" { + pd.write("required ") } pd.writeType(field) pd.write(" ") @@ -220,7 +220,11 @@ func (pd *ProtoDefinition) writeStringFileOptions(name string, value string) { } func (pd *ProtoDefinition) writeBoolFileOptions(name string, value bool) { - pd.writeStringFileOptions(name, strconv.FormatBool(value)) + pd.write("option ") + pd.write(name) + pd.write(" = ") + pd.write(strconv.FormatBool(value)) + pd.write(";\n") } func (pd *ProtoDefinition) writeFileOptions() { @@ -276,9 +280,12 @@ func (pd *ProtoDefinition) writeFileDescriptor() { pd.write(pd.descriptor.Syntax().String()) pd.write("\";\n\n") - pd.write("package ") - pd.write(string(pd.descriptor.Package().Name())) - pd.write(";\n\n") + packageName := pd.descriptor.FullName() + if packageName != "" { + pd.write("package ") + pd.write(string(packageName)) + pd.write(";\n\n") + } pd.writeFileOptions() @@ -310,12 +317,12 @@ func NewFromBytes(payload []byte) (*ProtoDefinition, error) { return nil, fmt.Errorf("Couldn't unmarshal proto: %w", err) } - return NewFromDescriptor(pb) + return NewFromDescriptor(&pb) } -func NewFromDescriptor(pb descriptorpb.FileDescriptorProto) (*ProtoDefinition, error) { +func NewFromDescriptor(pb *descriptorpb.FileDescriptorProto) (*ProtoDefinition, error) { fileOptions := protodesc.FileOptions{AllowUnresolvable: true} - descriptor, err := fileOptions.New(&pb, &protoregistry.Files{}) + descriptor, err := fileOptions.New(pb, &protoregistry.Files{}) if err != nil { return nil, fmt.Errorf("Couldn't create FileDescriptor: %w", err) diff --git a/pkg/protodump/proto_test.go b/pkg/protodump/proto_test.go index 49c767d..9b12a8f 100644 --- a/pkg/protodump/proto_test.go +++ b/pkg/protodump/proto_test.go @@ -1,59 +1,46 @@ package protodump import ( - "bytes" - "encoding/hex" "fmt" "os" "os/exec" "path" "regexp" - "strings" "testing" "github.com/stretchr/testify/assert" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/descriptorpb" ) const FIXTURES = "fixtures" var regex = regexp.MustCompile("(?m)var .* = \\[\\]byte{([^}]*)}") -func convertProtoToFileDescriptor(filePath string) ([]byte, error) { +func convertProtoToFileDescriptor(filePath string) (*descriptorpb.FileDescriptorProto, error) { dir, err := os.MkdirTemp("", "") if err != nil { return nil, err } defer os.RemoveAll(dir) - err = exec.Command("protoc", fmt.Sprintf("--go_out=%s", dir), filePath).Run() + filename := "proto.bin" + err = exec.Command("protoc", fmt.Sprintf("--go_out=%s", dir), fmt.Sprintf("--descriptor_set_out=%s", path.Join(dir, filename)), filePath).Run() if err != nil { return nil, err } - fileName := fmt.Sprintf("%s.pb.go", strings.TrimSuffix(path.Base(filePath), ".proto")) - data, err := os.ReadFile(path.Join(dir, fileName)) + data, err := os.ReadFile(path.Join(dir, filename)) if err != nil { return nil, err } - match := regex.FindSubmatch(data) - if len(match) != 2 { - return nil, fmt.Errorf("Could not find file descriptor bytes in %s", fileName) - } - - var hexString strings.Builder - for _, char := range strings.ToLower(string(bytes.ReplaceAll(match[1], []byte("0x"), []byte{}))) { - r := rune(char) - if (r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') { - hexString.WriteRune(r) - } - } - - result, err := hex.DecodeString(hexString.String()) + var descriptor descriptorpb.FileDescriptorSet + err = proto.Unmarshal(data, &descriptor) if err != nil { return nil, err } - return result, nil + return descriptor.GetFile()[0], nil } func TestDefinitions(t *testing.T) { @@ -69,7 +56,7 @@ func TestDefinitions(t *testing.T) { expected, err := os.ReadFile(filePath) assert.NoError(t, err) - actual, err := NewFromBytes(descriptor) + actual, err := NewFromDescriptor(descriptor) assert.NoError(t, err) assert.Equal(t, string(expected), actual.String())