git-lfs-transfer: Add support for batch download

Issue: https://gitlab.com/gitlab-org/gitlab/-/issues/336618

internal/command/lfstransfer/gitlab_backend.go

@@ -2,6 +2,10 @@ package lfstransfer
 
 import (
 	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"io"
 	"io/fs"
@@ -9,6 +13,7 @@ import (
 	"github.com/charmbracelet/git-lfs-transfer/transfer"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/lfstransfer"
 )
 
 type errCustom struct {
@@ -41,19 +46,102 @@ type GitlabBackend struct {
 	config *config.Config
 	args   *commandargs.Shell
 	auth   *GitlabAuthentication
+	client *lfstransfer.Client
+}
+
+type idData struct {
+	Operation string            `json:"operation"`
+	Oid       string            `json:"oid"`
+	Href      string            `json:"href"`
+	Headers   map[string]string `json:"headers,omitempty"`
 }
 
 func NewGitlabBackend(ctx context.Context, config *config.Config, args *commandargs.Shell, auth *GitlabAuthentication) (*GitlabBackend, error) {
+	client, err := lfstransfer.NewClient(config, args, auth.href, auth.auth)
+	if err != nil {
+		return nil, err
+	}
 	return &GitlabBackend{
 		ctx,
 		config,
 		args,
 		auth,
+		client,
 	}, nil
 }
 
+func (b *GitlabBackend) issueBatchArgs(op string, oid string, href string, headers map[string]string) (id string, token string, err error) {
+	data := &idData{
+		Operation: op,
+		Oid:       oid,
+		Href:      href,
+		Headers:   headers,
+	}
+	dataBinary, err := json.Marshal(data)
+	if err != nil {
+		return "", "", err
+	}
+
+	h := hmac.New(sha256.New, []byte(b.config.Secret))
+	_, err = h.Write(dataBinary)
+	if err != nil {
+		return "", "", err
+	}
+	id = base64.StdEncoding.EncodeToString(dataBinary)
+	token = base64.StdEncoding.EncodeToString(h.Sum(nil))
+	return id, token, nil
+}
+
 func (b *GitlabBackend) Batch(op string, pointers []transfer.BatchItem, args transfer.Args) ([]transfer.BatchItem, error) {
-	return nil, newErrUnsupported("batch")
+	if op != "download" {
+		return nil, newErrUnsupported("upload batch")
+	}
+
+	reqObjects := make([]*lfstransfer.BatchObject, 0)
+	for _, pointer := range pointers {
+		reqObject := &lfstransfer.BatchObject{
+			Oid:  pointer.Oid,
+			Size: pointer.Size,
+		}
+		reqObjects = append(reqObjects, reqObject)
+	}
+	refName := args["refname"]
+	hashAlgo := args["hash-algo"]
+	res, err := b.client.Batch(op, reqObjects, refName, hashAlgo)
+	if err != nil {
+		return nil, err
+	}
+
+	items := make([]transfer.BatchItem, 0)
+	for _, retObject := range res.Objects {
+		var present bool
+		var action *lfstransfer.BatchAction
+		var args transfer.Args
+		if action, present = retObject.Actions[op]; present {
+			id, token, err := b.issueBatchArgs(op, retObject.Oid, action.Href, action.Header)
+			if err != nil {
+				return nil, err
+			}
+			args = transfer.Args{
+				"id":    id,
+				"token": token,
+			}
+		}
+		if op == "upload" {
+			present = !present
+		}
+		batchItem := transfer.BatchItem{
+			Pointer: transfer.Pointer{
+				Oid:  retObject.Oid,
+				Size: retObject.Size,
+			},
+			Present: present,
+			Args:    args,
+		}
+		items = append(items, batchItem)
+	}
+
+	return items, nil
 }
 
 func (b *GitlabBackend) StartUpload(oid string, r io.Reader, args transfer.Args) (io.Closer, error) {

internal/command/lfstransfer/lfstransfer_test.go

@@ -2,10 +2,15 @@ package lfstransfer
 
 import (
 	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"sync"
 	"testing"
 
@@ -18,6 +23,21 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/lfsauthenticate"
 )
 
+const (
+	largeFileContents      = "This is a large file\n"
+	evenLargerFileContents = "This is an even larger file\n"
+)
+
+var (
+	largeFileLen  = len(largeFileContents)
+	largeFileHash = sha256.Sum256([]byte(largeFileContents))
+	largeFileOid  = hex.EncodeToString(largeFileHash[:])
+
+	evenLargerFileLen  = len(evenLargerFileContents)
+	evenLargerFileHash = sha256.Sum256([]byte(evenLargerFileContents))
+	evenLargerFileOid  = hex.EncodeToString(evenLargerFileHash[:])
+)
+
 func setupWaitGroup(t *testing.T, cmd *Command) *sync.WaitGroup {
 	wg := &sync.WaitGroup{}
 	wg.Add(1)
@@ -249,19 +269,64 @@ func TestLfsTransferNoPermissions(t *testing.T) {
 }
 
 func TestLfsTransferBatchDownload(t *testing.T) {
-	_, cmd, pl, _ := setup(t, "rw", "group/repo", "download")
+	url, cmd, pl, _ := setup(t, "rw", "group/repo", "download")
 	wg := setupWaitGroup(t, cmd)
 	negotiateVersion(t, pl)
 
 	writeCommandArgsAndTextData(t, pl, "batch", nil, []string{
 		"00000000 0",
+		fmt.Sprintf("%s %d", largeFileOid, largeFileLen),
+		fmt.Sprintf("%s %d", evenLargerFileOid, evenLargerFileLen),
 	})
 	status, args, data := readStatusArgsAndTextData(t, pl)
-	require.Equal(t, "status 405", status)
+	require.Equal(t, "status 200", status)
 	require.Empty(t, args)
-	require.Equal(t, []string{
-		"error: batch is not yet supported by git-lfs-transfer. See https://gitlab.com/groups/gitlab-org/-/epics/11872 to track progress.",
-	}, data)
+	require.Equal(t, "00000000 0 noop", data[0])
+
+	largeFileArgs := strings.Split(data[1], " ")
+	require.Equal(t, 5, len(largeFileArgs))
+	require.Equal(t, largeFileOid, largeFileArgs[0])
+	require.Equal(t, fmt.Sprint(largeFileLen), largeFileArgs[1])
+	require.Equal(t, "download", largeFileArgs[2])
+
+	var idArg string
+	var tokenArg string
+	for _, arg := range largeFileArgs[3:] {
+		switch {
+		case strings.HasPrefix(arg, "id="):
+			idArg = arg
+		case strings.HasPrefix(arg, "token="):
+			tokenArg = arg
+		default:
+			require.Fail(t, "Unexpected batch item argument: %v", arg)
+		}
+	}
+
+	idBase64, found := strings.CutPrefix(idArg, "id=")
+	require.True(t, found)
+	idBinary, err := base64.StdEncoding.DecodeString(idBase64)
+	require.NoError(t, err)
+	var id map[string]interface{}
+	require.NoError(t, json.Unmarshal(idBinary, &id))
+	require.Equal(t, map[string]interface{}{
+		"operation": "download",
+		"oid":       largeFileOid,
+		"href":      fmt.Sprintf("%s/group/repo/gitlab-lfs/objects/%s", url, largeFileOid),
+		"headers": map[string]interface{}{
+			"Authorization": "Basic 1234567890",
+			"Content-Type":  "application/octet-stream",
+		},
+	}, id)
+
+	h := hmac.New(sha256.New, []byte("very secret"))
+	h.Write(idBinary)
+	tokenBase64, found := strings.CutPrefix(tokenArg, "token=")
+	require.True(t, found)
+	tokenBinary, err := base64.StdEncoding.DecodeString(tokenBase64)
+	require.NoError(t, err)
+	require.Equal(t, h.Sum(nil), tokenBinary)
+
+	require.Equal(t, fmt.Sprintf("%s %d noop", evenLargerFileOid, evenLargerFileLen), data[2])
 
 	quit(t, pl)
 	wg.Wait()
@@ -279,7 +344,7 @@ func TestLfsTransferBatchUpload(t *testing.T) {
 	require.Equal(t, "status 405", status)
 	require.Empty(t, args)
 	require.Equal(t, []string{
-		"error: batch is not yet supported by git-lfs-transfer. See https://gitlab.com/groups/gitlab-org/-/epics/11872 to track progress.",
+		"error: upload batch is not yet supported by git-lfs-transfer. See https://gitlab.com/groups/gitlab-org/-/epics/11872 to track progress.",
 	}, data)
 
 	quit(t, pl)
@@ -461,6 +526,66 @@ func setup(t *testing.T, keyId string, repo string, op string) (string, *Command
 				}
 			},
 		},
+		{
+			Path: "/group/repo/info/lfs/objects/batch",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte("john:sometoken"))), r.Header.Get("Authorization"))
+
+				var requestBody map[string]interface{}
+				json.NewDecoder(r.Body).Decode(&requestBody)
+
+				reqObjects := requestBody["objects"].([]interface{})
+				retObjects := make([]map[string]interface{}, 0)
+				for _, o := range reqObjects {
+					reqObject := o.(map[string]interface{})
+					retObject := map[string]interface{}{
+						"oid": reqObject["oid"],
+					}
+					switch reqObject["oid"] {
+					case largeFileOid:
+						retObject["size"] = largeFileLen
+						if op == "download" {
+							retObject["actions"] = map[string]interface{}{
+								"download": map[string]interface{}{
+									"href": fmt.Sprintf("%s/group/repo/gitlab-lfs/objects/%s", url, largeFileOid),
+									"header": map[string]interface{}{
+										"Authorization": "Basic 1234567890",
+										"Content-Type":  "application/octet-stream",
+									},
+								},
+							}
+						}
+					case evenLargerFileOid:
+						require.Equal(t, evenLargerFileLen, int(reqObject["size"].(float64)))
+						retObject["size"] = evenLargerFileLen
+						if op == "upload" {
+							retObject["actions"] = map[string]interface{}{
+								"upload": map[string]interface{}{
+									"href": fmt.Sprintf("%s/group/repo/gitlab-lfs/objects/%s/%d", url, evenLargerFileOid, evenLargerFileLen),
+									"header": map[string]interface{}{
+										"Authorization": "Basic 1234567890",
+										"Content-Type":  "application/octet-stream",
+									},
+								},
+							}
+						}
+					default:
+						retObject["size"] = reqObject["size"]
+						retObject["error"] = map[string]interface{}{
+							"code":    404,
+							"message": "Not found",
+						}
+					}
+					retObjects = append(retObjects, retObject)
+				}
+
+				retBody := map[string]interface{}{
+					"objects": retObjects,
+				}
+				body, _ := json.Marshal(retBody)
+				w.Write(body)
+			},
+		},
 	}
 	url = testserver.StartHttpServer(t, requests)
 

