Merge branch '330494-geo-ssh-push-proxy-hanging-bug' into 'main'

Fix the Geo SSH push proxy hanging

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

internal/command/shared/customaction/customaction.go

@@ -112,10 +112,32 @@ func (c *Command) performRequest(ctx context.Context, client *client.GitlabNetCl
 }
 
 func (c *Command) readFromStdin() ([]byte, error) {
-	output := new(bytes.Buffer)
-	_, err := io.Copy(output, c.ReadWriter.In)
+	var output []byte
+	var needsPackData bool
+
+	scanner := pktline.NewScanner(c.ReadWriter.In)
+	for scanner.Scan() {
+		line := scanner.Bytes()
+		output = append(output, line...)
+
+		if pktline.IsFlush(line) {
+			break
+		}
 
-	return output.Bytes(), err
+		if !needsPackData && !pktline.IsRefRemoval(line) {
+			needsPackData = true
+		}
+	}
+
+	if needsPackData {
+		packData := new(bytes.Buffer)
+		_, err := io.Copy(packData, c.ReadWriter.In)
+
+		output = append(output, packData.Bytes()...)
+		return output, err
+	} else {
+		return output, nil
+	}
 }
 
 func (c *Command) readFromStdinNoEOF() []byte {

internal/command/shared/customaction/customaction_test.go

@@ -46,7 +46,7 @@ func TestExecuteEOFSent(t *testing.T) {
 				require.NoError(t, json.Unmarshal(b, &request))
 
 				require.Equal(t, request.Data.UserId, who)
-				require.Equal(t, "input", string(request.Output))
+				require.Equal(t, "0009input", string(request.Output))
 
 				err = json.NewEncoder(w).Encode(Response{Result: []byte("output")})
 				require.NoError(t, err)
@@ -58,7 +58,7 @@ func TestExecuteEOFSent(t *testing.T) {
 
 	outBuf := &bytes.Buffer{}
 	errBuf := &bytes.Buffer{}
-	input := bytes.NewBufferString("input")
+	input := bytes.NewBufferString("0009input")
 
 	response := &accessverifier.Response{
 		Who: who,

internal/pktline/pktline.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"regexp"
 	"strconv"
 )
 
@@ -16,6 +17,8 @@ const (
 	pktDelim   = "0001"
 )
 
+var branchRemovalPktRegexp = regexp.MustCompile(`\A[a-f0-9]{4}[a-f0-9]{40} 0{40} `)
+
 // NewScanner returns a bufio.Scanner that splits on Git pktline boundaries
 func NewScanner(r io.Reader) *bufio.Scanner {
 	scanner := bufio.NewScanner(r)
@@ -24,7 +27,16 @@ func NewScanner(r io.Reader) *bufio.Scanner {
 	return scanner
 }
 
-// IsDone detects the special flush packet '0009done\n'
+func IsRefRemoval(pkt []byte) bool {
+	return branchRemovalPktRegexp.Match(pkt)
+}
+
+// IsFlush detects the special flush packet '0000'
+func IsFlush(pkt []byte) bool {
+	return bytes.Equal(pkt, []byte("0000"))
+}
+
+// IsDone detects the special done packet '0009done\n'
 func IsDone(pkt []byte) bool {
 	return bytes.Equal(pkt, PktDone())
 }

internal/pktline/pktline_test.go

@@ -68,6 +68,40 @@ func TestScanner(t *testing.T) {
 	}
 }
 
+func TestIsRefRemoval(t *testing.T) {
+	testCases := []struct {
+		in        string
+		isRemoval bool
+	}{
+		{in: "003f7217a7c7e582c46cec22a130adf4b9d7d950fba0 7d1665144a3a975c05f1f43902ddaf084e784dbe refs/heads/debug", isRemoval: false},
+		{in: "003f0000000000000000000000000000000000000000 7d1665144a3a975c05f1f43902ddaf084e784dbe refs/heads/debug", isRemoval: false},
+		{in: "003f7217a7c7e582c46cec22a130adf4b9d7d950fba0 0000000000000000000000000000000000000000 refs/heads/debug", isRemoval: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.in, func(t *testing.T) {
+			require.Equal(t, tc.isRemoval, IsRefRemoval([]byte(tc.in)))
+		})
+	}
+}
+
+func TestIsFlush(t *testing.T) {
+	testCases := []struct {
+		in    string
+		flush bool
+	}{
+		{in: "0008abcd", flush: false},
+		{in: "invalid packet", flush: false},
+		{in: "0000", flush: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.in, func(t *testing.T) {
+			require.Equal(t, tc.flush, IsFlush([]byte(tc.in)))
+		})
+	}
+}
+
 func TestIsDone(t *testing.T) {
 	testCases := []struct {
 		in   string

spec/gitlab_shell_custom_git_receive_pack_spec.rb

@@ -74,11 +74,11 @@ describe 'Custom bin/gitlab-shell git-receive-pack' do
           expect(stderr.gets).to eq("remote: message\n")
           expect(stderr.gets).to eq(remote_blank_line)
 
-          stdin.puts("input")
+          stdin.puts("0009input")
           stdin.close
 
           expect(stdout.gets(6)).to eq("custom")
-          expect(stdout.flush.read).to eq("input\n")
+          expect(stdout.flush.read).to eq("0009input")
         end
       end