proto: Add `SetCustomHooks` RPC

Currently the `RestoreCustomHooks` RPC is used to set the custom git
hooks for a repository from a backup. In the future this RPC will be
used more generally to write hooks to a specified repository. To
accommodate this, this change implements a new RPC `SetCustomHooks`
which performs the exact same operation but which a more general name.

3 files changed, 255 insertions, 94 deletions

diff --git a/internal/gitaly/service/repository/restore_custom_hooks.go b/internal/gitaly/service/repository/restore_custom_hooks.go
index a3b24d392..5e231a6f9 100644
--- a/internal/gitaly/service/repository/restore_custom_hooks.go
+++ b/internal/gitaly/service/repository/restore_custom_hooks.go
@@ -26,6 +26,40 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/streamio"
 )
 
+// SetCustomHooks sets the git hooks for a repository. The hooks are sent in a
+// tar archive containing a `custom_hooks` directory. This directory is
+// ultimately extracted to the repository.
+func (s *server) SetCustomHooks(stream gitalypb.RepositoryService_SetCustomHooksServer) error {
+	ctx := stream.Context()
+
+	firstRequest, err := stream.Recv()
+	if err != nil {
+		return structerr.NewInternal("first request failed %w", err)
+	}
+
+	repo := firstRequest.GetRepository()
+	if err := service.ValidateRepository(repo); err != nil {
+		return structerr.NewInvalidArgument("%w", err)
+	}
+
+	reader := streamio.NewReader(func() ([]byte, error) {
+		if firstRequest != nil {
+			data := firstRequest.GetData()
+			firstRequest = nil
+			return data, nil
+		}
+
+		request, err := stream.Recv()
+		return request.GetData(), err
+	})
+
+	if err := s.setCustomHooks(ctx, reader, repo); err != nil {
+		return structerr.NewInternal("%w", err)
+	}
+
+	return stream.SendAndClose(&gitalypb.SetCustomHooksResponse{})
+}
+
 // RestoreCustomHooks sets the git hooks for a repository. The hooks are sent in
 // a tar archive containing a `custom_hooks` directory. This directory is
 // ultimately extracted to the repository.
@@ -53,29 +87,37 @@ func (s *server) RestoreCustomHooks(stream gitalypb.RepositoryService_RestoreCus
 		return request.GetData(), err
 	})
 
