Unable to extract target repo when nested in oneOf field

3 files changed, 70 insertions, 10 deletions

diff --git a/changelogs/unreleased/po-fix-target-repo.yml b/changelogs/unreleased/po-fix-target-repo.yml
new file mode 100644
index 000000000..d5b1a79f5
--- /dev/null
+++ b/changelogs/unreleased/po-fix-target-repo.yml
@@ -0,0 +1,5 @@
+---
+title: Unable to extract target repo when nested in oneOf field
+merge_request: 1377
+author:
+type: fixed
diff --git a/internal/praefect/protoregistry/targetrepo.go b/internal/praefect/protoregistry/targetrepo.go
index 74a482b30..9af538f9c 100644
--- a/internal/praefect/protoregistry/targetrepo.go
+++ b/internal/praefect/protoregistry/targetrepo.go
@@ -10,7 +10,10 @@ import (
 	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
 )
 
-const protobufTag = "protobuf"
+const (
+	protobufTag      = "protobuf"
+	protobufOneOfTag = "protobuf_oneof"
+)
 
 // reflectFindRepoTarget finds the target repository by using the OID to
 // navigate the struct tags
@@ -41,10 +44,10 @@ func reflectFindRepoTarget(pbMsg proto.Message, targetOID []int) (*gitalypb.Repo
 }
 
 // matches a tag string like "bytes,1,opt,name=repository,proto3"
-var protobufTagRegex = regexp.MustCompile(`^(.*?),(\d+),(.*?),name=(.*?),proto3$`)
+var protobufTagRegex = regexp.MustCompile(`^(.*?),(\d+),(.*?),name=(.*?),proto3(\,oneof)?$`)
 
 const (
-	protobufTagRegexGroups     = 5
+	protobufTagRegexGroups     = 6
 	protobufTagRegexFieldGroup = 2
 )
 
@@ -52,17 +55,20 @@ func findProtoField(msgV reflect.Value, protoField int) (reflect.Value, error) {
 	msgV = reflect.Indirect(msgV)
 	for i := 0; i < msgV.NumField(); i++ {
 		field := msgV.Type().Field(i)
-		tag := field.Tag.Get(protobufTag)
 
-		matches := protobufTagRegex.FindStringSubmatch(tag)
-		if len(matches) != protobufTagRegexGroups {
-			continue
+		ok, err := tryNumberedField(field, protoField)
+		if err != nil {
+			return reflect.Value{}, err
 		}
-
-		fieldStr := matches[protobufTagRegexFieldGroup]
-		if fieldStr == strconv.Itoa(protoField) {
+		if ok {
 			return msgV.FieldByName(field.Name), nil
 		}
+
+		oneofField, ok := tryOneOfField(msgV, field, protoField)
+		if !ok {
+			continue
+		}
+		return oneofField, nil
 	}
 
 	err := fmt.Errorf(
@@ -71,3 +77,39 @@ func findProtoField(msgV reflect.Value, protoField int) (reflect.Value, error) {
 	)
 	return reflect.Value{}, err
 }
+
+func tryNumberedField(field reflect.StructField, protoField int) (bool, error) {
+	tag := field.Tag.Get(protobufTag)
+	matches := protobufTagRegex.FindStringSubmatch(tag)
+	if len(matches) == protobufTagRegexGroups {
+		fieldStr := matches[protobufTagRegexFieldGroup]
+		if fieldStr == strconv.Itoa(protoField) {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+func tryOneOfField(msgV reflect.Value, field reflect.StructField, protoField int) (reflect.Value, bool) {
+	oneOfTag := field.Tag.Get(protobufOneOfTag)
+	if oneOfTag == "" {
+		return reflect.Value{}, false // empty tag means this is not a oneOf field
+	}
+
+	// try all of the oneOf fields until a match is found
+	msgV = msgV.FieldByName(field.Name).Elem().Elem()
+	for i := 0; i < msgV.NumField(); i++ {
+		field = msgV.Type().Field(i)
+
+		ok, err := tryNumberedField(field, protoField)
+		if err != nil {
+			return reflect.Value{}, false
+		}
+		if ok {
+			return msgV.FieldByName(field.Name), true
+		}
+	}
+
+	return reflect.Value{}, false
+}
diff --git a/internal/praefect/protoregistry/targetrepo_test.go b/internal/praefect/protoregistry/targetrepo_test.go
index f2c1f394e..66b39e2cc 100644
--- a/internal/praefect/protoregistry/targetrepo_test.go
+++ b/internal/praefect/protoregistry/targetrepo_test.go
@@ -58,6 +58,19 @@ func TestProtoRegistryTargetRepo(t *testing.T) {
 			pbMsg:     &gitalypb.RepackIncrementalResponse{},
 			expectErr: errors.New("proto message gitaly.RepackIncrementalResponse does not match expected RPC request message gitaly.RepackIncrementalRequest"),
 		},
+		{
+			desc:   "target nested in oneOf",
+			svc:    "OperationService",
+			method: "UserCommitFiles",
+			pbMsg: &gitalypb.UserCommitFilesRequest{
+				UserCommitFilesRequestPayload: &gitalypb.UserCommitFilesRequest_Header{
+					Header: &gitalypb.UserCommitFilesRequestHeader{
+						Repository: testRepos[1],
+					},
+				},
+			},
+			expectRepo: testRepos[1],
+		},
 	}
 
 	for _, tc := range testcases {