internal/gitlabnet/lfstransfer/client.go

+package lfstransfer
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"time"
+
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet"
+)
+
+type Client struct {
+	config *config.Config
+	args   *commandargs.Shell
+	href   string
+	auth   string
+}
+
+type BatchAction struct {
+	Href      string            `json:"href"`
+	Header    map[string]string `json:"header,omitempty"`
+	ExpiresAt time.Time         `json:"expires_at,omitempty"`
+	ExpiresIn int               `json:"expires_in,omitempty"`
+}
+
+type BatchObject struct {
+	Oid           string                  `json:"oid,omitempty"`
+	Size          int64                   `json:"size"`
+	Authenticated bool                    `json:"authenticated,omitempty"`
+	Actions       map[string]*BatchAction `json:"actions,omitempty"`
+}
+
+type batchRef struct {
+	Name string `json:"name,omitempty"`
+}
+
+type batchRequest struct {
+	Operation     string         `json:"operation"`
+	Objects       []*BatchObject `json:"objects"`
+	Ref           *batchRef      `json:"ref,omitempty"`
+	HashAlgorithm string         `json:"hash_algo,omitempty"`
+}
+
+type BatchResponse struct {
+	Objects       []*BatchObject `json:"objects"`
+	HashAlgorithm string         `json:"hash_algo,omitempty"`
+}
+
+func NewClient(config *config.Config, args *commandargs.Shell, href string, auth string) (*Client, error) {
+	return &Client{config: config, args: args, href: href, auth: auth}, nil
+}
+
+func (c *Client) Batch(operation string, reqObjects []*BatchObject, ref string, reqHashAlgo string) (*BatchResponse, error) {
+	var bref *batchRef
+	if ref != "" {
+		bref = &batchRef{Name: ref}
+	}
+	body := batchRequest{
+		Operation:     operation,
+		Objects:       reqObjects,
+		Ref:           bref,
+		HashAlgorithm: reqHashAlgo,
+	}
+	jsonData, err := json.Marshal(body)
+	if err != nil {
+		return nil, err
+	}
+
+	jsonReader := bytes.NewReader(jsonData)
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/objects/batch", c.href), jsonReader)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Content-Type", "application/vnd.git-lfs+json")
+	req.Header.Set("Authorization", c.auth)
+
+	client := http.Client{}
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	defer res.Body.Close()
+
+	response := &BatchResponse{}
+	if err := gitlabnet.ParseJSON(res, response); err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}