git: stats: Extract ReferenceDiscovery parsing from the analyzehttp code

The reference discovery analysis code is currently contained in
"analyzehttp.go", implying that it's useful for HTTP transports only.
This commit extracts the code into its own file and provides a nice
interface that's decoupled from the HTTP parsing code in order to make
it reusable for SSH transports.

3 files changed, 218 insertions, 103 deletions

diff --git a/internal/git/stats/analyzehttp.go b/internal/git/stats/analyzehttp.go
index c5c776a78..9bd13617e 100644
--- a/internal/git/stats/analyzehttp.go
+++ b/internal/git/stats/analyzehttp.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"os"
 	"strings"
@@ -48,116 +49,17 @@ func ctxErr(ctx context.Context, err error) error {
 	return err
 }
 
-type Reference struct {
-	Oid, Name string
-}
-
 type Get struct {
 	start          time.Time
 	responseHeader time.Duration
 	httpStatus     int
-	firstGitPacket time.Duration
-	responseBody   time.Duration
-	payloadSize    int64
-	packets        int
-	refs           []Reference
-	caps           []string
+	ReferenceDiscovery
 }
 
 func (g *Get) ResponseHeader() time.Duration { return g.responseHeader }
 func (g *Get) HTTPStatus() int               { return g.httpStatus }
