Geo Pull custom action support

4 files changed, 136 insertions, 9 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)
 }