Move go code up one level

2 files changed, 138 insertions, 0 deletions

diff --git a/internal/handler/exec.go b/internal/handler/exec.go
new file mode 100644
index 0000000..1f3177d
--- /dev/null
+++ b/internal/handler/exec.go
@@ -0,0 +1,96 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"gitlab.com/gitlab-org/gitaly/auth"
+	"gitlab.com/gitlab-org/gitaly/client"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/labkit/tracing"
+	"google.golang.org/grpc"
+)
+
+// GitalyHandlerFunc implementations are responsible for making
+// an appropriate Gitaly call using the provided client and context
+// and returning an error from the Gitaly call.
+type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn) (int32, error)
+
+type GitalyConn struct {
+	ctx   context.Context
+	conn  *grpc.ClientConn
+	close func()
+}
+
+type GitalyCommand struct {
+	Config      *config.Config
+	ServiceName string
+	Address     string
+	Token       string
+}
+
+// RunGitalyCommand provides a bootstrap for Gitaly commands executed
+// through GitLab-Shell. It ensures that logging, tracing and other
+// common concerns are configured before executing the `handler`.
+func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error {
+	gitalyConn, err := getConn(gc)
+
+	if err != nil {
+		return err
+	}
+
+	_, err = handler(gitalyConn.ctx, gitalyConn.conn)
+
+	gitalyConn.close()
+
+	return err
+}
+
+func getConn(gc *GitalyCommand) (*GitalyConn, error) {
+	if gc.Address == "" {
+		return nil, fmt.Errorf("no gitaly_address given")
+	}
+
+	connOpts := client.DefaultDialOpts
+	if gc.Token != "" {
+		connOpts = append(client.DefaultDialOpts, grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(gc.Token)))
+	}
+
+	// Use a working directory that won't get removed or unmounted.
+	if err := os.Chdir("/"); err != nil {
+		return nil, err
+	}
+
+	// Configure distributed tracing
+	serviceName := fmt.Sprintf("gitlab-shell-%v", gc.ServiceName)
+	closer := tracing.Initialize(
+		tracing.WithServiceName(serviceName),
+
+		// For GitLab-Shell, we explicitly initialize tracing from a config file
+		// instead of the default environment variable (using GITLAB_TRACING)
+		// This decision was made owing to the difficulty in passing environment
+		// variables into GitLab-Shell processes.
+		//
+		// Processes are spawned as children of the SSH daemon, which tightly
+		// controls environment variables; doing this means we don't have to
+		// enable PermitUserEnvironment
+		tracing.WithConnectionString(gc.Config.GitlabTracing),
+	)
+
+	ctx, finished := tracing.ExtractFromEnv(context.Background())
+
+	conn, err := client.Dial(gc.Address, connOpts)
+	if err != nil {
+		return nil, err
+	}
+
+	finish := func() {
+		finished()
+		closer.Close()
+		conn.Close()
+	}
+
+	return &GitalyConn{ctx: ctx, conn: conn, close: finish}, nil
+}
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
new file mode 100644
index 0000000..6c7d3f5
--- /dev/null
+++ b/internal/handler/exec_test.go
@@ -0,0 +1,42 @@
+package handler
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
+)
+
+func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
+	return func(ctx context.Context, client *grpc.ClientConn) (int32, error) {
+		require.NotNil(t, ctx)
+		require.NotNil(t, client)
+
+		return 0, err
+	}
+}
+
+func TestRunGitalyCommand(t *testing.T) {
+	cmd := GitalyCommand{
+		Config:  &config.Config{},
+		Address: "tcp://localhost:9999",
+	}
+
+	err := cmd.RunGitalyCommand(makeHandler(t, nil))
+	require.NoError(t, err)
+
+	expectedErr := errors.New("error")
+	err = cmd.RunGitalyCommand(makeHandler(t, expectedErr))
+	require.Equal(t, err, expectedErr)
+}
+
+func TestMissingGitalyAddress(t *testing.T) {
+	cmd := GitalyCommand{Config: &config.Config{}}
+
+	err := cmd.RunGitalyCommand(makeHandler(t, nil))
+	require.EqualError(t, err, "no gitaly_address given")
+}