feat: repsect oid batch arguments

Preserve oid batch arguments in requset/response per [specs](https://github.com/git-lfs/git-lfs/blob/main/docs/proposals/ssh_adapter.md#requests-to-transfer-objects)
charmbracelet · Oct 26, 2023 · c778ab1 · c778ab1
1 parent ef1f8d6
commit c778ab1
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 32 deletions.
diff --git a/internal/local/backend.go b/internal/local/backend.go
@@ -43,26 +43,22 @@ func New(lfsPath string, umask os.FileMode, timestamp *time.Time) *LocalBackend
 }
 
 // Batch implements main.Backend.
-func (l *LocalBackend) Batch(_ string, pointers []transfer.Pointer, _ map[string]string) ([]transfer.BatchItem, error) {
-	items := make([]transfer.BatchItem, 0)
+func (l *LocalBackend) Batch(_ string, pointers []transfer.BatchItem, _ transfer.Args) ([]transfer.BatchItem, error) {
 	for _, o := range pointers {
 		present := false
 		stat, err := os.Stat(oidExpectedPath(l.lfsPath, o.Oid))
 		if err == nil {
 			o.Size = stat.Size()
 			present = true
 		}
-		items = append(items, transfer.BatchItem{
-			Pointer: o,
-			Present: present,
-		})
+		o.Present = present
 	}
-	return items, nil
+	return pointers, nil
 }
 
 // Download implements main.Backend. The returned reader must be closed by the
 // caller.
-func (l *LocalBackend) Download(oid string, _ map[string]string) (fs.File, error) {
+func (l *LocalBackend) Download(oid string, _ transfer.Args) (fs.File, error) {
 	f, err := os.Open(oidExpectedPath(l.lfsPath, oid))
 	if err != nil {
 		return nil, err
@@ -71,7 +67,7 @@ func (l *LocalBackend) Download(oid string, _ map[string]string) (fs.File, error
 }
 
 // FinishUpload implements main.Backend.
-func (l *LocalBackend) FinishUpload(state interface{}, _ map[string]string) error {
+func (l *LocalBackend) FinishUpload(state interface{}, _ transfer.Args) error {
 	switch state := state.(type) {
 	case *UploadState:
 		destPath := oidExpectedPath(l.lfsPath, state.Oid)
@@ -94,7 +90,7 @@ func (l *LocalBackend) FinishUpload(state interface{}, _ map[string]string) erro
 }
 
 // LockBackend implements main.Backend.
-func (l *LocalBackend) LockBackend(_ map[string]string) transfer.LockBackend {
+func (l *LocalBackend) LockBackend(_ transfer.Args) transfer.LockBackend {
 	path := filepath.Join(l.lfsPath, "locks")
 	return NewLockBackend(l, path)
 }
@@ -106,7 +102,7 @@ type UploadState struct {
 }
 
 // StartUpload implements main.Backend. The returned temp file should be closed.
-func (l *LocalBackend) StartUpload(oid string, r io.Reader, _ map[string]string) (interface{}, error) {
+func (l *LocalBackend) StartUpload(oid string, r io.Reader, _ transfer.Args) (interface{}, error) {
 	if r == nil {
 		return nil, fmt.Errorf("%w: received null data", transfer.ErrMissingData)
 	}
@@ -132,7 +128,7 @@ func (l *LocalBackend) StartUpload(oid string, r io.Reader, _ map[string]string)
 }
 
 // Verify implements main.Backend.
-func (l *LocalBackend) Verify(oid string, args map[string]string) (transfer.Status, error) {
+func (l *LocalBackend) Verify(oid string, args transfer.Args) (transfer.Status, error) {
 	var expectedSize int
 	size, ok := args[transfer.SizeKey]
 	if ok {

diff --git a/transfer/args.go b/transfer/args.go
@@ -19,12 +19,12 @@ const (
 )
 
 // ParseArgs parses the given args.
-func ParseArgs(lines []string) (map[string]string, error) {
-	args := make(map[string]string, 0)
-	for _, line := range lines {
+func ParseArgs(parts []string) (Args, error) {
+	args := make(Args, 0)
+	for _, line := range parts {
 		parts := strings.SplitN(line, "=", 2)
 		if len(parts) != 2 {
-			return nil, fmt.Errorf("invalid argument: %q", line)
+			continue
 		}
 		key, value := parts[0], parts[1]
 		args[key] = value
@@ -34,10 +34,18 @@ func ParseArgs(lines []string) (map[string]string, error) {
 }
 
 // ArgsToList converts the given args to a list.
-func ArgsToList(args map[string]string) []string {
+func ArgsToList(args Args) []string {
 	list := make([]string, 0)
 	for key, value := range args {
 		list = append(list, fmt.Sprintf("%s=%s", key, value))
 	}
 	return list
 }
+
+// Args is a key-value pair of arguments.
+type Args map[string]string
+
+// String returns the string representation of the arguments.
+func (a Args) String() string {
+	return strings.Join(ArgsToList(a), " ")
+}
diff --git a/transfer/backend.go b/transfer/backend.go
@@ -14,12 +14,12 @@ const (
 
 // Backend is a Git LFS backend.
 type Backend interface {
-	Batch(op string, pointers []Pointer, args map[string]string) ([]BatchItem, error)
-	StartUpload(oid string, r io.Reader, args map[string]string) (interface{}, error)
-	FinishUpload(state interface{}, args map[string]string) error
-	Verify(oid string, args map[string]string) (Status, error)
-	Download(oid string, args map[string]string) (fs.File, error)
-	LockBackend(args map[string]string) LockBackend
+	Batch(op string, pointers []BatchItem, args Args) ([]BatchItem, error)
+	StartUpload(oid string, r io.Reader, args Args) (interface{}, error)
+	FinishUpload(state interface{}, args Args) error
+	Verify(oid string, args Args) (Status, error)
+	Download(oid string, args Args) (fs.File, error)
+	LockBackend(args Args) LockBackend
 }
 
 // Lock is a Git LFS lock.

diff --git a/transfer/batch.go b/transfer/batch.go
@@ -3,5 +3,6 @@ package transfer
 // BatchItem is a Git LFS batch item.
 type BatchItem struct {
 	Pointer
+	Args    Args
 	Present bool
 }
diff --git a/transfer/oid.go b/transfer/oid.go
@@ -1,6 +1,7 @@
 package transfer
 
 import (
+	"fmt"
 	"path"
 	"regexp"
 )
@@ -11,6 +12,11 @@ type Pointer struct {
 	Size int64  `json:"size"`
 }
 
+// String returns the string representation of the pointer.
+func (p Pointer) String() string {
+	return fmt.Sprintf("%s %d", p.Oid, p.Size)
+}
+
 var oidPattern = regexp.MustCompile(`^[a-f\d]{64}$`)
 
 // IsValid checks if the pointer has a valid structure.

diff --git a/transfer/processor.go b/transfer/processor.go
@@ -61,22 +61,32 @@ func (p *Processor) ReadBatch(op string) ([]BatchItem, error) {
 	}
 	Logf("data: %d %v", len(data), data)
 	Logf("batch: %s args: %d %v data: %d %v", op, len(args), args, len(data), data)
-	items := make([]Pointer, 0)
+	items := make([]BatchItem, 0)
 	for _, line := range data {
 		if line == "" {
 			return nil, ErrInvalidPacket
 		}
-		parts := strings.SplitN(line, " ", 2)
-		if len(parts) != 2 || parts[1] == "" {
+		parts := strings.Split(line, " ")
+		if len(parts) < 2 || parts[1] == "" {
 			return nil, ErrParseError
 		}
 		size, err := strconv.ParseInt(parts[1], 10, 64)
 		if err != nil {
 			return nil, fmt.Errorf("%w: invalid integer, got: %q", ErrParseError, parts[1])
 		}
-		item := Pointer{
-			parts[0],
-			size,
+		var oidArgs Args
+		if len(parts) > 2 {
+			oidArgs, err = ParseArgs(parts[2:])
+			if err != nil {
+				return nil, fmt.Errorf("%w: %s", ErrParseError, err)
+			}
+		}
+		item := BatchItem{
+			Pointer: Pointer{
+				Oid:  parts[0],
+				Size: size,
+			},
+			Args: oidArgs,
 		}
 		items = append(items, item)
 	}
@@ -101,7 +111,11 @@ func (p *Processor) BatchData(op string, presentAction string, missingAction str
 		if item.Present {
 			action = presentAction
 		}
-		oids = append(oids, fmt.Sprintf("%s %d %s", item.Oid, item.Size, action))
+		line := fmt.Sprintf("%s %s", item.Pointer, action)
+		if len(item.Args) > 0 {
+			line = fmt.Sprintf("%s %s", line, item.Args)
+		}
+		oids = append(oids, line)
 	}
 	return NewSuccessStatus(oids), nil
 }
@@ -117,7 +131,7 @@ func (p *Processor) DownloadBatch() (Status, error) {
 }
 
 // SizeFromArgs returns the size from the given args.
-func (Processor) SizeFromArgs(args map[string]string) (int64, error) {
+func SizeFromArgs(args Args) (int64, error) {
 	size, ok := args[SizeKey]
 	if !ok {
 		return 0, fmt.Errorf("missing required size header")
@@ -139,7 +153,7 @@ func (p *Processor) PutObject(oid string) (Status, error) {
 	if err != nil {
 		return nil, fmt.Errorf("%w: %s", ErrParseError, err)
 	}
-	expectedSize, err := p.SizeFromArgs(args)
+	expectedSize, err := SizeFromArgs(args)
 	if err != nil {
 		return nil, err
 	}