Praefect to throw error on non-repository scoped rpcsjc-praefect-error-on-non-repository-rpcs

5 files changed, 92 insertions, 73 deletions

diff --git a/internal/praefect/auth_test.go b/internal/praefect/auth_test.go
index fcfb17fa8..b2d0f5dd5 100644
--- a/internal/praefect/auth_test.go
+++ b/internal/praefect/auth_test.go
@@ -5,6 +5,8 @@ import (
 	"net"
 	"testing"
 
+	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
+
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -136,6 +138,7 @@ func TestAuthSuccess(t *testing.T) {
 
 			_, err = cli.ServerAccessor(ctx, &mock.SimpleRequest{
 				Value: 1,
+				Repo:  &gitalypb.Repository{StorageName: "praefect", RelativePath: "/doesnt/actually/exist"},
 			})
 
 			assert.NoError(t, err, tc.desc)
@@ -168,10 +171,10 @@ func runServer(t *testing.T, token string, required bool) (*Server, string, func
 	conf := config.Config{
 		Auth: auth.Config{Token: token, Transitioning: !required},
 		VirtualStorages: []*config.VirtualStorage{
-			{
+			&config.VirtualStorage{
 				Name: "praefect",
 				Nodes: []*models.Node{
-					{
+					&models.Node{
 						Storage:        "praefect-internal-0",
 						DefaultPrimary: true,
 						Address:        backend,
diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 04ecd072a..6d0684acf 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -2,6 +2,7 @@ package praefect
 
 import (
 	"context"
+	"errors"
 	"os"
 	"os/signal"
 	"sync"
@@ -78,53 +79,19 @@ func (c *Coordinator) streamDirector(ctx context.Context, fullMethodName string,
 
 	var requestFinalizer func()
 
-	if mi.Scope == protoregistry.ScopeRepository {
-		targetRepo, err := mi.TargetRepo(m)
-		if err != nil {
-			if err == protoregistry.ErrTargetRepoMissing {
-				return nil, status.Errorf(codes.InvalidArgument, err.Error())
-			}
-			return nil, err
-		}
-
-		shard, err := c.nodeMgr.GetShard(targetRepo.GetStorageName())
-		if err != nil {
-			return nil, err
-		}
-
-		primary, err := shard.GetPrimary()
-		if err != nil {
-			return nil, err
-		}
-
-		if err = c.rewriteStorageForRepositoryMessage(mi, m, peeker, primary.GetStorage()); err != nil {
-			if err == protoregistry.ErrTargetRepoMissing {
-				return nil, status.Errorf(codes.InvalidArgument, err.Error())
-			}
-
-			return nil, err
-		}
-
-		if mi.Operation == protoregistry.OpMutator {
-			change := datastore.UpdateRepo
-			if isDestructive(fullMethodName) {
-				change = datastore.DeleteRepo
-			}
-
-			secondaries, err := shard.GetSecondaries()
-			if err != nil {
-				return nil, err
-			}
+	if mi.Scope != protoregistry.ScopeRepository {
+		return nil, errors.New("only repository scoped requests are allowed")
+	}
 
-			if requestFinalizer, err = c.createReplicaJobs(targetRepo, primary, secondaries, change); err != nil {
-				return nil, err
-			}
+	targetRepo, err := mi.TargetRepo(m)
+	if err != nil {
+		if err == protoregistry.ErrTargetRepoMissing {
+			return nil, status.Errorf(codes.InvalidArgument, err.Error())
 		}
-
-		return proxy.NewStreamParameters(ctx, primary.GetConnection(), requestFinalizer, nil), nil
+		return nil, err
 	}
 
-	shard, err := c.nodeMgr.GetShard(c.conf.VirtualStorages[0].Name)
+	shard, err := c.nodeMgr.GetShard(targetRepo.GetStorageName())
 	if err != nil {
 		return nil, err
 	}
@@ -134,10 +101,32 @@ func (c *Coordinator) streamDirector(ctx context.Context, fullMethodName string,
 		return nil, err
 	}
 
-	return proxy.NewStreamParameters(ctx, primary.GetConnection(), noopRequestFinalizer, nil), nil
-}
+	if err = c.rewriteStorageForRepositoryMessage(mi, m, peeker, primary.GetStorage()); err != nil {
+		if err == protoregistry.ErrTargetRepoMissing {
+			return nil, status.Errorf(codes.InvalidArgument, err.Error())
+		}
 
-var noopRequestFinalizer = func() {}
+		return nil, err
+	}
+
+	if mi.Operation == protoregistry.OpMutator {
+		change := datastore.UpdateRepo
+		if isDestructive(fullMethodName) {
+			change = datastore.DeleteRepo
+		}
+
+		secondaries, err := shard.GetSecondaries()
+		if err != nil {
+			return nil, err
+		}
+
+		if requestFinalizer, err = c.createReplicaJobs(targetRepo, primary, secondaries, change); err != nil {
+			return nil, err
+		}
+	}
+
+	return proxy.NewStreamParameters(ctx, primary.GetConnection(), requestFinalizer, nil), nil
+}
 
 func (c *Coordinator) rewriteStorageForRepositoryMessage(mi protoregistry.MethodInfo, m proto.Message, peeker proxy.StreamModifier, primaryStorage string) error {
 	targetRepo, err := mi.TargetRepo(m)
diff --git a/internal/praefect/mock/mock.pb.go b/internal/praefect/mock/mock.pb.go
index edc09f622..b3fc533bb 100644
--- a/internal/praefect/mock/mock.pb.go
+++ b/internal/praefect/mock/mock.pb.go
@@ -12,6 +12,8 @@ import (
 	empty "github.com/golang/protobuf/ptypes/empty"
 	gitalypb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	grpc "google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
 )
 
 // Reference imports to suppress errors if they are not otherwise used.
@@ -26,10 +28,11 @@ var _ = math.Inf
 const _ = proto.ProtoPackageIsVersion3 // please upgrade the proto package
 
 type SimpleRequest struct {
-	Value                int32    `protobuf:"varint,1,opt,name=value,proto3" json:"value,omitempty"`
-	XXX_NoUnkeyedLiteral struct{} `json:"-"`
-	XXX_unrecognized     []byte   `json:"-"`
-	XXX_sizecache        int32    `json:"-"`
+	Value                int32                `protobuf:"varint,1,opt,name=value,proto3" json:"value,omitempty"`
+	Repo                 *gitalypb.Repository `protobuf:"bytes,2,opt,name=repo,proto3" json:"repo,omitempty"`
+	XXX_NoUnkeyedLiteral struct{}             `json:"-"`
+	XXX_unrecognized     []byte               `json:"-"`
+	XXX_sizecache        int32                `json:"-"`
 }
 
 func (m *SimpleRequest) Reset()         { *m = SimpleRequest{} }
@@ -64,6 +67,13 @@ func (m *SimpleRequest) GetValue() int32 {
 	return 0
 }
 
+func (m *SimpleRequest) GetRepo() *gitalypb.Repository {
+	if m != nil {
+		return m.Repo
+	}
+	return nil
+}
+
 type SimpleResponse struct {
 	Value                int32    `protobuf:"varint,2,opt,name=value,proto3" json:"value,omitempty"`
 	XXX_NoUnkeyedLiteral struct{} `json:"-"`
@@ -151,25 +161,25 @@ func init() {
 func init() { proto.RegisterFile("mock.proto", fileDescriptor_6fa4806c90f7156d) }
 
 var fileDescriptor_6fa4806c90f7156d = []byte{
-	// 275 bytes of a gzipped FileDescriptorProto
-	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x9c, 0x50, 0xcd, 0x4a, 0xc3, 0x40,
-	0x10, 0x66, 0x4b, 0x22, 0x71, 0xaa, 0xa5, 0x5d, 0x8b, 0xc8, 0x7a, 0x91, 0x80, 0x25, 0xa7, 0x2d,
-	0xad, 0xf8, 0x00, 0x1e, 0x0a, 0x7a, 0xf0, 0x92, 0xe2, 0x03, 0xa4, 0x71, 0x8c, 0xc1, 0xa4, 0xb3,
-	0xee, 0x6e, 0x0a, 0x79, 0x92, 0x3e, 0xa4, 0x6f, 0xd0, 0x93, 0x6c, 0x96, 0x60, 0x7b, 0xf5, 0x36,
-	0x33, 0xdf, 0xcf, 0x7e, 0xfb, 0x01, 0xd4, 0x94, 0x7f, 0x49, 0xa5, 0xc9, 0x12, 0x0f, 0xdc, 0x2c,
-	0x2e, 0xcc, 0x67, 0xa6, 0xf1, 0xdd, 0xdf, 0xc4, 0x6d, 0x41, 0x54, 0x54, 0x38, 0xef, 0xb6, 0x4d,
-	0xf3, 0x31, 0xc7, 0x5a, 0xd9, 0xd6, 0x83, 0xf1, 0x3d, 0x5c, 0xae, 0xcb, 0x5a, 0x55, 0x98, 0xe2,
-	0x77, 0x83, 0xc6, 0xf2, 0x29, 0x84, 0xbb, 0xac, 0x6a, 0xf0, 0x86, 0xdd, 0xb1, 0x24, 0x4c, 0xfd,
-	0x12, 0xcf, 0x60, 0xd4, 0xd3, 0x8c, 0xa2, 0xad, 0xc1, 0x3f, 0xde, 0xe0, 0x98, 0xf7, 0x08, 0xc3,
-	0x14, 0x15, 0xf5, 0x66, 0x33, 0x08, 0x34, 0x2a, 0xea, 0xbc, 0x86, 0x4b, 0x2e, 0x8b, 0xd2, 0x66,
-	0x55, 0x2b, 0x1d, 0xc5, 0x94, 0x96, 0x74, 0x9b, 0x76, 0xf8, 0xf2, 0x87, 0xf5, 0x31, 0xd6, 0xa8,
-	0x77, 0x65, 0x8e, 0x7c, 0x05, 0x23, 0x37, 0xa2, 0x7e, 0xca, 0x73, 0x34, 0x86, 0x34, 0xbf, 0x92,
-	0xdd, 0x3f, 0x4f, 0xd2, 0x8a, 0xe9, 0xe9, 0xd1, 0x67, 0x8b, 0xa3, 0xc3, 0x3e, 0x09, 0xa2, 0xc1,
-	0x98, 0xf1, 0x17, 0x98, 0xb8, 0xc7, 0x7a, 0x93, 0xb7, 0x6d, 0xa6, 0x5b, 0x3e, 0xf1, 0xa2, 0xa3,
-	0xa0, 0xe2, 0x5a, 0xfa, 0x92, 0x64, 0x5f, 0x92, 0x5c, 0xb9, 0x92, 0xe2, 0xf3, 0xc3, 0x3e, 0x09,
-	0xa3, 0x81, 0x60, 0x0b, 0xfe, 0x0c, 0x63, 0xa7, 0x78, 0x6d, 0x6c, 0x66, 0xff, 0xed, 0xc4, 0x04,
-	0x5b, 0x6c, 0xce, 0x3a, 0xe8, 0xe1, 0x37, 0x00, 0x00, 0xff, 0xff, 0x3c, 0x0c, 0xc7, 0x57, 0xb9,
-	0x01, 0x00, 0x00,
+	// 280 bytes of a gzipped FileDescriptorProto
+	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x9c, 0x90, 0x31, 0x4e, 0xc3, 0x30,
+	0x14, 0x86, 0xe5, 0xa8, 0x41, 0xe0, 0x42, 0x45, 0x4d, 0x85, 0x2a, 0xb3, 0x54, 0x19, 0xaa, 0x4c,
+	0xae, 0x1a, 0xc4, 0x01, 0x18, 0x40, 0x30, 0x74, 0x49, 0xc5, 0x01, 0xdc, 0xf0, 0x08, 0x11, 0x49,
+	0x9f, 0xb1, 0x9d, 0x4a, 0x39, 0x09, 0x97, 0xe4, 0x04, 0x4c, 0xc8, 0xb1, 0xac, 0xb6, 0x0b, 0x03,
+	0x9b, 0xdf, 0xfb, 0xf5, 0x7f, 0xef, 0xf7, 0x4f, 0x69, 0x83, 0xc5, 0x87, 0x50, 0x1a, 0x2d, 0xb2,
+	0x81, 0x7b, 0xf3, 0x73, 0xf3, 0x2e, 0x35, 0xbc, 0xfa, 0x1d, 0xbf, 0x29, 0x11, 0xcb, 0x1a, 0x16,
+	0xfd, 0xb4, 0x69, 0xdf, 0x16, 0xd0, 0x28, 0xdb, 0x79, 0x31, 0x59, 0xd1, 0x8b, 0x75, 0xd5, 0xa8,
+	0x1a, 0x72, 0xf8, 0x6c, 0xc1, 0x58, 0x36, 0xa1, 0xf1, 0x4e, 0xd6, 0x2d, 0x4c, 0xc9, 0x8c, 0xa4,
+	0x71, 0xee, 0x07, 0x36, 0xa7, 0x03, 0x0d, 0x0a, 0xa7, 0xd1, 0x8c, 0xa4, 0xc3, 0x8c, 0x89, 0xb2,
+	0xb2, 0xb2, 0xee, 0x44, 0x0e, 0x0a, 0x4d, 0x65, 0x51, 0x77, 0x79, 0xaf, 0x27, 0x73, 0x3a, 0x0a,
+	0x38, 0xa3, 0x70, 0x6b, 0x60, 0xcf, 0x8b, 0x0e, 0x78, 0xc9, 0x1d, 0x1d, 0x3a, 0x6f, 0x38, 0x1a,
+	0xf0, 0xe4, 0x6f, 0x7c, 0xf6, 0x4d, 0x42, 0xdc, 0x35, 0xe8, 0x5d, 0x55, 0x00, 0x7b, 0xa4, 0x23,
+	0xf7, 0x04, 0x7d, 0x5f, 0x14, 0x60, 0x0c, 0x6a, 0x76, 0x25, 0xfa, 0x3e, 0x8e, 0x7e, 0xc5, 0x27,
+	0xc7, 0x4b, 0x9f, 0x2d, 0x39, 0xfb, 0xf9, 0x4a, 0xe3, 0xd3, 0x88, 0x93, 0x8c, 0x3d, 0xd3, 0xb1,
+	0xbb, 0x16, 0x28, 0x2f, 0x5b, 0xa9, 0x3b, 0x36, 0xf6, 0xae, 0x83, 0xa4, 0xfc, 0x5a, 0xf8, 0x36,
+	0x45, 0x68, 0x53, 0x3c, 0xb8, 0x36, 0xf7, 0xa8, 0x25, 0x7b, 0xa2, 0x97, 0xce, 0xb1, 0x6a, 0xad,
+	0xb4, 0xff, 0x26, 0x11, 0x4e, 0x96, 0x9b, 0x93, 0x5e, 0xba, 0xfd, 0x0d, 0x00, 0x00, 0xff, 0xff,
+	0x83, 0xbf, 0x8a, 0x39, 0xe2, 0x01, 0x00, 0x00,
 }
 
 // Reference imports to suppress errors if they are not otherwise used.
@@ -237,6 +247,20 @@ type SimpleServiceServer interface {
 	RepoMutatorUnary(context.Context, *RepoRequest) (*empty.Empty, error)
 }
 
+// UnimplementedSimpleServiceServer can be embedded to have forward compatible implementations.
+type UnimplementedSimpleServiceServer struct {
+}
+
+func (*UnimplementedSimpleServiceServer) ServerAccessor(ctx context.Context, req *SimpleRequest) (*SimpleResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method ServerAccessor not implemented")
+}
+func (*UnimplementedSimpleServiceServer) RepoAccessorUnary(ctx context.Context, req *RepoRequest) (*empty.Empty, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RepoAccessorUnary not implemented")
+}
+func (*UnimplementedSimpleServiceServer) RepoMutatorUnary(ctx context.Context, req *RepoRequest) (*empty.Empty, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RepoMutatorUnary not implemented")
+}
+
 func RegisterSimpleServiceServer(s *grpc.Server, srv SimpleServiceServer) {
 	s.RegisterService(&_SimpleService_serviceDesc, srv)
 }
diff --git a/internal/praefect/mock/mock.proto b/internal/praefect/mock/mock.proto
index a1718fdd7..e7dcd43df 100644
--- a/internal/praefect/mock/mock.proto
+++ b/internal/praefect/mock/mock.proto
@@ -12,6 +12,7 @@ import "google/protobuf/empty.proto";
 
 message SimpleRequest {
 	int32 value = 1;
+    gitaly.Repository repo = 2;
 }
 
 message SimpleResponse {
@@ -27,7 +28,8 @@ service SimpleService {
   rpc ServerAccessor(SimpleRequest) returns (SimpleResponse) {
     option (gitaly.op_type) = {
       op: ACCESSOR
-      scope_level: SERVER
+      scope_level: REPOSITORY
+      target_repository_field: "2"
     };
   }
 	
diff --git a/internal/praefect/server_test.go b/internal/praefect/server_test.go
index f8f6012be..aecca90d8 100644
--- a/internal/praefect/server_test.go
+++ b/internal/praefect/server_test.go
@@ -51,13 +51,14 @@ func TestServerRouteServerAccessor(t *testing.T) {
 	cli, _, cleanup := runPraefectServerWithMock(t, conf, backends)
 	defer cleanup()
 
-	expectReq := &mock.SimpleRequest{Value: 1}
+	expectReq := &mock.SimpleRequest{Value: 1, Repo: &gitalypb.Repository{StorageName: "praefect", RelativePath: "/doesnt/actually/exist"}}
 
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
 
 		actualReq := <-reqQ
+		actualReq.Repo.StorageName = "praefect"
 		assert.True(t, proto.Equal(expectReq, actualReq),
 			"received unexpected request value: %+v instead of %+v", actualReq, expectReq)
 	}()