Use Tar hash as checksum

pkg/fileutils/tarxfer.go

@@ -19,8 +19,6 @@ package fileutils
 import (
 	"archive/tar"
 	"context"
-	"crypto/sha256"
-	"encoding/hex"
 	"fmt"
 	"io"
 	"io/fs"
@@ -44,15 +42,15 @@ func NewTarReceiver(cacheBase string, demux *stream.Demultiplexer) *Receiver {
 
 func (r *Receiver) Receive(ctx context.Context, fn fs.WalkDirFunc) (string, error) {
 	errCh := make(chan error, 1)
+	hashCh := make(chan string, 1)
 	dataCh := make(chan []byte)
-	go startTar(r.demux, errCh, dataCh)
+	go startTar(r.demux, errCh, hashCh, dataCh)
 
-	header, err := readTarHeader(ctx, errCh, dataCh)
+	checksum, err := readTarHash(ctx, errCh, hashCh)
 	if err != nil {
 		return "", err
 	}
 
-	checksum := sha256Hex(header)
 	cacheDir := filepath.Join(r.cacheBase, checksum)
 	tarFile := cacheDir + ".tar"
 
@@ -61,6 +59,11 @@ func (r *Receiver) Receive(ctx context.Context, fn fs.WalkDirFunc) (string, erro
 		return "", err
 	}
 
+	header, err := readTarHeader(ctx, errCh, dataCh)
+	if err != nil {
+		return "", err
+	}
+
 	if !cached {
 		if err := writeTar(tarFile, header); err != nil {
 			return "", err
@@ -109,13 +112,9 @@ func (r *Receiver) Receive(ctx context.Context, fn fs.WalkDirFunc) (string, erro
 	})
 }
 
-func sha256Hex(b []byte) string {
-	h := sha256.Sum256(b)
-	return hex.EncodeToString(h[:])
-}
-
-func startTar(demux *stream.Demultiplexer, errCh chan<- error, dataCh chan<- []byte) {
+func startTar(demux *stream.Demultiplexer, errCh chan<- error, hashCh chan<- string, dataCh chan<- []byte) {
 	defer close(errCh)
+	defer close(hashCh)
 	defer close(dataCh)
 
 	for {
@@ -133,6 +132,12 @@ func startTar(demux *stream.Demultiplexer, errCh chan<- error, dataCh chan<- []b
 				errCh <- fmt.Errorf("server error in TAR mode: %s", errMsg)
 				return
 			}
+
+			if hash, ok := bt.Metadata["hash"]; ok {
+				hashCh <- hash
+				continue
+			}
+
 			dataCh <- bt.Data
 			if bt.Complete {
 				errCh <- nil
@@ -156,14 +161,31 @@ func startTar(demux *stream.Demultiplexer, errCh chan<- error, dataCh chan<- []b
 	}
 }
 
+func readTarHash(ctx context.Context, errCh <-chan error, hashCh <-chan string) (string, error) {
+	select {
+	case h, ok := <-hashCh:
+		if !ok {
+			if e := <-errCh; e != nil {
+				return "", e
+			}
+			return "", fmt.Errorf("hash channel closed, no hash received")
+		}
+		return h, nil
+	case e := <-errCh:
+		return "", e
+	case <-ctx.Done():
+		return "", ctx.Err()
+	}
+}
+
 func readTarHeader(ctx context.Context, errCh <-chan error, dataCh <-chan []byte) ([]byte, error) {
 	select {
 	case d, ok := <-dataCh:
 		if !ok {
 			if e := <-errCh; e != nil {
 				return nil, e
 			}
-			return nil, nil
+			return nil, fmt.Errorf("data channel closed")
 		}
 		return d, nil
 	case e := <-errCh:

pkg/fileutils/tarxfer_test.go

@@ -21,7 +21,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
-	"fmt"
+	"encoding/hex"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -72,13 +72,17 @@ func TestReceiver_Receive_Success(t *testing.T) {
 	if len(archive) < 512 {
 		t.Fatalf("tar archive too small: %d", len(archive))
 	}
+
+	hashBytes := sha256.Sum256(archive)
+	hash := hex.EncodeToString(hashBytes[:])
 	header := archive[:512]
 	body := archive[512:]
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	demux := newDemux(ctx)
 
+	_ = demux.Accept(btPacket([]byte{}, false, map[string]string{"hash": hash}))
 	_ = demux.Accept(btPacket(header, false, nil))
 	_ = demux.Accept(btPacket(body, true, nil))
 
@@ -96,9 +100,8 @@ func TestReceiver_Receive_Success(t *testing.T) {
 		t.Fatalf("Receive returned error: %v", err)
 	}
 
-	wantChecksum := func(b []byte) string { h := sha256.Sum256(b); return fmt.Sprintf("%x", h) }(header)
-	if checksum != wantChecksum {
-		t.Fatalf("checksum mismatch: want %s, got %s", wantChecksum, checksum)
+	if checksum != hash {
+		t.Fatalf("checksum mismatch: want %s, got %s", hash, checksum)
 	}
 
 	if len(visited) != 1 || visited[0] != "file1" {

pkg/fssync/walk_test.go

@@ -23,8 +23,6 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	gofs "io/fs"
-	"os"
-	"path/filepath"
 	"testing"
 	"time"
 
@@ -42,7 +40,7 @@ var (
 	headers = map[string][]byte{}
 )
 
-func makeNestedTarHeaderAndBody() (hdr, rest []byte) {
+func makeNestedTarHeaderAndBody() (checksum string, full []byte) {
 	const payload = "hello from tar with nesting\n"
 
 	var buf bytes.Buffer
@@ -68,22 +66,34 @@ func makeNestedTarHeaderAndBody() (hdr, rest []byte) {
 
 	_ = tw.Close()
 
-	full := buf.Bytes()
-	return full[:512], full[512:]
+	full = buf.Bytes()
+	header := sha256.Sum256(full)
+	return hex.EncodeToString(header[:]), full
 }
 
 func (p *FSSyncProxy) Send(s *api.ServerStream) error {
 	id := s.BuildId
 	d := demuxes[id]
-	header, body := makeNestedTarHeaderAndBody()
+	checksum, full := makeNestedTarHeaderAndBody()
 	go func() {
 		_ = d.Accept(&api.ClientStream{
 			BuildId: id,
 			PacketType: &api.ClientStream_BuildTransfer{
 				BuildTransfer: &api.BuildTransfer{
 					Id:        id,
 					Direction: api.TransferDirection_INTO,
-					Data:      append(header, body...),
+					Metadata:  map[string]string{"hash": checksum},
+				},
+			},
+		})
+
+		_ = d.Accept(&api.ClientStream{
+			BuildId: id,
+			PacketType: &api.ClientStream_BuildTransfer{
+				BuildTransfer: &api.BuildTransfer{
+					Id:        id,
+					Direction: api.TransferDirection_INTO,
+					Data:      full,
 				},
 			},
 		})
@@ -133,13 +143,7 @@ func TestWalk_UnsupportedMode(t *testing.T) {
 func TestWalk_TarModeSuccess(t *testing.T) {
 	tmp := t.TempDir()
 
-	header, body := makeNestedTarHeaderAndBody()
-	sum := sha256.Sum256(header)
-	checksum := hex.EncodeToString(sum[:])
-	cacheDir := filepath.Join(tmp, checksum)
-	if err := os.MkdirAll(cacheDir, 0o755); err != nil {
-		t.Fatalf("mkdir cacheDir: %v", err)
-	}
+	_, full := makeNestedTarHeaderAndBody()
 
 	fs := NewFS(context.Background(), &FSSyncProxy{}, "/", tmp)
 
@@ -151,10 +155,10 @@ func TestWalk_TarModeSuccess(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Walk returned err=%v", err)
 	}
-	fsSum := sha256.Sum256(append(header, body...))
+	fsSum := sha256.Sum256(full)
 	fsChecksum := hex.EncodeToString(fsSum[:])
 	if fs.getChecksum() != fsChecksum {
-		t.Errorf("checksum = %q, want %q", fs.getChecksum(), checksum)
+		t.Errorf("checksum = %q, want %q", fs.getChecksum(), fsChecksum)
 	}
 	if len(walked) == 0 {
 		t.Errorf("walk callback not invoked")