Fix the Geo SSH push proxy hanging

Geo SSH proxy push currently impossible when the only
action that happens is branch removal. This fix
works in a way that it waits for flush packet from git
and then checks pkt lines to determine is pack data is expected.
The thing is that git doesnt send pack data when only
branch removal happens. Explanation is in
https://gitlab.com/gitlab-org/gitlab/-/issues/330494

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