Add editions

Fix overly strict error handling causing valid protos to be skipped

cmd/protodump/main.go

@@ -44,7 +44,7 @@ func writeFile(outputDir string, filename string, content []byte) (string, error
 		return "", fmt.Errorf("failed to evalsymlinks on %s: %v", eval, err)
 	}
 
-	if !strings.HasPrefix(base, outputDirAbs) {
+	if base != outputDirAbs && !strings.HasPrefix(base, outputDirAbs+string(filepath.Separator)) {
 		return "", fmt.Errorf("invalid filepath: %s", base)
 	}
 
@@ -55,10 +55,17 @@ func writeFile(outputDir string, filename string, content []byte) (string, error
 	}
 
 	final := filepath.Join(base, rest, fileBase)
-	err = os.WriteFile(final, content, 0700)
+
+	file, err := os.OpenFile(final, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
+	if err != nil {
+		return "", fmt.Errorf("failed to open file %s: %v", final, err)
+	}
+	defer file.Close()
+	_, err = file.Write(content)
 	if err != nil {
 		return "", fmt.Errorf("failed to write file %s: %v", final, err)
 	}
+
 	return final, nil
 }
 

pkg/protodump/fixtures/edition.proto

@@ -0,0 +1,12 @@
+edition = "2023";
+
+package demo;
+
+option go_package = "./;editions";
+
+message Sample {
+  int32 id = 1;
+  string name = 2;
+  repeated int32 values = 3;
+}
+

pkg/protodump/fixtures/escaping.proto2

@@ -0,0 +1,13 @@
+syntax = "proto2";
+
+package escaping;
+
+option java_package = "a\"b\\c";
+option go_package = "./;escaping";
+
+message M {
+  reserved "ev\"il";
+  optional string note = 1 [default = "evil\";\n}\nmessage Injected {\n  string pwned = 1;\n}\n//"];
+  optional bytes blob = 2 [default = "A\377\001hi"];
+}
+

pkg/protodump/proto.go

@@ -23,6 +23,40 @@ type ProtoDefinition struct {
 	filename    string
 }
 
+func escapeProtoString(s string) string {
+	var b strings.Builder
+	for i := 0; i < len(s); i++ {
+		switch c := s[i]; {
+		case c == '\\':
+			b.WriteString(`\\`)
+		case c == '"':
+			b.WriteString(`\"`)
+		case c == '\n':
+			b.WriteString(`\n`)
+		case c == '\r':
+			b.WriteString(`\r`)
+		case c == '\t':
+			b.WriteString(`\t`)
+		case c < 0x20 || c >= 0x7f:
+			fmt.Fprintf(&b, `\%03o`, c)
+		default:
+			b.WriteByte(c)
+		}
+	}
+	return b.String()
+}
+
+func editionString(d protoreflect.FileDescriptor) (string, error) {
+	edition := protodesc.ToFileDescriptorProto(d).GetEdition().String()
+	if strings.HasPrefix(edition, "EDITION_") {
+		rest := strings.TrimPrefix(edition, "EDITION_")
+		if _, err := strconv.Atoi(rest); err == nil {
+			return rest, nil
+		}
+	}
+	return "unknown", fmt.Errorf("not a source edition: %v", edition)
+}
+
 func (pd *ProtoDefinition) indent() {
 	pd.indendation += 1
 }
@@ -40,6 +74,12 @@ func (pd *ProtoDefinition) write(s string) {
 	pd.builder.WriteString(s)
 }
 
+func (pd *ProtoDefinition) writeQuoted(s string) {
+	pd.write("\"")
+	pd.write(escapeProtoString(s))
+	pd.write("\"")
+}
+
 func (pd *ProtoDefinition) String() string {
 	return pd.builder.String()
 }
@@ -143,7 +183,9 @@ func (pd *ProtoDefinition) writeField(field protoreflect.FieldDescriptor) {
 		pd.write(" [default = ")
 		kind := field.Kind().String()
 		if kind == "string" {
-			pd.write(fmt.Sprintf("\"%s\"", field.Default().String()))
+			pd.writeQuoted(field.Default().String())
+		} else if kind == "bytes" {
+			pd.writeQuoted(string(field.Default().Bytes()))
 		} else if kind == "enum" {
 			pd.write(string(field.DefaultEnumValue().Name()))
 		} else {
@@ -181,9 +223,9 @@ func (pd *ProtoDefinition) writeMessage(message protoreflect.MessageDescriptor)
 
 	for i := 0; i < message.ReservedNames().Len(); i++ {
 		name := message.ReservedNames().Get(i)
-		pd.writeIndented("reserved \"")
-		pd.write(string(name))
-		pd.write("\";\n")
+		pd.writeIndented("reserved ")
+		pd.writeQuoted(string(name))
+		pd.write(";\n")
 	}
 
 	for i := 0; i < message.ReservedRanges().Len(); i++ {
@@ -234,17 +276,16 @@ func (pd *ProtoDefinition) writeImport(fileImport protoreflect.FileImport) {
 	if fileImport.IsPublic {
 		pd.write("public ")
 	}
-	pd.write("\"")
-	pd.write(fileImport.Path())
-	pd.write("\";\n")
+	pd.writeQuoted(fileImport.Path())
+	pd.write(";\n")
 }
 
 func (pd *ProtoDefinition) writeStringFileOptions(name string, value string) {
 	pd.write("option ")
 	pd.write(name)
-	pd.write(" = \"")
-	pd.write(strings.ReplaceAll(value, "\\", "\\\\"))
-	pd.write("\";\n")
+	pd.write(" = ")
+	pd.writeQuoted(value)
+	pd.write(";\n")
 }
 
 func (pd *ProtoDefinition) writeBoolFileOptions(name string, value bool) {
@@ -304,9 +345,16 @@ func (pd *ProtoDefinition) writeFileOptions() {
 }
 
 func (pd *ProtoDefinition) writeFileDescriptor() {
+	if pd.descriptor.Syntax() == protoreflect.Editions {
+		pd.write("edition = \"")
+		edition, _ := editionString(pd.descriptor)
+		pd.write(edition)
+		pd.write("\";\n\n")
+	} else {
 		pd.write("syntax = \"")
 		pd.write(pd.descriptor.Syntax().String())
 		pd.write("\";\n\n")
+	}
 
 	packageName := pd.descriptor.FullName()
 	if packageName != "" {

pkg/protodump/proto_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/descriptorpb"
 )
@@ -44,19 +45,19 @@ func convertProtoToFileDescriptor(filePath string) (*descriptorpb.FileDescriptor
 
 func TestDefinitions(t *testing.T) {
 	files, err := os.ReadDir(FIXTURES)
-	assert.NoError(t, err)
+	require.NoError(t, err)
 
 	for _, file := range files {
 		t.Run(file.Name(), func(t *testing.T) {
 			filePath := path.Join(FIXTURES, file.Name())
 			descriptor, err := convertProtoToFileDescriptor(filePath)
-			assert.NoError(t, err)
+			require.NoError(t, err)
 
 			expected, err := os.ReadFile(filePath)
-			assert.NoError(t, err)
+			require.NoError(t, err)
 
 			actual, err := NewFromDescriptor(descriptor)
-			assert.NoError(t, err)
+			require.NoError(t, err)
 
 			assert.Equal(t, string(expected), actual.String())
 		})

pkg/protodump/scan.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"os"
+	"strings"
 
 	"google.golang.org/protobuf/encoding/protowire"
 )
@@ -15,12 +16,22 @@ func consumeBytes(data []byte, position int) (int, error) {
 	start := position
 	consumedFieldOne := false
 	for {
-		number, _, len := protowire.ConsumeField(data[position:])
-		if len < 0 {
-			err := protowire.ParseError(len)
+		number, _, length := protowire.ConsumeField(data[position:])
+		if length < 0 {
+			err := protowire.ParseError(length)
+			// Treat "invalid field number" as end of data, not an error
+			if strings.Contains(err.Error(), "invalid field number") {
+				return position - start, nil
+			}
+			// Return other parse errors as actual errors
 			return position - start, fmt.Errorf("couldn't consume proto bytes: %w", err)
 		}
 
+		// Prevent infinite loop - if we can't consume any bytes, we're done
+		if length == 0 {
+			return position - start, nil
+		}
+
 		// Only consume Field 1 once (to handle the case where protobuf definitions are adjacent
 		// in program memory)
 		if number == 1 {
@@ -30,7 +41,12 @@ func consumeBytes(data []byte, position int) (int, error) {
 			consumedFieldOne = true
 		}
 
-		position += len
+		position += length
+		
+		// Additional safety check - don't read beyond data bounds
+		if position-start >= len(data[start:]) {
+			return position - start, nil
+		}
 	}
 }