Geo: Add `done` pktline when git clone --depth option is given

internal/command/githttp/pull.go

@@ -4,12 +4,13 @@ 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"
-	"io"
 )
 
 const pullService = "git-upload-pack"
@@ -40,7 +41,7 @@ func (c *PullCommand) Execute(ctx context.Context) error {
 		return err
 	}
 
-	return c.requestUploadPack(ctx, client)
+	return c.requestUploadPack(ctx, client, data.GeoProxyFetchDirectToPrimaryWithOptions)
 }
 
 func (c *PullCommand) requestInfoRefs(ctx context.Context, client *git.Client) error {
@@ -63,9 +64,9 @@ func (c *PullCommand) requestInfoRefs(ctx context.Context, client *git.Client) e
 	return err
 }
 
-func (c *PullCommand) requestUploadPack(ctx context.Context, client *git.Client) error {
+func (c *PullCommand) requestUploadPack(ctx context.Context, client *git.Client, geoProxyFetchDirectToPrimaryWithOptions bool) error {
 	pipeReader, pipeWriter := io.Pipe()
-	go c.readFromStdin(pipeWriter)
+	go c.readFromStdin(pipeWriter, geoProxyFetchDirectToPrimaryWithOptions)
 
 	response, err := client.UploadPack(ctx, pipeReader)
 	if err != nil {
@@ -78,19 +79,23 @@ func (c *PullCommand) requestUploadPack(ctx context.Context, client *git.Client)
 	return err
 }
 
-func (c *PullCommand) readFromStdin(pw *io.PipeWriter) {
+func (c *PullCommand) readFromStdin(pw *io.PipeWriter, geoProxyFetchDirectToPrimaryWithOptions bool) {
 	scanner := pktline.NewScanner(c.ReadWriter.In)
 
 	for scanner.Scan() {
 		line := scanner.Bytes()
 
-		if pktline.IsDone(line) {
-			pw.Write(line)
+		pw.Write(line)
 
+		if pktline.IsDone(line) {
 			break
 		}
 
-		pw.Write(line)
+		if pktline.IsFlush(line) && geoProxyFetchDirectToPrimaryWithOptions {
+			pw.Write(pktline.PktDone())
+
+			break
+		}
 	}
 
 	pw.Close()

internal/command/githttp/pull_test.go

@@ -18,13 +18,12 @@ import (
 
 var cloneResponse = `0090want 11d731b83788cd556abea7b465c6bee52d89923c multi_ack_detailed side-band-64k thin-pack ofs-delta deepen-since deepen-not agent=git/2.41.0
 0032want e56497bb5f03a90a51293fc6d516788730953899
-00000009done
 `
 
 func TestPullExecute(t *testing.T) {
 	url := setupPull(t, http.StatusOK)
 	output := &bytes.Buffer{}
-	input := strings.NewReader(cloneResponse)
+	input := strings.NewReader(cloneResponse + "00000009done\n")
 
 	cmd := &PullCommand{
 		Config:     &config.Config{GitlabUrl: url},
@@ -40,6 +39,25 @@ func TestPullExecute(t *testing.T) {
 	require.Equal(t, infoRefsWithoutPrefix, output.String())
 }
 
+func TestPullExecuteWithDepth(t *testing.T) {
+	url := setupPull(t, http.StatusOK)
+	output := &bytes.Buffer{}
+	input := strings.NewReader(cloneResponse + "0000\n")
+
+	cmd := &PullCommand{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: output, In: input},
+		Response: &accessverifier.Response{
+			Payload: accessverifier.CustomPayload{
+				Data: accessverifier.CustomPayloadData{PrimaryRepo: url, GeoProxyFetchDirectToPrimaryWithOptions: true},
+			},
+		},
+	}
+
+	require.NoError(t, cmd.Execute(context.Background()))
+	require.Equal(t, infoRefsWithoutPrefix, output.String())
+}
+
 func TestPullExecuteWithFailedInfoRefs(t *testing.T) {
 	testCases := []struct {
 		desc            string
@@ -94,7 +112,7 @@ func TestPullExecuteWithFailedInfoRefs(t *testing.T) {
 func TestExecuteWithFailedUploadPack(t *testing.T) {
 	url := setupPull(t, http.StatusForbidden)
 	output := &bytes.Buffer{}
-	input := strings.NewReader(cloneResponse)
+	input := strings.NewReader(cloneResponse + "00000009done\n")
 
 	cmd := &PullCommand{
 		Config:     &config.Config{GitlabUrl: url},

internal/gitlabnet/accessverifier/client.go

@@ -43,13 +43,14 @@ 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"`
-	RequestHeaders               map[string]string `json:"request_headers"`
-	GeoProxyDirectToPrimary      bool              `json:"geo_proxy_direct_to_primary"`
-	GeoProxyFetchDirectToPrimary bool              `json:"geo_proxy_fetch_direct_to_primary"`
+	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"`
+	GeoProxyFetchDirectToPrimary            bool              `json:"geo_proxy_fetch_direct_to_primary"`
+	GeoProxyFetchDirectToPrimaryWithOptions bool              `json:"geo_proxy_fetch_direct_to_primary_with_options"`
 }
 
 type CustomPayload struct {