Merge branch 'pks-praefect-info-ruby-sidecar' into 'master'

Fix connecting to Praefect service from Ruby sidecar

Closes #3017

See merge request gitlab-org/gitaly!2451

7 files changed, 143 insertions, 6 deletions

diff --git a/changelogs/unreleased/pks-praefect-info-ruby-sidecar.yml b/changelogs/unreleased/pks-praefect-info-ruby-sidecar.yml
new file mode 100644
index 000000000..fc6ab6a4b
--- /dev/null
+++ b/changelogs/unreleased/pks-praefect-info-ruby-sidecar.yml
@@ -0,0 +1,5 @@
+---
+title: Fix connecting to Praefect service from Ruby sidecar
+merge_request: 2451
+author:
+type: fixed
diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 48c1dec25..dc0564653 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -140,7 +140,12 @@ func (c *Coordinator) directRepositoryScopedMessage(ctx context.Context, call gr
 		return nil, helper.ErrInvalidArgumentf("repo scoped: target repo is invalid")
 	}
 
-	if ctx, err = metadata.InjectPraefectServer(ctx, c.conf); err != nil {
+	praefectServer, err := metadata.PraefectFromConfig(c.conf)
+	if err != nil {
+		return nil, fmt.Errorf("repo scoped: could not create Praefect configuration: %w", err)
+	}
+
+	if ctx, err = praefectServer.Inject(ctx); err != nil {
 		return nil, fmt.Errorf("repo scoped: could not inject Praefect server: %w", err)
 	}
 
diff --git a/internal/praefect/metadata/server.go b/internal/praefect/metadata/server.go
index 90838b34a..fa6c36bde 100644
--- a/internal/praefect/metadata/server.go
+++ b/internal/praefect/metadata/server.go
@@ -44,8 +44,8 @@ type PraefectServer struct {
 	Token string `json:"token"`
 }
 
-// InjectPraefectServer injects Praefect connection metadata into an incoming context
-func InjectPraefectServer(ctx context.Context, conf config.Config) (context.Context, error) {
+// PraefectFromConfig creates a Praefect server for a given configuration.
+func PraefectFromConfig(conf config.Config) (*PraefectServer, error) {
 	praefectServer := PraefectServer{Token: conf.Auth.Token}
 
 	addrBySchema := map[string]*string{
@@ -82,7 +82,12 @@ func InjectPraefectServer(ctx context.Context, conf config.Config) (context.Cont
 		*addrBySchema[endpoint.schema] = addr
 	}
 
-	serialized, err := praefectServer.serialize()
+	return &praefectServer, nil
+}
+
+// Inject injects Praefect connection metadata into an incoming context
+func (p *PraefectServer) Inject(ctx context.Context) (context.Context, error) {
+	serialized, err := p.serialize()
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/praefect/metadata/server_test.go b/internal/praefect/metadata/server_test.go
index 9518d5d54..44ab3af2b 100644
--- a/internal/praefect/metadata/server_test.go
+++ b/internal/praefect/metadata/server_test.go
@@ -153,7 +153,11 @@ func TestPraefect_InjectMetadata(t *testing.T) {
 			}
 
 			ctx = peer.NewContext(ctx, tc.peer)
-			ctx, err := InjectPraefectServer(ctx, cfg)
+
+			praefectServer, err := PraefectFromConfig(cfg)
+			require.NoError(t, err)
+
+			ctx, err = praefectServer.Inject(ctx)
 			require.NoError(t, err)
 
 			server, err := PraefectFromContext(ctx)
diff --git a/internal/rubyserver/proxy.go b/internal/rubyserver/proxy.go
index f0b8e3f9b..17c64c6ba 100644
--- a/internal/rubyserver/proxy.go
+++ b/internal/rubyserver/proxy.go
@@ -62,6 +62,19 @@ func setHeaders(ctx context.Context, repo *gitalypb.Repository, mustExist bool)
 		repoAltDirsHeader, repoAltDirsCombined,
 	)
 
+	// While it looks weird that we're extracting and then re-injecting the
+	// Praefect server info into the context, `PraefectFromContext()` will
+	// also resolve connection information from the context's peer info.
+	// Thus the re-injected connection info will contain resolved addresses.
+	if praefectServer, err := praefect_metadata.PraefectFromContext(ctx); err == nil {
+		ctx, err = praefectServer.Inject(ctx)
+		if err != nil {
+			return nil, err
+		}
+	} else if err != praefect_metadata.ErrPraefectServerNotFound {
+		return nil, err
+	}
+
 	// list of http/2 headers that will be forwarded as-is to gitaly-ruby
 	proxyHeaderWhitelist := []string{
 		"gitaly-servers",
diff --git a/internal/service/operations/branches_test.go b/internal/service/operations/branches_test.go
index 78d475103..bd235e572 100644
--- a/internal/service/operations/branches_test.go
+++ b/internal/service/operations/branches_test.go
@@ -2,17 +2,34 @@ package operations
 
 import (
 	"context"
+	"fmt"
+	"net"
 	"os/exec"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/internal/config"
 	"gitlab.com/gitlab-org/gitaly/internal/git/log"
+	"gitlab.com/gitlab-org/gitaly/internal/helper"
 	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
+	"gitlab.com/gitlab-org/gitaly/internal/praefect/metadata"
+	"gitlab.com/gitlab-org/gitaly/internal/service/hook"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
 )
 
+type testTransactionServer struct {
+	called int
+}
+
+func (s *testTransactionServer) VoteTransaction(ctx context.Context, in *gitalypb.VoteTransactionRequest) (*gitalypb.VoteTransactionResponse, error) {
+	s.called++
+	return &gitalypb.VoteTransactionResponse{
+		State: gitalypb.VoteTransactionResponse_COMMIT,
+	}, nil
+}
+
 func TestSuccessfulCreateBranchRequest(t *testing.T) {
 	testRepo, testRepoPath, cleanupFn := testhelper.NewTestRepo(t)
 	defer cleanupFn()
@@ -89,6 +106,91 @@ func TestSuccessfulGitHooksForUserCreateBranchRequest(t *testing.T) {
 	}
 }
 
+func TestUserCreateBranchWithTransaction(t *testing.T) {
+	testRepo, testRepoPath, cleanup := testhelper.NewTestRepo(t)
+	defer cleanup()
+
+	internalSocket := config.GitalyInternalSocketPath()
+	internalListener, err := net.Listen("unix", internalSocket)
+	require.NoError(t, err)
+
+	tcpSocket, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	transactionServer := &testTransactionServer{}
+	srv := testhelper.NewServerWithAuth(t, nil, nil, config.Config.Auth.Token)
+	gitalypb.RegisterOperationServiceServer(srv.GrpcServer(), &server{ruby: RubyServer})
+	gitalypb.RegisterHookServiceServer(srv.GrpcServer(), hook.NewServer(hook.GitlabAPIStub, config.Config.Hooks))
+	gitalypb.RegisterRefTransactionServer(srv.GrpcServer(), transactionServer)
+
+	require.NoError(t, srv.Start())
+	defer srv.Stop()
+
+	go srv.GrpcServer().Serve(internalListener)
+	go srv.GrpcServer().Serve(tcpSocket)
+
+	testcases := []struct {
+		desc    string
+		address string
+		server  metadata.PraefectServer
+	}{
+		{
+			desc:    "explicit TCP address",
+			address: tcpSocket.Addr().String(),
+			server: metadata.PraefectServer{
+				ListenAddr: fmt.Sprintf("tcp://" + tcpSocket.Addr().String()),
+				Token:      config.Config.Auth.Token,
+			},
+		},
+		{
+			desc:    "catch-all TCP address",
+			address: tcpSocket.Addr().String(),
+			server: metadata.PraefectServer{
+				ListenAddr: fmt.Sprintf("tcp://0.0.0.0:%d", tcpSocket.Addr().(*net.TCPAddr).Port),
+				Token:      config.Config.Auth.Token,
+			},
+		},
+		{
+			desc:    "Unix socket",
+			address: "unix://" + internalSocket,
+			server: metadata.PraefectServer{
+				SocketPath: "unix://" + internalSocket,
+				Token:      config.Config.Auth.Token,
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			defer exec.Command("git", "-C", testRepoPath, "branch", "-D", "new-branch").Run()
+
+			client, conn := newOperationClient(t, tc.address)
+			defer conn.Close()
+
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+			ctx, err = tc.server.Inject(ctx)
+			require.NoError(t, err)
+			ctx, err = metadata.InjectTransaction(ctx, 1, "node", true)
+			require.NoError(t, err)
+			ctx = helper.IncomingToOutgoing(ctx)
+
+			request := &gitalypb.UserCreateBranchRequest{
+				Repository: testRepo,
+				BranchName: []byte("new-branch"),
+				StartPoint: []byte("c7fbe50c7c7419d9701eebe64b1fdacc3df5b9dd"),
+				User:       testhelper.TestUser,
+			}
+
+			transactionServer.called = 0
+			response, err := client.UserCreateBranch(ctx, request)
+			require.NoError(t, err)
+			require.Empty(t, response.PreReceiveError)
+			require.Equal(t, 1, transactionServer.called)
+		})
+	}
+}
+
 func testSuccessfulGitHooksForUserCreateBranchRequest(t *testing.T, ctx context.Context) {
 	testRepo, testRepoPath, cleanupFn := testhelper.NewTestRepo(t)
 	defer cleanupFn()
diff --git a/internal/service/smarthttp/receive_pack_test.go b/internal/service/smarthttp/receive_pack_test.go
index b6a306957..a2bb352a3 100644
--- a/internal/service/smarthttp/receive_pack_test.go
+++ b/internal/service/smarthttp/receive_pack_test.go
@@ -623,11 +623,14 @@ func TestPostReceiveWithReferenceTransactionHook(t *testing.T) {
 			// RefTransaction server for Gitaly itself. As this is the only Praefect
 			// service required in this context, we can just pretend that
 			// Gitaly is the Praefect server and inject it.
-			ctx, err = metadata.InjectPraefectServer(ctx, pconfig.Config{
+			praefectServer, err := metadata.PraefectFromConfig(pconfig.Config{
 				SocketPath: "unix://" + gitalySocketPath,
 			})
 			require.NoError(t, err)
 
+			ctx, err = praefectServer.Inject(ctx)
+			require.NoError(t, err)
+
 			ctx = helper.IncomingToOutgoing(ctx)
 			ctx = features.Disable(ctx)