-func (g *Get) FirstGitPacket() time.Duration { return g.firstGitPacket }
-func (g *Get) ResponseBody() time.Duration   { return g.responseBody }
-func (g *Get) PayloadSize() int64            { return g.payloadSize }
-func (g *Get) Packets() int                  { return g.packets }
-func (g *Get) Refs() []Reference             { return g.refs }
-func (g *Get) Caps() []string                { return g.caps }
-
-type uploadPackState int
-
-const (
-	uploadPackExpectService uploadPackState = iota
-	uploadPackExpectFlush
-	uploadPackExpectRefWithCaps
-	uploadPackExpectRef
-	uploadPackExpectEnd
-)
-
-// Expected response:
-// - "# service=git-upload-pack\n"
-// - FLUSH
-// - "<OID> <ref>\x00<capabilities>\n"
-// - "<OID> <ref>\n"
-// - ...
-// - FLUSH
-func (g *Get) Parse(body io.Reader) error {
-	state := uploadPackExpectService
-	scanner := pktline.NewScanner(body)
-
-	for ; scanner.Scan(); g.packets++ {
-		pkt := scanner.Bytes()
-		data := pktline.Data(pkt)
-		g.payloadSize += int64(len(data))
-
-		switch state {
-		case uploadPackExpectService:
-			g.firstGitPacket = time.Since(g.start)
-			header := string(data)
-			if header != "# service=git-upload-pack\n" {
-				return fmt.Errorf("unexpected header %q", header)
-			}
-
-			state = uploadPackExpectFlush
-		case uploadPackExpectFlush:
-			if !pktline.IsFlush(pkt) {
-				return errors.New("missing flush after service announcement")
-			}
-
-			state = uploadPackExpectRefWithCaps
-		case uploadPackExpectRefWithCaps:
-			split := bytes.Split(data, []byte{0})
-			if len(split) != 2 {
-				return errors.New("invalid first reference line")
-			}
-			g.caps = strings.Split(string(split[1]), " ")
-
-			ref := strings.SplitN(string(split[0]), " ", 2)
-			if len(ref) != 2 {
-				continue
-			}
-			g.refs = append(g.refs, Reference{Oid: ref[0], Name: ref[1]})
-
-			state = uploadPackExpectRef
-		case uploadPackExpectRef:
-			if pktline.IsFlush(pkt) {
-				state = uploadPackExpectEnd
-				continue
-			}
-
-			split := strings.SplitN(string(data), " ", 2)
-			if len(split) != 2 {
-				continue
-			}
-			g.refs = append(g.refs, Reference{Oid: split[0], Name: split[1]})
-		case uploadPackExpectEnd:
-			return errors.New("received packet after flush")
-		}
-	}
-
-	if err := scanner.Err(); err != nil {
-		return err
-	}
-	if len(g.refs) == 0 {
-		return errors.New("received no references")
-	}
-	if len(g.caps) == 0 {
-		return errors.New("received no capabilities")
-	}
-
-	g.responseBody = time.Since(g.start)
-
-	return nil
-}
+func (g *Get) FirstGitPacket() time.Duration { return g.FirstPacket().Sub(g.start) }
+func (g *Get) ResponseBody() time.Duration   { return g.LastPacket().Sub(g.start) }
 
 func (cl *Clone) doGet(ctx context.Context) error {
 	req, err := http.NewRequest("GET", cl.URL+"/info/refs?service=git-upload-pack", nil)
@@ -188,7 +90,10 @@ func (cl *Clone) doGet(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+	defer func() {
+		io.Copy(ioutil.Discard, resp.Body)
+		resp.Body.Close()
+	}()
 
 	if code := resp.StatusCode; code < 200 || code >= 400 {
 		return fmt.Errorf("git http get: unexpected http status: %d", code)
diff --git a/internal/git/stats/reference_discovery.go b/internal/git/stats/reference_discovery.go
new file mode 100644
index 000000000..032009fa1
--- /dev/null
+++ b/internal/git/stats/reference_discovery.go
@@ -0,0 +1,141 @@
+package stats
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+	"time"
+
+	"gitlab.com/gitlab-org/gitaly/internal/git/pktline"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/text"
+)
+
+// Reference as used by the reference discovery protocol
+type Reference struct {
+	// Oid is the object ID the reference points to
+	Oid string
+	// Name of the reference. The name will be suffixed with ^{} in case
+	// the reference is the peeled commit.
+	Name string
+}
+
+// ReferenceDiscovery contains information about a reference discovery session.
+type ReferenceDiscovery struct {
+	// firstPacket tracks the time when the first pktline was received
+	firstPacket time.Time
+	// lastPacket tracks the time when the last pktline was received
+	lastPacket time.Time
+	// payloadSize tracks the size of all pktlines' data
+	payloadSize int64
+	// packets tracks the total number of packets consumed
+	packets int
+	// refs contains all announced references
+	refs []Reference
+	// caps contains all supported capabilities
+	caps []string
+}
+
+func (d *ReferenceDiscovery) FirstPacket() time.Time { return d.firstPacket }
+func (d *ReferenceDiscovery) LastPacket() time.Time  { return d.lastPacket }
+func (d *ReferenceDiscovery) PayloadSize() int64     { return d.payloadSize }
+func (d *ReferenceDiscovery) Packets() int           { return d.packets }
+func (d *ReferenceDiscovery) Refs() []Reference      { return d.refs }
+func (d *ReferenceDiscovery) Caps() []string         { return d.caps }
+
+type referenceDiscoveryState int
+
+const (
+	referenceDiscoveryExpectService referenceDiscoveryState = iota
+	referenceDiscoveryExpectFlush
+	referenceDiscoveryExpectRefWithCaps
+	referenceDiscoveryExpectRef
+	referenceDiscoveryExpectEnd
+)
+
+// ParseReferenceDiscovery parses a client's reference discovery stream and
+// returns either information about the reference discovery or an error in case
+// it couldn't make sense of the client's request.
+func ParseReferenceDiscovery(body io.Reader) (ReferenceDiscovery, error) {
+	d := ReferenceDiscovery{}
+	return d, d.Parse(body)
+}
+
+// Parse parses a client's reference discovery stream into the given
+// ReferenceDiscovery struct or returns an error in case it couldn't make sense
+// of the client's request.
+//
+// Expected protocol:
+// - "# service=git-upload-pack\n"
+// - FLUSH
+// - "<OID> <ref>\x00<capabilities>\n"
+// - "<OID> <ref>\n"
+// - ...
+// - FLUSH
+func (d *ReferenceDiscovery) Parse(body io.Reader) error {
+	state := referenceDiscoveryExpectService
+	scanner := pktline.NewScanner(body)
+
+	for ; scanner.Scan(); d.packets++ {
+		pkt := scanner.Bytes()
+		data := text.ChompBytes(pktline.Data(pkt))
+		d.payloadSize += int64(len(data))
+
+		switch state {
+		case referenceDiscoveryExpectService:
+			d.firstPacket = time.Now()
+			if data != "# service=git-upload-pack" {
+				return fmt.Errorf("unexpected header %q", data)
+			}
+
+			state = referenceDiscoveryExpectFlush
+		case referenceDiscoveryExpectFlush:
+			if !pktline.IsFlush(pkt) {
+				return errors.New("missing flush after service announcement")
+			}
+
+			state = referenceDiscoveryExpectRefWithCaps
+		case referenceDiscoveryExpectRefWithCaps:
+			split := strings.SplitN(data, "\000", 2)
+			if len(split) != 2 {
+				return errors.New("invalid first reference line")
+			}
+
+			ref := strings.SplitN(string(split[0]), " ", 2)
+			if len(ref) != 2 {
+				return errors.New("invalid reference line")
+			}
+			d.refs = append(d.refs, Reference{Oid: ref[0], Name: ref[1]})
+			d.caps = strings.Split(string(split[1]), " ")
+
+			state = referenceDiscoveryExpectRef
+		case referenceDiscoveryExpectRef:
+			if pktline.IsFlush(pkt) {
+				state = referenceDiscoveryExpectEnd
+				continue
+			}
+
+			split := strings.SplitN(data, " ", 2)
+			if len(split) != 2 {
+				return errors.New("invalid reference line")
+			}
+			d.refs = append(d.refs, Reference{Oid: split[0], Name: split[1]})
+		case referenceDiscoveryExpectEnd:
+			return errors.New("received packet after flush")
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+	if len(d.refs) == 0 {
+		return errors.New("received no references")
+	}
+	if len(d.caps) == 0 {
+		return errors.New("received no capabilities")
+	}
+
+	d.lastPacket = time.Now()
+
+	return nil
+}
diff --git a/internal/git/stats/reference_discovery_test.go b/internal/git/stats/reference_discovery_test.go
new file mode 100644
index 000000000..1bd998094
--- /dev/null
+++ b/internal/git/stats/reference_discovery_test.go
@@ -0,0 +1,69 @@
+package stats
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/internal/git/pktline"
+)
+
+func TestSingleRefParses(t *testing.T) {
+	buf := &bytes.Buffer{}
+	pktline.WriteString(buf, "# service=git-upload-pack\n")
+	pktline.WriteFlush(buf)
+	pktline.WriteString(buf, oid1+" HEAD\x00capability")
+	pktline.WriteFlush(buf)
+
+	d, err := ParseReferenceDiscovery(buf)
+	require.NoError(t, err)
+	require.Equal(t, []string{"capability"}, d.Caps())
+	require.Equal(t, []Reference{{Oid: oid1, Name: "HEAD"}}, d.Refs())
+}
+
+func TestMultipleRefsAndCapsParse(t *testing.T) {
+	buf := &bytes.Buffer{}
+	pktline.WriteString(buf, "# service=git-upload-pack\n")
+	pktline.WriteFlush(buf)
+	pktline.WriteString(buf, oid1+" HEAD\x00first second")
+	pktline.WriteString(buf, oid2+" refs/heads/master")
+	pktline.WriteFlush(buf)
+
+	d, err := ParseReferenceDiscovery(buf)
+	require.NoError(t, err)
+	require.Equal(t, []string{"first", "second"}, d.Caps())
+	require.Equal(t, []Reference{{Oid: oid1, Name: "HEAD"}, {Oid: oid2, Name: "refs/heads/master"}}, d.Refs())
+}
+
+func TestInvalidHeaderFails(t *testing.T) {
+	buf := &bytes.Buffer{}
+	pktline.WriteString(buf, "# service=invalid\n")
+	pktline.WriteFlush(buf)
+	pktline.WriteString(buf, oid1+" HEAD\x00caps")
+	pktline.WriteFlush(buf)
+
+	_, err := ParseReferenceDiscovery(buf)
+	require.Error(t, err)
+}
+
+func TestMissingRefsFail(t *testing.T) {
+	buf := &bytes.Buffer{}
+	pktline.WriteString(buf, "# service=git-upload-pack\n")
+	pktline.WriteFlush(buf)
+	pktline.WriteFlush(buf)
+
+	_, err := ParseReferenceDiscovery(buf)
+	require.Error(t, err)
+}
+
+func TestInvalidRefFail(t *testing.T) {
+	buf := &bytes.Buffer{}
+	pktline.WriteString(buf, "# service=git-upload-pack\n")
+	pktline.WriteFlush(buf)
+	pktline.WriteString(buf, oid1+" HEAD\x00caps")
+	pktline.WriteString(buf, oid2)
+	pktline.WriteFlush(buf)
+
+	_, err := ParseReferenceDiscovery(buf)
+	require.Error(t, err)
+}