+	if err := s.setCustomHooks(ctx, reader, repo); err != nil {
+		return structerr.NewInternal("%w", err)
+	}
+
+	return stream.SendAndClose(&gitalypb.RestoreCustomHooksResponse{})
+}
+
+func (s *server) setCustomHooks(ctx context.Context, reader io.Reader, repo repository.GitRepo) error {
 	if featureflag.TransactionalRestoreCustomHooks.IsEnabled(ctx) {
-		if err := s.restoreCustomHooks(ctx, reader, repo); err != nil {
-			return structerr.NewInternal("setting custom hooks: %w", err)
+		if err := s.setCustomHooksTransaction(ctx, reader, repo); err != nil {
+			return fmt.Errorf("writing custom hooks: %w", err)
 		}
 
-		return stream.SendAndClose(&gitalypb.RestoreCustomHooksResponse{})
+		return nil
 	}
 
 	repoPath, err := s.locator.GetPath(repo)
 	if err != nil {
-		return structerr.NewInternal("getting repo path failed %w", err)
+		return fmt.Errorf("getting repo path failed %w", err)
 	}
 
 	if err := extractHooks(ctx, reader, repoPath); err != nil {
-		return structerr.NewInternal("extracting hooks: %w", err)
+		return fmt.Errorf("extracting hooks: %w", err)
 	}
 
-	return stream.SendAndClose(&gitalypb.RestoreCustomHooksResponse{})
+	return nil
 }
 
-// restoreCustomHooks transactionally and atomically sets the provided custom
-// hooks for the specified repository.
-func (s *server) restoreCustomHooks(ctx context.Context, tar io.Reader, repo repository.GitRepo) error {
+// setCustomHooksTransaction transactionally and atomically sets the provided
+// custom hooks for the specified repository.
+func (s *server) setCustomHooksTransaction(ctx context.Context, tar io.Reader, repo repository.GitRepo) error {
 	repoPath, err := s.locator.GetRepoPath(repo)
 	if err != nil {
 		return fmt.Errorf("getting repo path: %w", err)
diff --git a/internal/gitaly/service/repository/restore_custom_hooks_test.go b/internal/gitaly/service/repository/restore_custom_hooks_test.go
index 10a3a5c5e..c020b95dd 100644
--- a/internal/gitaly/service/repository/restore_custom_hooks_test.go
+++ b/internal/gitaly/service/repository/restore_custom_hooks_test.go
@@ -28,129 +28,247 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-func TestSuccessfulRestoreCustomHooksRequest(t *testing.T) {
+func TestSetCustomHooksRequest_success(t *testing.T) {
 	t.Parallel()
 	testhelper.NewFeatureSets(featureflag.TransactionalRestoreCustomHooks).
-		Run(t, testSuccessfulRestoreCustomHooksRequest)
+		Run(t, testSuccessfulSetCustomHooksRequest)
 }
 
-func testSuccessfulRestoreCustomHooksRequest(t *testing.T, ctx context.Context) {
+func testSuccessfulSetCustomHooksRequest(t *testing.T, ctx context.Context) {
 	t.Parallel()
 
-	cfg := testcfg.Build(t)
-	testcfg.BuildGitalyHooks(t, cfg)
-	txManager := transaction.NewTrackingManager()
-
-	client, addr := runRepositoryService(t, cfg, nil, testserver.WithTransactionManager(txManager))
-	cfg.SocketPath = addr
-
-	ctx, err := txinfo.InjectTransaction(ctx, 1, "node", true)
-	require.NoError(t, err)
+	for _, tc := range []struct {
+		desc         string
+		streamWriter func(*testing.T, context.Context, *gitalypb.Repository, gitalypb.RepositoryServiceClient) (io.Writer, func())
+	}{
+		{
+			desc: "SetCustomHooks",
+			streamWriter: func(t *testing.T, ctx context.Context, repo *gitalypb.Repository, client gitalypb.RepositoryServiceClient) (io.Writer, func()) {
+				stream, err := client.SetCustomHooks(ctx)
+				require.NoError(t, err)
+
+				request := &gitalypb.SetCustomHooksRequest{Repository: repo}
+				writer := streamio.NewWriter(func(p []byte) error {
+					request.Data = p
+					if err := stream.Send(request); err != nil {
+						return err
+					}
+
+					request = &gitalypb.SetCustomHooksRequest{}
+					return nil
+				})
+
+				closeFunc := func() {
+					_, err = stream.CloseAndRecv()
+					require.NoError(t, err)
+				}
+
+				return writer, closeFunc
+			},
+		},
+		{
+			desc: "RestoreCustomHooks",
+			streamWriter: func(t *testing.T, ctx context.Context, repo *gitalypb.Repository, client gitalypb.RepositoryServiceClient) (io.Writer, func()) {
+				stream, err := client.RestoreCustomHooks(ctx)
+				require.NoError(t, err)
+
+				request := &gitalypb.RestoreCustomHooksRequest{Repository: repo}
+				writer := streamio.NewWriter(func(p []byte) error {
+					request.Data = p
+					if err := stream.Send(request); err != nil {
+						return err
+					}
+
+					request = &gitalypb.RestoreCustomHooksRequest{}
+					return nil
+				})
+
+				closeFunc := func() {
+					_, err = stream.CloseAndRecv()
+					require.NoError(t, err)
+				}
+
+				return writer, closeFunc
+			},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			cfg := testcfg.Build(t)
+			testcfg.BuildGitalyHooks(t, cfg)
+			txManager := transaction.NewTrackingManager()
 
-	ctx = metadata.IncomingToOutgoing(ctx)
-	repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
+			client, addr := runRepositoryService(t, cfg, nil, testserver.WithTransactionManager(txManager))
+			cfg.SocketPath = addr
 
-	// reset the txManager since CreateRepository would have done
-	// voting
-	txManager.Reset()
-	stream, err := client.RestoreCustomHooks(ctx)
-	require.NoError(t, err)
+			ctx, err := txinfo.InjectTransaction(ctx, 1, "node", true)
+			require.NoError(t, err)
 
-	request := &gitalypb.RestoreCustomHooksRequest{Repository: repo}
+			ctx = metadata.IncomingToOutgoing(ctx)
+			repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
 
-	writer := streamio.NewWriter(func(p []byte) error {
-		request.Data = p
-		if err := stream.Send(request); err != nil {
-			return err
-		}
+			// reset the txManager since CreateRepository would have done
+			// voting
+			txManager.Reset()
 
-		request = &gitalypb.RestoreCustomHooksRequest{}
-		return nil
-	})
+			writer, closeStream := tc.streamWriter(t, ctx, repo, client)
 
-	file, err := os.Open("testdata/custom_hooks.tar")
-	require.NoError(t, err)
+			file, err := os.Open("testdata/custom_hooks.tar")
+			require.NoError(t, err)
 
-	_, err = io.Copy(writer, file)
-	require.NoError(t, err)
-	_, err = stream.CloseAndRecv()
-	require.NoError(t, err)
+			_, err = io.Copy(writer, file)
+			require.NoError(t, err)
+			closeStream()
 
-	voteHash, err := newDirectoryVote(filepath.Join(repoPath, customHooksDir))
-	require.NoError(t, err)
+			voteHash, err := newDirectoryVote(filepath.Join(repoPath, customHooksDir))
+			require.NoError(t, err)
 
-	testhelper.MustClose(t, file)
+			testhelper.MustClose(t, file)
 
-	expectedVote, err := voteHash.Vote()
-	require.NoError(t, err)
+			expectedVote, err := voteHash.Vote()
+			require.NoError(t, err)
 
-	require.FileExists(t, filepath.Join(repoPath, "custom_hooks", "pre-push.sample"))
+			require.FileExists(t, filepath.Join(repoPath, "custom_hooks", "pre-push.sample"))
 
-	if featureflag.TransactionalRestoreCustomHooks.IsEnabled(ctx) {
-		require.Equal(t, 2, len(txManager.Votes()))
-		assert.Equal(t, voting.Prepared, txManager.Votes()[0].Phase)
-		assert.Equal(t, expectedVote, txManager.Votes()[1].Vote)
-		assert.Equal(t, voting.Committed, txManager.Votes()[1].Phase)
-	} else {
-		require.Equal(t, 0, len(txManager.Votes()))
+			if featureflag.TransactionalRestoreCustomHooks.IsEnabled(ctx) {
+				require.Equal(t, 2, len(txManager.Votes()))
+				assert.Equal(t, voting.Prepared, txManager.Votes()[0].Phase)
+				assert.Equal(t, expectedVote, txManager.Votes()[1].Vote)
+				assert.Equal(t, voting.Committed, txManager.Votes()[1].Phase)
+			} else {
+				require.Equal(t, 0, len(txManager.Votes()))
+			}
+		})
 	}
 }
 
-func TestFailedRestoreCustomHooksDueToValidations(t *testing.T) {
+func TestSetCustomHooks_failedValidation(t *testing.T) {
 	t.Parallel()
 	testhelper.NewFeatureSets(featureflag.TransactionalRestoreCustomHooks).
-		Run(t, testFailedRestoreCustomHooksDueToValidations)
+		Run(t, testFailedSetCustomHooksDueToValidations)
 }
 
-func testFailedRestoreCustomHooksDueToValidations(t *testing.T, ctx context.Context) {
+func testFailedSetCustomHooksDueToValidations(t *testing.T, ctx context.Context) {
 	t.Parallel()
-	_, client := setupRepositoryServiceWithoutRepo(t)
 
-	stream, err := client.RestoreCustomHooks(ctx)
-	require.NoError(t, err)
+	for _, tc := range []struct {
+		desc         string
+		streamSender func(*testing.T, context.Context, gitalypb.RepositoryServiceClient) error
+	}{
+		{
+			desc: "SetCustomHooks",
+			streamSender: func(t *testing.T, ctx context.Context, client gitalypb.RepositoryServiceClient) error {
+				stream, err := client.SetCustomHooks(ctx)
+				require.NoError(t, err)
+
+				require.NoError(t, stream.Send(&gitalypb.SetCustomHooksRequest{}))
+
+				_, err = stream.CloseAndRecv()
+				return err
+			},
+		},
+		{
+			desc: "RestoreCustomHooks",
+			streamSender: func(t *testing.T, ctx context.Context, client gitalypb.RepositoryServiceClient) error {
+				stream, err := client.RestoreCustomHooks(ctx)
+				require.NoError(t, err)
 
-	require.NoError(t, stream.Send(&gitalypb.RestoreCustomHooksRequest{}))
+				require.NoError(t, stream.Send(&gitalypb.RestoreCustomHooksRequest{}))
 
-	_, err = stream.CloseAndRecv()
-	testhelper.RequireGrpcError(t, err, status.Error(codes.InvalidArgument, testhelper.GitalyOrPraefect(
-		"empty Repository",
-		"repo scoped: empty Repository",
-	)))
+				_, err = stream.CloseAndRecv()
+				return err
+			},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			_, client := setupRepositoryServiceWithoutRepo(t)
+
+			err := tc.streamSender(t, ctx, client)
+			testhelper.RequireGrpcError(t, err, status.Error(codes.InvalidArgument, testhelper.GitalyOrPraefect(
+				"empty Repository",
+				"repo scoped: empty Repository",
+			)))
+		})
+	}
 }
 
-func TestFailedRestoreCustomHooksDueToBadTar(t *testing.T) {
+func TestSetCustomHooks_corruptTar(t *testing.T) {
 	t.Parallel()
 	testhelper.NewFeatureSets(featureflag.TransactionalRestoreCustomHooks).
-		Run(t, testFailedRestoreCustomHooksDueToBadTar)
+		Run(t, testFailedSetCustomHooksDueToBadTar)
 }
 
-func testFailedRestoreCustomHooksDueToBadTar(t *testing.T, ctx context.Context) {
-	_, repo, _, client := setupRepositoryService(t, ctx)
-
-	stream, err := client.RestoreCustomHooks(ctx)
-
-	require.NoError(t, err)
-
-	request := &gitalypb.RestoreCustomHooksRequest{Repository: repo}
-	writer := streamio.NewWriter(func(p []byte) error {
-		request.Data = p
-		if err := stream.Send(request); err != nil {
-			return err
-		}
-
-		request = &gitalypb.RestoreCustomHooksRequest{}
-		return nil
-	})
+func testFailedSetCustomHooksDueToBadTar(t *testing.T, ctx context.Context) {
+	t.Parallel()
 
-	file, err := os.Open("testdata/corrupted_hooks.tar")
-	require.NoError(t, err)
-	defer file.Close()
+	for _, tc := range []struct {
+		desc         string
+		streamWriter func(*testing.T, context.Context, *gitalypb.Repository, gitalypb.RepositoryServiceClient) (io.Writer, func() error)
+	}{
+		{
+			desc: "SetCustomHooks",
+			streamWriter: func(t *testing.T, ctx context.Context, repo *gitalypb.Repository, client gitalypb.RepositoryServiceClient) (io.Writer, func() error) {
+				stream, err := client.SetCustomHooks(ctx)
+				require.NoError(t, err)
+
+				request := &gitalypb.SetCustomHooksRequest{Repository: repo}
+				writer := streamio.NewWriter(func(p []byte) error {
+					request.Data = p
+					if err := stream.Send(request); err != nil {
+						return err
+					}
+
+					request = &gitalypb.SetCustomHooksRequest{}
+					return nil
+				})
+
+				closeFunc := func() error {
+					_, err = stream.CloseAndRecv()
+					return err
+				}
+
+				return writer, closeFunc
+			},
+		},
+		{
+			desc: "RestoreCustomHooks",
+			streamWriter: func(t *testing.T, ctx context.Context, repo *gitalypb.Repository, client gitalypb.RepositoryServiceClient) (io.Writer, func() error) {
+				stream, err := client.RestoreCustomHooks(ctx)
+				require.NoError(t, err)
+
+				request := &gitalypb.RestoreCustomHooksRequest{Repository: repo}
+				writer := streamio.NewWriter(func(p []byte) error {
+					request.Data = p
+					if err := stream.Send(request); err != nil {
+						return err
+					}
+
+					request = &gitalypb.RestoreCustomHooksRequest{}
+					return nil
+				})
+
+				closeFunc := func() error {
+					_, err = stream.CloseAndRecv()
+					return err
+				}
+
+				return writer, closeFunc
+			},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			_, repo, _, client := setupRepositoryService(t, ctx)
+			writer, closeStream := tc.streamWriter(t, ctx, repo, client)
 
-	_, err = io.Copy(writer, file)
-	require.NoError(t, err)
-	_, err = stream.CloseAndRecv()
+			file, err := os.Open("testdata/corrupted_hooks.tar")
+			require.NoError(t, err)
+			defer testhelper.MustClose(t, file)
 
-	testhelper.RequireGrpcCode(t, err, codes.Internal)
+			_, err = io.Copy(writer, file)
+			require.NoError(t, err)
+			err = closeStream()
+			testhelper.RequireGrpcCode(t, err, codes.Internal)
+		})
+	}
 }
 
 type testFile struct {
diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index b2157a889..3a2c00497 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -99,6 +99,7 @@ var transactionRPCs = map[string]transactionsCondition{
 	// The `RestoreCustomHooks` RPC can be make transactional by enabling the
 	// `TransactionalRestoreCustomHooks` feature flag.
 	"/gitaly.RepositoryService/RestoreCustomHooks": transactionsFlag(featureflag.TransactionalRestoreCustomHooks),
+	"/gitaly.RepositoryService/SetCustomHooks":     transactionsFlag(featureflag.TransactionalRestoreCustomHooks),
 }
 
 // forcePrimaryRoutingRPCs tracks RPCs which need to always get routed to the primary. This should