Merge branch '202037-geo-ssh-clone-pull-redirect-to-primary-when-selective-sync-enabled-and-project-not-selected' into 'master'

Geo: Add custom action support for clone/pull

See merge request gitlab-org/gitlab-shell!369

11 files changed, 502 insertions, 27 deletions

diff --git a/internal/command/receivepack/receivepack.go b/internal/command/receivepack/receivepack.go
index 3af3941..7271264 100644
--- a/internal/command/receivepack/receivepack.go
+++ b/internal/command/receivepack/receivepack.go
@@ -28,7 +28,11 @@ func (c *Command) Execute() error {
 	}
 
 	if response.IsCustomAction() {
-		customAction := customaction.Command{c.Config, c.ReadWriter}
+		customAction := customaction.Command{
+			Config:     c.Config,
+			ReadWriter: c.ReadWriter,
+			EOFSent:    true,
+		}
 		return customAction.Execute(response)
 	}
 
diff --git a/internal/command/shared/customaction/customaction.go b/internal/command/shared/customaction/customaction.go
index c4b6647..801ad63 100644
--- a/internal/command/shared/customaction/customaction.go
+++ b/internal/command/shared/customaction/customaction.go
@@ -5,13 +5,14 @@ import (
 	"errors"
 
 	"io"
-	"io/ioutil"
 	"net/http"
 
+	log "github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/pktline"
 )
 
 type Request struct {
@@ -28,6 +29,7 @@ type Response struct {
 type Command struct {
 	Config     *config.Config
 	ReadWriter *readwriter.ReadWriter
+	EOFSent    bool
 }
 
 func (c *Command) Execute(response *accessverifier.Response) error {
@@ -53,21 +55,38 @@ func (c *Command) processApiEndpoints(response *accessverifier.Response) error {
 	request.Data.UserId = response.Who
 
 	for _, endpoint := range data.ApiEndpoints {
+		fields := log.Fields{
+			"primary_repo": data.PrimaryRepo,
+			"endpoint":     endpoint,
+		}
+
+		log.WithFields(fields).Info("Performing custom action")
+
 		response, err := c.performRequest(client, endpoint, request)
 		if err != nil {
 			return err
 		}
 
+		// Print to os.Stdout the result contained in the response
+		//
 		if err = c.displayResult(response.Result); err != nil {
 			return err
 		}
 
 		// In the context of the git push sequence of events, it's necessary to read
 		// stdin in order to capture output to pass onto subsequent commands
-		output, err := ioutil.ReadAll(c.ReadWriter.In)
-		if err != nil {
-			return err
+		//
+		var output []byte
+
+		if c.EOFSent {
+			output, err = c.readFromStdin()
+			if err != nil {
+				return err
+			}
+		} else {
+			output = c.readFromStdinNoEOF()
 		}
+
 		request.Output = output
 	}
 
@@ -89,6 +108,29 @@ func (c *Command) performRequest(client *gitlabnet.GitlabClient, endpoint string
 	return cr, nil
 }
 
+func (c *Command) readFromStdin() ([]byte, error) {
+	output := new(bytes.Buffer)
+	_, err := io.Copy(output, c.ReadWriter.In)
+
+	return output.Bytes(), err
+}
+
+func (c *Command) readFromStdinNoEOF() []byte {
+	var output []byte
+
+	scanner := pktline.NewScanner(c.ReadWriter.In)
+	for scanner.Scan() {
+		line := scanner.Bytes()
+		output = append(output, line...)
+
+		if pktline.IsDone(line) {
+			break
+		}
+	}
+
+	return output
+}
+
 func (c *Command) displayResult(result []byte) error {
 	_, err := io.Copy(c.ReadWriter.Out, bytes.NewReader(result))
 	return err
diff --git a/internal/command/shared/customaction/customaction_test.go b/internal/command/shared/customaction/customaction_test.go
index 3dfe288..31044f9 100644
--- a/internal/command/shared/customaction/customaction_test.go
+++ b/internal/command/shared/customaction/customaction_test.go
@@ -15,12 +15,12 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/testserver"
 )
 
-func TestExecute(t *testing.T) {
+func TestExecuteEOFSent(t *testing.T) {
 	who := "key-1"
 
 	requests := []testserver.TestRequestHandler{
 		{
-			Path: "/geo/proxy/info_refs",
+			Path: "/geo/proxy/info_refs_receive_pack",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				b, err := ioutil.ReadAll(r.Body)
 				require.NoError(t, err)
@@ -36,7 +36,7 @@ func TestExecute(t *testing.T) {
 			},
 		},
 		{
-			Path: "/geo/proxy/push",
+			Path: "/geo/proxy/receive_pack",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				b, err := ioutil.ReadAll(r.Body)
 				require.NoError(t, err)
@@ -65,7 +65,7 @@ func TestExecute(t *testing.T) {
 		Payload: accessverifier.CustomPayload{
 			Action: "geo_proxy_to_primary",
 			Data: accessverifier.CustomPayloadData{
-				ApiEndpoints: []string{"/geo/proxy/info_refs", "/geo/proxy/push"},
+				ApiEndpoints: []string{"/geo/proxy/info_refs_receive_pack", "/geo/proxy/receive_pack"},
 				Username:     "custom",
 				PrimaryRepo:  "https://repo/path",
 			},
@@ -75,6 +75,77 @@ func TestExecute(t *testing.T) {
 	cmd := &Command{
 		Config:     &config.Config{GitlabUrl: url},
 		ReadWriter: &readwriter.ReadWriter{ErrOut: errBuf, Out: outBuf, In: input},
+		EOFSent:    true,
+	}
+
+	require.NoError(t, cmd.Execute(response))
+
+	// expect printing of info message, "custom" string from the first request
+	// and "output" string from the second request
+	require.Equal(t, "customoutput", outBuf.String())
+}
+
+func TestExecuteNoEOFSent(t *testing.T) {
+	who := "key-1"
+
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/geo/proxy/info_refs_upload_pack",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, err := ioutil.ReadAll(r.Body)
+				require.NoError(t, err)
+
+				var request *Request
+				require.NoError(t, json.Unmarshal(b, &request))
+
+				require.Equal(t, request.Data.UserId, who)
+				require.Empty(t, request.Output)
+
+				err = json.NewEncoder(w).Encode(Response{Result: []byte("custom")})
+				require.NoError(t, err)
+			},
+		},
+		{
+			Path: "/geo/proxy/upload_pack",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, err := ioutil.ReadAll(r.Body)
+				require.NoError(t, err)
+
+				var request *Request
+				require.NoError(t, json.Unmarshal(b, &request))
+
+				require.Equal(t, request.Data.UserId, who)
+				require.Equal(t, "0032want 343d70886785dc1f98aaf70f3b4ca87c93a5d0dd\n", string(request.Output))
+
+				err = json.NewEncoder(w).Encode(Response{Result: []byte("output")})
+				require.NoError(t, err)
+			},
+		},
+	}
+
+	url, cleanup := testserver.StartSocketHttpServer(t, requests)
+	defer cleanup()
+
+	outBuf := &bytes.Buffer{}
+	errBuf := &bytes.Buffer{}
+	input := bytes.NewBufferString("0032want 343d70886785dc1f98aaf70f3b4ca87c93a5d0dd\n")
+
+	response := &accessverifier.Response{
+		Who: who,
+		Payload: accessverifier.CustomPayload{
+			Action: "geo_proxy_to_primary",
+			Data: accessverifier.CustomPayloadData{
+				ApiEndpoints: []string{"/geo/proxy/info_refs_upload_pack", "/geo/proxy/upload_pack"},
+				Username:     "custom",
+				PrimaryRepo:  "https://repo/path",
+			},
+		},
+	}
+
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{ErrOut: errBuf, Out: outBuf, In: input},
+		EOFSent:    false,
 	}
 
 	require.NoError(t, cmd.Execute(response))
diff --git a/internal/command/uploadpack/uploadpack.go b/internal/command/uploadpack/uploadpack.go
index a5c71b2..56814d7 100644
--- a/internal/command/uploadpack/uploadpack.go
+++ b/internal/command/uploadpack/uploadpack.go
@@ -4,6 +4,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/shared/accessverifier"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/command/shared/customaction"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/shared/disallowedcommand"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 )
@@ -26,6 +27,15 @@ func (c *Command) Execute() error {
 		return err
 	}
 
+	if response.IsCustomAction() {
+		customAction := customaction.Command{
+			Config:     c.Config,
+			ReadWriter: c.ReadWriter,
+			EOFSent:    false,
+		}
+		return customAction.Execute(response)
+	}
+
 	return c.performGitalyCall(response)
 }
 
diff --git a/internal/gitlabnet/accessverifier/client_test.go b/internal/gitlabnet/accessverifier/client_test.go
index 0b6c7b0..009dcc0 100644
--- a/internal/gitlabnet/accessverifier/client_test.go
+++ b/internal/gitlabnet/accessverifier/client_test.go
@@ -18,8 +18,9 @@ import (
 )
 
 var (
-	repo   = "group/private"
-	action = commandargs.ReceivePack
+	repo              = "group/private"
+	receivePackAction = commandargs.ReceivePack
+	uploadPackAction  = commandargs.UploadPack
 )
 
 func buildExpectedResponse(who string) *Response {
@@ -52,7 +53,7 @@ func buildExpectedResponse(who string) *Response {
 }
 
 func TestSuccessfulResponses(t *testing.T) {
-	client, cleanup := setup(t)
+	client, cleanup := setup(t, "")
 	defer cleanup()
 
 	testCases := []struct {
@@ -73,7 +74,7 @@ func TestSuccessfulResponses(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			result, err := client.Verify(tc.args, action, repo)
+			result, err := client.Verify(tc.args, receivePackAction, repo)
 			require.NoError(t, err)
 
 			response := buildExpectedResponse(tc.who)
@@ -82,19 +83,42 @@ func TestSuccessfulResponses(t *testing.T) {
 	}
 }
 
-func TestGetCustomAction(t *testing.T) {
-	client, cleanup := setup(t)
+func TestGeoPushGetCustomAction(t *testing.T) {
+	client, cleanup := setup(t, "responses/allowed_with_push_payload.json")
 	defer cleanup()
 
 	args := &commandargs.Shell{GitlabUsername: "custom"}
-	result, err := client.Verify(args, action, repo)
+	result, err := client.Verify(args, receivePackAction, repo)
 	require.NoError(t, err)
 
 	response := buildExpectedResponse("user-1")
 	response.Payload = CustomPayload{
 		Action: "geo_proxy_to_primary",
 		Data: CustomPayloadData{
-			ApiEndpoints: []string{"geo/proxy_git_push_ssh/info_refs", "geo/proxy_git_push_ssh/push"},
+			ApiEndpoints: []string{"geo/proxy_git_ssh/info_refs_receive_pack", "geo/proxy_git_ssh/receive_pack"},
+			Username:     "custom",
+			PrimaryRepo:  "https://repo/path",
+		},
+	}
+	response.StatusCode = 300
+
+	require.True(t, response.IsCustomAction())
+	require.Equal(t, response, result)
+}
+
+func TestGeoPullGetCustomAction(t *testing.T) {
+	client, cleanup := setup(t, "responses/allowed_with_pull_payload.json")
+	defer cleanup()
+
+	args := &commandargs.Shell{GitlabUsername: "custom"}
+	result, err := client.Verify(args, uploadPackAction, repo)
+	require.NoError(t, err)
+
+	response := buildExpectedResponse("user-1")
+	response.Payload = CustomPayload{
+		Action: "geo_proxy_to_primary",
+		Data: CustomPayloadData{
+			ApiEndpoints: []string{"geo/proxy_git_ssh/info_refs_upload_pack", "geo/proxy_git_ssh/upload_pack"},
 			Username:     "custom",
 			PrimaryRepo:  "https://repo/path",
 		},
@@ -106,7 +130,7 @@ func TestGetCustomAction(t *testing.T) {
 }
 
 func TestErrorResponses(t *testing.T) {
-	client, cleanup := setup(t)
+	client, cleanup := setup(t, "")
 	defer cleanup()
 
 	testCases := []struct {
@@ -134,7 +158,7 @@ func TestErrorResponses(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			args := &commandargs.Shell{GitlabKeyId: tc.fakeId}
-			resp, err := client.Verify(args, action, repo)
+			resp, err := client.Verify(args, receivePackAction, repo)
 
 			require.EqualError(t, err, tc.expectedError)
 			require.Nil(t, resp)
@@ -142,7 +166,7 @@ func TestErrorResponses(t *testing.T) {
 	}
 }
 
-func setup(t *testing.T) (*Client, func()) {
+func setup(t *testing.T, allowedPayload string) (*Client, func()) {
 	testDirCleanup, err := testhelper.PrepareTestRootDir()
 	require.NoError(t, err)
 	defer testDirCleanup()
@@ -150,9 +174,13 @@ func setup(t *testing.T) (*Client, func()) {
 	body, err := ioutil.ReadFile(path.Join(testhelper.TestRoot, "responses/allowed.json"))
 	require.NoError(t, err)
 
-	allowedWithPayloadPath := path.Join(testhelper.TestRoot, "responses/allowed_with_payload.json")
-	bodyWithPayload, err := ioutil.ReadFile(allowedWithPayloadPath)
-	require.NoError(t, err)
+	var bodyWithPayload []byte
+
+	if allowedPayload != "" {
+		allowedWithPayloadPath := path.Join(testhelper.TestRoot, allowedPayload)
+		bodyWithPayload, err = ioutil.ReadFile(allowedWithPayloadPath)
+		require.NoError(t, err)
+	}
 
 	requests := []testserver.TestRequestHandler{
 		{
diff --git a/internal/pktline/pktline.go b/internal/pktline/pktline.go
new file mode 100644
index 0000000..35fceb2
--- /dev/null
+++ b/internal/pktline/pktline.go
@@ -0,0 +1,75 @@
+package pktline
+
+// Utility functions for working with the Git pkt-line format. See
+// https://github.com/git/git/blob/master/Documentation/technical/protocol-common.txt
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"io"
+	"strconv"
+)
+
+const (
+	maxPktSize = 0xffff
+	pktDelim   = "0001"
+)
+
+// NewScanner returns a bufio.Scanner that splits on Git pktline boundaries
+func NewScanner(r io.Reader) *bufio.Scanner {
+	scanner := bufio.NewScanner(r)
+	scanner.Buffer(make([]byte, maxPktSize), maxPktSize)
+	scanner.Split(pktLineSplitter)
+	return scanner
+}
+
+// IsDone detects the special flush packet '0009done\n'
+func IsDone(pkt []byte) bool {
+	return bytes.Equal(pkt, PktDone())
+}
+
+// PktDone returns the bytes for a "done" packet.
+func PktDone() []byte {
+	return []byte("0009done\n")
+}
+
+func pktLineSplitter(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	if len(data) < 4 {
+		if atEOF && len(data) > 0 {
+			return 0, nil, fmt.Errorf("pktLineSplitter: incomplete length prefix on %q", data)
+		}
+		return 0, nil, nil // want more data
+	}
+
+	// We have at least 4 bytes available so we can decode the 4-hex digit
+	// length prefix of the packet line.
+	pktLength64, err := strconv.ParseInt(string(data[:4]), 16, 0)
+	if err != nil {
+		return 0, nil, fmt.Errorf("pktLineSplitter: decode length: %v", err)
+	}
+
+	// Cast is safe because we requested an int-size number from strconv.ParseInt
+	pktLength := int(pktLength64)
+
+	if pktLength < 0 {
+		return 0, nil, fmt.Errorf("pktLineSplitter: invalid length: %d", pktLength)
+	}
+
+	if pktLength < 4 {
+		// Special case: magic empty packet 0000, 0001, 0002 or 0003.
+		return 4, data[:4], nil
+	}
+
+	if len(data) < pktLength {
+		// data contains incomplete packet
+
+		if atEOF {
+			return 0, nil, fmt.Errorf("pktLineSplitter: less than %d bytes in input %q", pktLength, data)
+		}
+
+		return 0, nil, nil // want more data
+	}
+
+	return pktLength, data[:pktLength], nil
+}
diff --git a/internal/pktline/pktline_test.go b/internal/pktline/pktline_test.go
new file mode 100644
index 0000000..cf3f6fd
--- /dev/null
+++ b/internal/pktline/pktline_test.go
@@ -0,0 +1,87 @@
+package pktline
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+var (
+	largestString = strings.Repeat("z", 0xffff-4)
+)
+
+func TestScanner(t *testing.T) {
+	largestPacket := "ffff" + largestString
+	testCases := []struct {
+		desc string
+		in   string
+		out  []string
+		fail bool
+	}{
+		{
+			desc: "happy path",
+			in:   "0010hello world!000000010010hello world!",
+			out:  []string{"0010hello world!", "0000", "0001", "0010hello world!"},
+		},
+		{
+			desc: "large input",
+			in:   "0010hello world!0000" + largestPacket + "0000",
+			out:  []string{"0010hello world!", "0000", largestPacket, "0000"},
+		},
+		{
+			desc: "missing byte middle",
+			in:   "0010hello world!00000010010hello world!",
+			out:  []string{"0010hello world!", "0000", "0010010hello wor"},
+			fail: true,
+		},
+		{
+			desc: "unfinished prefix",
+			in:   "0010hello world!000",
+			out:  []string{"0010hello world!"},
+			fail: true,
+		},
+		{
+			desc: "short read in data, only prefix",
+			in:   "0010hello world!0005",
+			out:  []string{"0010hello world!"},
+			fail: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			scanner := NewScanner(strings.NewReader(tc.in))
+			var output []string
+			for scanner.Scan() {
+				output = append(output, scanner.Text())
+			}
+
+			if tc.fail {
+				require.Error(t, scanner.Err())
+			} else {
+				require.NoError(t, scanner.Err())
+			}
+
+			require.Equal(t, tc.out, output)
+		})
+	}
+}
+
+func TestIsDone(t *testing.T) {
+	testCases := []struct {
+		in    string
+		done bool
+	}{
+		{in: "0008abcd", done: false},
+		{in: "invalid packet", done: false},
+		{in: "0009done\n", done: true},
+		{in: "0001", done: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.in, func(t *testing.T) {
+			require.Equal(t, tc.done, IsDone([]byte(tc.in)))
+		})
+	}
+}
diff --git a/internal/testhelper/testdata/testroot/responses/allowed_with_pull_payload.json b/internal/testhelper/testdata/testroot/responses/allowed_with_pull_payload.json
new file mode 100644
index 0000000..c97c18d
--- /dev/null
+++ b/internal/testhelper/testdata/testroot/responses/allowed_with_pull_payload.json
@@ -0,0 +1,40 @@
+{
+  "status": true,
+  "gl_repository": "project-26",
+  "gl_project_path": "group/private",
+  "gl_id": "user-1",
+  "gl_username": "root",
+  "git_config_options": [
+    "option"
+  ],
+  "gitaly": {
+    "repository": {
+      "storage_name": "default",
+      "relative_path": "@hashed/5f/9c/5f9c4ab08cac7457e9111a30e4664920607ea2c115a1433d7be98e97e64244ca.git",
+      "git_object_directory": "path/to/git_object_directory",
+      "git_alternate_object_directories": [
+        "path/to/git_alternate_object_directory"
+      ],
+      "gl_repository": "project-26",
+      "gl_project_path": "group/private"
+    },
+    "address": "unix:gitaly.socket",
+    "token": "token"
+  },
+  "payload": {
+    "action": "geo_proxy_to_primary",
+    "data": {
+      "api_endpoints": [
+        "geo/proxy_git_ssh/info_refs_upload_pack",
+        "geo/proxy_git_ssh/upload_pack"
+      ],
+      "gl_username": "custom",
+      "primary_repo": "https://repo/path"
+    }
+  },
+  "git_protocol": "protocol",
+  "gl_console_messages": [
+    "console",
+    "message"
+  ]
+}
diff --git a/internal/testhelper/testdata/testroot/responses/allowed_with_payload.json b/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json
index 4996b2a..318d42d 100644
--- a/internal/testhelper/testdata/testroot/responses/allowed_with_payload.json
+++ b/internal/testhelper/testdata/testroot/responses/allowed_with_push_payload.json
@@ -20,7 +20,7 @@
   "payload" : {
     "action": "geo_proxy_to_primary",
 		"data": {
-			"api_endpoints": ["geo/proxy_git_push_ssh/info_refs", "geo/proxy_git_push_ssh/push"],
+			"api_endpoints": ["geo/proxy_git_ssh/info_refs_receive_pack", "geo/proxy_git_ssh/receive_pack"],
 			"gl_username":   "custom",
 			"primary_repo":  "https://repo/path"
     }
diff --git a/spec/gitlab_shell_custom_git_receive_pack_spec.rb b/spec/gitlab_shell_custom_git_receive_pack_spec.rb
index 771ef1e..1979738 100644
--- a/spec/gitlab_shell_custom_git_receive_pack_spec.rb
+++ b/spec/gitlab_shell_custom_git_receive_pack_spec.rb
@@ -15,14 +15,14 @@ describe 'Custom bin/gitlab-shell git-receive-pack' do
   end
 
   def mock_server(server)
-    server.mount_proc('/geo/proxy_git_push_ssh/info_refs') do |req, res|
+    server.mount_proc('/geo/proxy_git_ssh/info_refs_receive_pack') do |req, res|
       res.content_type = 'application/json'
       res.status = 200
 
       res.body = {"result" => "#{Base64.encode64('custom')}"}.to_json
     end
 
-    server.mount_proc('/geo/proxy_git_push_ssh/push') do |req, res|
+    server.mount_proc('/geo/proxy_git_ssh/receive_pack') do |req, res|
       res.content_type = 'application/json'
       res.status = 200
 
@@ -50,7 +50,7 @@ describe 'Custom bin/gitlab-shell git-receive-pack' do
           "payload" => {
             "action" => "geo_proxy_to_primary",
             "data" => {
-              "api_endpoints" => ["/geo/proxy_git_push_ssh/info_refs", "/geo/proxy_git_push_ssh/push"],
+              "api_endpoints" => ["/geo/proxy_git_ssh/info_refs_receive_pack", "/geo/proxy_git_ssh/receive_pack"],
               "gl_username" =>   "custom",
               "primary_repo" =>  "https://repo/path"
             },
diff --git a/spec/gitlab_shell_custom_git_upload_pack_spec.rb b/spec/gitlab_shell_custom_git_upload_pack_spec.rb
new file mode 100644
index 0000000..6770344
--- /dev/null
+++ b/spec/gitlab_shell_custom_git_upload_pack_spec.rb
@@ -0,0 +1,118 @@
+require_relative 'spec_helper'
+
+require 'open3'
+require 'json'
+require 'base64'
+
+describe 'Custom bin/gitlab-shell git-upload-pack' do
+  include_context 'gitlab shell'
+
+  let(:env) { {'SSH_CONNECTION' => 'fake', 'SSH_ORIGINAL_COMMAND' => 'git-upload-pack group/repo' } }
+  let(:divider) { "remote: ========================================================================\n" }
+
+  before(:context) do
+    write_config("gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}")
+  end
+
+  def mock_server(server)
+    server.mount_proc('/geo/proxy_git_ssh/info_refs_upload_pack') do |req, res|
+      res.content_type = 'application/json'
+      res.status = 200
+
+      res.body = {"result" => "#{Base64.encode64('custom')}"}.to_json
+    end
+
+    server.mount_proc('/geo/proxy_git_ssh/upload_pack') do |req, res|
+      res.content_type = 'application/json'
+      res.status = 200
+
+      output = JSON.parse(req.body)['output']
+
+      res.body = {"result" => output}.to_json
+    end
+
+    server.mount_proc('/api/v4/internal/allowed') do |req, res|
+      res.content_type = 'application/json'
+
+      key_id = req.query['key_id'] || req.query['username']
+
+      unless key_id
+        body = JSON.parse(req.body)
+        key_id = body['key_id'] || body['username'].to_s
+      end
+
+      case key_id
+      when '100', 'someone' then
+        res.status = 300
+        body = {
+          "gl_id" => "user-100",
+          "status" => true,
+          "payload" => {
+            "action" => "geo_proxy_to_primary",
+            "data" => {
+              "api_endpoints" => ["/geo/proxy_git_ssh/info_refs_upload_pack", "/geo/proxy_git_ssh/upload_pack"],
+              "gl_username" =>   "custom",
+              "primary_repo" =>  "https://repo/path"
+            },
+          },
+          "gl_console_messages" => ["console", "message"]
+        }
+        res.body = body.to_json
+      else
+        res.status = 403
+      end
+    end
+  end
+
+  describe 'dialog for performing a custom action' do
+    context 'when API calls perform successfully' do
+      let(:remote_blank_line) { "remote: \n" }
+      def verify_successful_call!(cmd)
+        Open3.popen3(env, cmd) do |stdin, stdout, stderr|
+          expect(stderr.gets).to eq(remote_blank_line)
+          expect(stderr.gets).to eq("remote: console\n")
+          expect(stderr.gets).to eq("remote: message\n")
+          expect(stderr.gets).to eq(remote_blank_line)
+
+          stdin.puts("0032want 343d70886785dc1f98aaf70f3b4ca87c93a5d0dd\n")
+          stdin.close
+
+          expect(stdout.gets(6)).to eq("custom")
+          expect(stdout.flush.read).to eq("0032want 343d70886785dc1f98aaf70f3b4ca87c93a5d0dd\n")
+        end
+      end
+
+      context 'when key is provided' do
+        let(:cmd) { "#{gitlab_shell_path} key-100" }
+
+        it 'custom action is performed' do
+          verify_successful_call!(cmd)
+        end
+      end
+
+      context 'when username is provided' do
+        let(:cmd) { "#{gitlab_shell_path} username-someone" }
+
+        it 'custom action is performed' do
+          verify_successful_call!(cmd)
+        end
+      end
+    end
+
+    context 'when API error occurs' do
+      let(:cmd) { "#{gitlab_shell_path} key-101" }
+
+      it 'custom action is not performed' do
+        Open3.popen2e(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("remote: \n")
+          expect(stdout.gets).to eq(divider)
+          expect(stdout.gets).to eq("remote: \n")
+          expect(stdout.gets).to eq("remote: Internal API error (403)\n")
+          expect(stdout.gets).to eq("remote: \n")
+          expect(stdout.gets).to eq(divider)
+          expect(stdout.gets).to eq("remote: \n")
+        end
+      end
+    end
+  end
+end