Geo Pull custom action support

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)
 	}
 

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

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))

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)
 }
 

internal/gitlabnet/accessverifier/client_test.go

@@ -20,6 +20,7 @@ import (
 var (
 	repo              = "group/private"
 	receivePackAction = commandargs.ReceivePack
+	uploadPackAction  = commandargs.UploadPack
 )
 
 func buildExpectedResponse(who string) *Response {
@@ -94,7 +95,30 @@ func TestGeoPushGetCustomAction(t *testing.T) {
 	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",
 		},

internal/testhelper/testdata/testroot/responses/allowed_with_pull_payload.json

+{
+  "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"
+  ]
+}

spec/gitlab_shell_custom_git_upload_pack_spec.rb

+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