Perform HTTP request to primary on Geo push

Currently, we perform a request to Gitlab Rails that proxies
the request to primary

However, it causes timeouts on big pushes and consumes large
amount of memory. We can perform an HTTP request directly
from Gitlab Shell instead and stream the response to the user

9 files changed, 431 insertions, 7 deletions

diff --git a/internal/command/githttp/push.go b/internal/command/githttp/push.go
new file mode 100644
index 0000000..3377baf
--- /dev/null
+++ b/internal/command/githttp/push.go
@@ -0,0 +1,106 @@
+package githttp
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/accessverifier"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/git"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/pktline"
+)
+
+const service = "git-receive-pack"
+
+var receivePackHttpPrefix = []byte("001f# service=git-receive-pack\n0000")
+
+type PushCommand struct {
+	Config     *config.Config
+	ReadWriter *readwriter.ReadWriter
+	Response   *accessverifier.Response
+}
+
+// See Uploading Data > HTTP(S) section at:
+// https://git-scm.com/book/en/v2/Git-Internals-Transfer-Protocols
+//
+// 1. Perform /info/refs?service=git-receive-pack request
+// 2. Remove the header to make it consumable by SSH protocol
+// 3. Send the result to the user via SSH (writeToStdout)
+// 4. Read the send-pack data provided by user via SSH (stdinReader)
+// 5. Perform /git-receive-pack request and send this data
+// 6. Return the output to the user
+func (c *PushCommand) Execute(ctx context.Context) error {
+	data := c.Response.Payload.Data
+	client, err := git.NewClient(c.Config, data.PrimaryRepo, data.RequestHeaders)
+	if err != nil {
+		return err
+	}
+
+	if err := c.requestInfoRefs(ctx, client); err != nil {
+		return err
+	}
+
+	return c.requestReceivePack(ctx, client)
+}
+
+func (c *PushCommand) requestInfoRefs(ctx context.Context, client *git.Client) error {
+	response, err := client.InfoRefs(ctx, service)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+
+	// Read the first bytes that contain 001f# service=git-receive-pack\n0000 string
+	// to convert HTTP(S) Git response to the one expected by SSH
+	p := make([]byte, len(receivePackHttpPrefix))
+	_, err = response.Body.Read(p)
+	if err != nil || !bytes.Equal(p, receivePackHttpPrefix) {
+		return fmt.Errorf("Unexpected git-receive-pack response")
+	}
+
+	_, err = io.Copy(c.ReadWriter.Out, response.Body)
+
+	return err
+}
+
+func (c *PushCommand) requestReceivePack(ctx context.Context, client *git.Client) error {
+	pipeReader, pipeWriter := io.Pipe()
+	go c.readFromStdin(pipeWriter)
+
+	response, err := client.ReceivePack(ctx, pipeReader)
+	if err != nil {
+		return err
+	}
+	defer response.Body.Close()
+
+	_, err = io.Copy(c.ReadWriter.Out, response.Body)
+
+	return err
+}
+
+func (c *PushCommand) readFromStdin(pw *io.PipeWriter) {
+	var needsPackData bool
+
+	scanner := pktline.NewScanner(c.ReadWriter.In)
+	for scanner.Scan() {
+		line := scanner.Bytes()
+		pw.Write(line)
+
+		if pktline.IsFlush(line) {
+			break
+		}
+
+		if !needsPackData && !pktline.IsRefRemoval(line) {
+			needsPackData = true
+		}
+	}
+
+	if needsPackData {
+		io.Copy(pw, c.ReadWriter.In)
+	}
+
+	pw.Close()
+}
diff --git a/internal/command/githttp/push_test.go b/internal/command/githttp/push_test.go
new file mode 100644
index 0000000..1363abd
--- /dev/null
+++ b/internal/command/githttp/push_test.go
@@ -0,0 +1,150 @@
+package githttp
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/accessverifier"
+)
+
+var (
+	flush                 = "0000"
+	infoRefsWithoutPrefix = "00c4e56497bb5f03a90a51293fc6d516788730953899 refs/heads/'test'report-status " +
+		"report-status-v2 delete-refs side-band-64k quiet atomic ofs-delta push-options object-format=sha1 " +
+		"agent=git/2.38.3.gl200\n" + flush
+)
+
+func TestExecute(t *testing.T) {
+	url, input := setup(t, http.StatusOK)
+	output := &bytes.Buffer{}
+
+	cmd := &PushCommand{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
+		Response: &accessverifier.Response{
+			Payload: accessverifier.CustomPayload{
+				Data: accessverifier.CustomPayloadData{PrimaryRepo: url},
+			},
+		},
+	}
+
+	require.NoError(t, cmd.Execute(context.Background()))
+	require.Equal(t, infoRefsWithoutPrefix, output.String())
+}
+
+func TestExecuteWithFailedInfoRefs(t *testing.T) {
+	testCases := []struct {
+		desc            string
+		statusCode      int
+		responseContent string
+		expectedErr     string
+	}{
+		{
+			desc:        "request failed",
+			statusCode:  http.StatusForbidden,
+			expectedErr: "Internal API error (403)",
+		}, {
+			desc:            "unexpected response",
+			statusCode:      http.StatusOK,
+			responseContent: "unexpected response",
+			expectedErr:     "Unexpected git-receive-pack response",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			requests := []testserver.TestRequestHandler{
+				{
+					Path: "/info/refs",
+					Handler: func(w http.ResponseWriter, r *http.Request) {
+						require.Equal(t, "git-receive-pack", r.URL.Query().Get("service"))
+
+						w.WriteHeader(tc.statusCode)
+						w.Write([]byte(tc.responseContent))
+					},
+				},
+			}
+
+			url := testserver.StartHttpServer(t, requests)
+
+			cmd := &PushCommand{
+				Config: &config.Config{GitlabUrl: url},
+				Response: &accessverifier.Response{
+					Payload: accessverifier.CustomPayload{
+						Data: accessverifier.CustomPayloadData{PrimaryRepo: url},
+					},
+				},
+			}
+
+			err := cmd.Execute(context.Background())
+			require.Error(t, err)
+			require.Equal(t, tc.expectedErr, err.Error())
+		})
+	}
+}
+
+func TestExecuteWithFailedReceivePack(t *testing.T) {
+	url, input := setup(t, http.StatusForbidden)
+	output := &bytes.Buffer{}
+
+	cmd := &PushCommand{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
+		Response: &accessverifier.Response{
+			Payload: accessverifier.CustomPayload{
+				Data: accessverifier.CustomPayloadData{PrimaryRepo: url},
+			},
+		},
+	}
+
+	err := cmd.Execute(context.Background())
+	require.Error(t, err)
+	require.Equal(t, "Internal API error (403)", err.Error())
+}
+
+func setup(t *testing.T, receivePackStatusCode int) (string, io.Reader) {
+	infoRefs := "001f# service=git-receive-pack\n" + flush + infoRefsWithoutPrefix
+	receivePackPrefix := "00ab4c9d98d7750fa65db8ddcc60a89ef919f7a179f9 df505c066e4e63a801268a84627d7e8f7e033c7a " +
+		"refs/heads/main123 report-status-v2 side-band-64k object-format=sha1 agent=git/2.39.1"
+	receivePackData := "PACK some data"
+
+	// Imitate sending data via multiple packets
+	input := io.MultiReader(
+		strings.NewReader(receivePackPrefix),
+		strings.NewReader(flush),
+		strings.NewReader(receivePackData),
+	)
+
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/info/refs",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, "git-receive-pack", r.URL.Query().Get("service"))
+
+				w.Write([]byte(infoRefs))
+			},
+		},
+		{
+			Path: "/git-receive-pack",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				body, err := io.ReadAll(r.Body)
+				require.NoError(t, err)
+				defer r.Body.Close()
+
+				require.Equal(t, receivePackPrefix+flush+receivePackData, string(body))
+				w.WriteHeader(receivePackStatusCode)
+			},
+		},
+	}
+
+	return testserver.StartHttpServer(t, requests), input
+}
diff --git a/internal/command/receivepack/receivepack.go b/internal/command/receivepack/receivepack.go
index 976ab39..c9ef7cd 100644
--- a/internal/command/receivepack/receivepack.go
+++ b/internal/command/receivepack/receivepack.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/githttp"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/shared/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/shared/customaction"
@@ -30,6 +31,20 @@ func (c *Command) Execute(ctx context.Context) error {
 	}
 
 	if response.IsCustomAction() {
+		// When `geo_proxy_direct_to_primary` feature flag is enabled, a Git over HTTP direct request
+		// to primary repo is performed instead of proxying the request through Gitlab Rails.
+		// After the feature flag is enabled by default and removed,
+		// custom action functionality will be removed along with it.
+		if response.Payload.Data.GeoProxyDirectToPrimary {
+			cmd := githttp.PushCommand{
+				Config:     c.Config,
+				ReadWriter: c.ReadWriter,
+				Response:   response,
+			}
+
+			return cmd.Execute(ctx)
+		}
+
 		customAction := customaction.Command{
 			Config:     c.Config,
 			ReadWriter: c.ReadWriter,
diff --git a/internal/command/shared/customaction/customaction.go b/internal/command/shared/customaction/customaction.go
index c12d685..e8b5ede 100644
--- a/internal/command/shared/customaction/customaction.go
+++ b/internal/command/shared/customaction/customaction.go
@@ -34,6 +34,9 @@ type Command struct {
 	EOFSent    bool
 }
 
+// When `geo_proxy_direct_to_primary` feature flag is enabled, a Git over HTTP direct request
+// to primary repo is performed instead of proxying the request through Gitlab Rails.
+// After the feature flag is enabled by default and removed, this package will be removed along with it.
 func (c *Command) Execute(ctx context.Context, response *accessverifier.Response) error {
 	data := response.Payload.Data
 	apiEndpoints := data.ApiEndpoints
diff --git a/internal/gitlabnet/accessverifier/client.go b/internal/gitlabnet/accessverifier/client.go
index 8554d9b..b2383dc 100644
--- a/internal/gitlabnet/accessverifier/client.go
+++ b/internal/gitlabnet/accessverifier/client.go
@@ -40,10 +40,12 @@ type Gitaly struct {
 }
 
 type CustomPayloadData struct {
-	ApiEndpoints []string `json:"api_endpoints"`
-	Username     string   `json:"gl_username"`
-	PrimaryRepo  string   `json:"primary_repo"`
-	UserId       string   `json:"gl_id,omitempty"`
+	ApiEndpoints            []string          `json:"api_endpoints"`
+	Username                string            `json:"gl_username"`
+	PrimaryRepo             string            `json:"primary_repo"`
+	UserId                  string            `json:"gl_id,omitempty"`
+	RequestHeaders          map[string]string `json:"request_headers"`
+	GeoProxyDirectToPrimary bool              `json:"geo_proxy_direct_to_primary"`
 }
 
 type CustomPayload struct {
diff --git a/internal/gitlabnet/accessverifier/client_test.go b/internal/gitlabnet/accessverifier/client_test.go
index f2c88a5..e203200 100644
--- a/internal/gitlabnet/accessverifier/client_test.go
+++ b/internal/gitlabnet/accessverifier/client_test.go
@@ -107,9 +107,11 @@ func TestGeoPushGetCustomAction(t *testing.T) {
 	response.Payload = CustomPayload{
 		Action: "geo_proxy_to_primary",
 		Data: CustomPayloadData{
-			ApiEndpoints: []string{"geo/proxy_git_ssh/info_refs_receive_pack", "geo/proxy_git_ssh/receive_pack"},
-			Username:     "custom",
-			PrimaryRepo:  "https://repo/path",
+			ApiEndpoints:            []string{"geo/proxy_git_ssh/info_refs_receive_pack", "geo/proxy_git_ssh/receive_pack"},
+			GeoProxyDirectToPrimary: true,
+			RequestHeaders:          map[string]string{"Authorization": "Bearer token"},
+			Username:                "custom",
+			PrimaryRepo:             "https://repo/path",
 		},
 	}
 	response.StatusCode = 300
diff --git a/internal/gitlabnet/git/client.go b/internal/gitlabnet/git/client.go
new file mode 100644
index 0000000..db71e3f
--- /dev/null
+++ b/internal/gitlabnet/git/client.go
@@ -0,0 +1,56 @@
+package git
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+
+	"gitlab.com/gitlab-org/gitlab-shell/v14/client"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet"
+)
+
+type Client struct {
+	url     string
+	headers map[string]string
+	client  *client.GitlabNetClient
+}
+
+func NewClient(cfg *config.Config, url string, headers map[string]string) (*Client, error) {
+	client, err := gitlabnet.GetClient(cfg)
+	if err != nil {
+		return nil, fmt.Errorf("Error creating http client: %v", err)
+	}
+
+	return &Client{client: client, headers: headers, url: url}, nil
+}
+
+func (c *Client) InfoRefs(ctx context.Context, service string) (*http.Response, error) {
+	request, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url+"/info/refs?service="+service, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return c.do(request)
+}
+
+func (c *Client) ReceivePack(ctx context.Context, body io.Reader) (*http.Response, error) {
+	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url+"/git-receive-pack", body)
+	if err != nil {
+		return nil, err
+	}
+	request.Header.Add("Content-Type", "application/x-git-receive-pack-request")
+	request.Header.Add("Accept", "application/x-git-receive-pack-result")
+
+	return c.do(request)
+}
+
+func (c *Client) do(request *http.Request) (*http.Response, error) {
+
+	for k, v := range c.headers {
+		request.Header.Add(k, v)
+	}
+
+	return c.client.Do(request)
+}
diff --git a/internal/gitlabnet/git/client_test.go b/internal/gitlabnet/git/client_test.go
new file mode 100644
index 0000000..b550e99
--- /dev/null
+++ b/internal/gitlabnet/git/client_test.go
@@ -0,0 +1,88 @@
+package git
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/client/testserver"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+)
+
+var customHeaders = map[string]string{
+	"Authorization": "Bearer: token",
+	"Header-One":    "Value-Two",
+}
+
+func TestInfoRefs(t *testing.T) {
+	client := setup(t)
+
+	for _, service := range []string{
+		"git-receive-pack",
+		"git-upload-pack",
+		"git-archive-pack",
+	} {
+		response, err := client.InfoRefs(context.Background(), service)
+		require.NoError(t, err)
+
+		body, err := io.ReadAll(response.Body)
+		require.NoError(t, err)
+		defer response.Body.Close()
+
+		require.Equal(t, service, string(body))
+	}
+}
+
+func TestReceivePack(t *testing.T) {
+	client := setup(t)
+
+	content := "content"
+	response, err := client.ReceivePack(context.Background(), bytes.NewReader([]byte(content)))
+	require.NoError(t, err)
+	defer response.Body.Close()
+
+	body, err := io.ReadAll(response.Body)
+	require.NoError(t, err)
+
+	require.Equal(t, "git-receive-pack: content", string(body))
+}
+
+func setup(t *testing.T) *Client {
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/info/refs",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, customHeaders["Authorization"], r.Header.Get("Authorization"))
+				require.Equal(t, customHeaders["Header-One"], r.Header.Get("Header-One"))
+
+				w.Write([]byte(r.URL.Query().Get("service")))
+			},
+		},
+		{
+			Path: "/git-receive-pack",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, customHeaders["Authorization"], r.Header.Get("Authorization"))
+				require.Equal(t, customHeaders["Header-One"], r.Header.Get("Header-One"))
+				require.Equal(t, "application/x-git-receive-pack-request", r.Header.Get("Content-Type"))
+				require.Equal(t, "application/x-git-receive-pack-result", r.Header.Get("Accept"))
+				require.Equal(t, customHeaders["Header-One"], r.Header.Get("Header-One"))
+
+				body, err := io.ReadAll(r.Body)
+				require.NoError(t, err)
+				defer r.Body.Close()
+
+				w.Write([]byte("git-receive-pack: "))
+				w.Write(body)
+			},
+		},
+	}
+
+	url := testserver.StartHttpServer(t, requests)
+	client, err := NewClient(&config.Config{GitlabUrl: url}, url, customHeaders)
+	require.NoError(t, err)
+
+	return client
+}
diff --git a/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json b/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json
index 318d42d..adedda4 100644
--- a/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json
+++ b/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json
@@ -21,6 +21,8 @@
     "action": "geo_proxy_to_primary",
 		"data": {
 			"api_endpoints": ["geo/proxy_git_ssh/info_refs_receive_pack", "geo/proxy_git_ssh/receive_pack"],
+			"geo_proxy_direct_to_primary": true,
+			"request_headers": { "Authorization": "Bearer token" },
 			"gl_username":   "custom",
 			"primary_repo":  "https://repo/path"
     }