Fix mirror padding pickup

mirror_lifecycle.go

@@ -259,7 +259,7 @@ func (mt *MirrorTarget) AddEntries(ctx context.Context, uploadStart, uploadEnd u
 	}
 
 	bundleIdx := uploadStart / layout.EntryBundleWidth
-	nextEntry, newRoot, err := mt.writer.IntegrateBundles(ctx, bundleIdx, mt.bundleIterator(ctx, next, uploadStart, pendingCP))
+	nextEntry, newRoot, err := mt.writer.IntegrateBundles(ctx, bundleIdx, mt.bundleIterator(ctx, next, uploadStart, nextEntry, pendingCP))
 	switch {
 	case err != nil:
 		return 0, 0, nil, nil, err
@@ -291,30 +291,34 @@ func (mt *MirrorTarget) AddEntries(ctx context.Context, uploadStart, uploadEnd u
 //
 // Yielded entry bundles are always aligned to bundle boundaries. Specifically, this means that if the provided start is _not_ bundle aligned, then we will
 // fetch entries from the bundle at start/256 and use those entries to left-pad the first yielded bundle.
-func (mt *MirrorTarget) bundleIterator(ctx context.Context, next func() (*MirrorPackage, error), start uint64, pendingCP *log.Checkpoint) func(func(*api.EntryBundle, error) bool) {
+func (mt *MirrorTarget) bundleIterator(ctx context.Context, next func() (*MirrorPackage, error), start, treeSize uint64, pendingCP *log.Checkpoint) func(func(*api.EntryBundle, error) bool) {
 	crf := compact.RangeFactory{Hash: rfc6962.DefaultHasher.HashChildren}
 	return func(yield func(*api.EntryBundle, error) bool) {
 		// Check for unaligned upload start, and fetch entries from the start of the bundle to use to pad.
 		// This is necessary for the subtree proof for such an unaligned first bundle to validate.
 		var padEntries [][]byte
-		if p := start % layout.EntryBundleWidth; p != 0 {
-			// non-aligned starting bundle
-			br, err := mt.reader.ReadEntryBundle(ctx, start/layout.EntryBundleWidth, uint8(p))
+		if startPad := start % layout.EntryBundleWidth; startPad != 0 {
+			bIdx := start / layout.EntryBundleWidth
+			// The tlog-mirror spec allows an upload to start some way below the current tree size, so the
+			// stored bundle we need to read could well contain more entries than we're going to use for padding.
+			// We request a partial tile of the correct size to fetch just the slice of entries we need for padding.
+			storedPartial := layout.PartialTileSize(0, bIdx, treeSize)
+			br, err := mt.reader.ReadEntryBundle(ctx, bIdx, storedPartial)
 			if err != nil {
 				yield(nil, fmt.Errorf("failed to read bundle containing uploadStart (%d): %v", start, err))
 				return
 			}
-			// Parse and clip, just in case we were returned data from a full tile.
+			// Parse and clip, since the bundle we read may contain more entries than we need for padding.
 			b := &api.EntryBundle{}
 			if err := b.UnmarshalText(br); err != nil {
 				yield(nil, fmt.Errorf("failed to unmarshal bundle containing uploadStart (%d): %v", start, err))
 				return
 			}
-			if l := len(b.Entries); l < int(p) {
-				yield(nil, fmt.Errorf("POTENTIAL CORRUPTION: partial bundle at index %d.%d has only %d entries", start/layout.EntryBundleWidth, p, l))
+			if l := len(b.Entries); l < int(startPad) {
+				yield(nil, fmt.Errorf("POTENTIAL CORRUPTION: bundle at index %d.%d has only %d entries, want at least %d", bIdx, storedPartial, l, startPad))
 				return
 			}
-			padEntries = b.Entries[:p]
+			padEntries = b.Entries[:startPad]
 			// SPEC: The subtree consistency proof is computed from the subtree defined by [rounded_start + i * 256, end), and the log
 			//       checkpoint with tree size upload_end
 			start &= ^uint64(0xff) // floor to bundle boundary

mirror_lifecycle_test.go

@@ -18,13 +18,16 @@ import (
 	"bytes"
 	"context"
 	"encoding/base64"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
 	"iter"
+	"os"
 	"strings"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	fnote "github.com/transparency-dev/formats/note"
 	"github.com/transparency-dev/merkle"
 	"github.com/transparency-dev/merkle/compact"
@@ -543,79 +546,182 @@ func TestMirrorTarget_AddEntries_VerifySubtreeProof(t *testing.T) {
 	}
 }
 
-func TestMirrorTarget_AddEntries_Unaligned_PadsFirstBundle(t *testing.T) {
-	const (
-		testIntegratedSize = uint64(270)
-		testUploadStart    = uint64(270) // not aligned: 270 % 256 = 14
-		testUploadEnd      = testPendingCPSize
-	)
+// TestMirrorTarget_AddEntries_Unaligned tests handling of uploads whose start index is not aligned
+// to an entry bundle boundary:
+//   - the entry bundle resource we read for padding must have the shape (full, or partial of a given
+//     size) which the mirror's current tree size implies is actually stored, and
+//   - the entries we read must be used to left-pad the first bundle yielded for integration, with the
+//     subtree proof verified against the floored (bundle aligned) start index.
+func TestMirrorTarget_AddEntries_Unaligned(t *testing.T) {
+	// storedEntries are the entries the mirror already has integrated, i.e. what a read of the
+	// bundle containing uploadStart will return.
+	storedEntries := make([][]byte, layout.EntryBundleWidth)
+	for i := range storedEntries {
+		storedEntries[i] = fmt.Appendf(nil, "stored-%d", i)
+	}
+	uploadedEntries := [][]byte{[]byte("uploaded-0"), []byte("uploaded-1")}
 
-	var readEntryBundleCalled bool
+	for _, test := range []struct {
+		desc           string
+		integratedSize uint64
+		uploadStart    uint64
+		// wantRead is true if a read of the bundle containing uploadStart is expected, in which
+		// case wantReadIdx/wantReadPartial describe the resource which must be requested.
+		wantRead        bool
+		wantReadIdx     uint64
+		wantReadPartial uint8
+		// wantNumPad is the number of stored entries expected to be prepended to the first
+		// bundle yielded for integration.
+		wantNumPad int
+	}{
+		{
+			desc:           "aligned start, no padding needed",
+			integratedSize: 256,
+			uploadStart:    256,
+			wantRead:       false,
+			wantNumPad:     0,
+		}, {
+			desc:            "unaligned start == tree size",
+			integratedSize:  270,
+			uploadStart:     270,
+			wantRead:        true,
+			wantReadIdx:     1,
+			wantReadPartial: 14,
+			wantNumPad:      14,
+		}, {
+			desc:            "unaligned start behind tree size, same bundle",
+			integratedSize:  300,
+			uploadStart:     270,
+			wantRead:        true,
+			wantReadIdx:     1,
+			wantReadPartial: 44,
+			wantNumPad:      14,
+		}, {
+			desc:            "unaligned start behind tree size, bundle is full",
+			integratedSize:  testPendingCPSize,
+			uploadStart:     270,
+			wantRead:        true,
+			wantReadIdx:     1,
+			wantReadPartial: 0,
+			wantNumPad:      14,
+		},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			wantBundleIdx := test.uploadStart / layout.EntryBundleWidth
+			// Padding entries are taken from the start of the bundle, and the subtree proof is
+			// verified against the bundle aligned start index.
+			wantEntries := append(append([][]byte{}, storedEntries[:test.wantNumPad]...), uploadedEntries...)
+			wantProofStart := test.uploadStart - uint64(test.wantNumPad)
+
+			var (
+				gotRead     bool
+				gotBundles  [][][]byte
+				gotStart    uint64
+				gotEnd      uint64
+				gotSize     uint64
+				gotVerified bool
+			)
 
-	padEntries := testUploadStart % layout.EntryBundleWidth
-	padBundleRaw := make([]byte, 2*padEntries)
+			drv := &fakeDriver{
+				writer: &fakeMirrorWriter{
+					integrateFunc: func(ctx context.Context, fromBundleIdx uint64, bundles iter.Seq2[*api.EntryBundle, error]) (uint64, []byte, error) {
+						if fromBundleIdx != wantBundleIdx {
+							return 0, nil, fmt.Errorf("got fromBundleIdx %d, want %d", fromBundleIdx, wantBundleIdx)
+						}
+						for b, err := range bundles {
+							if err != nil {
+								return 0, nil, err
+							}
+							gotBundles = append(gotBundles, b.Entries)
+						}
+						pendingCPRoot, err := base64.StdEncoding.DecodeString(testPendingCPRoot)
+						return testPendingCPSize, pendingCPRoot, err
+					},
+					updateCheckpointFunc: func(ctx context.Context, f func(oldCP []byte) (newCP []byte, err error)) error {
+						_, err := f(nil)
+						return err
+					},
+				},
+				reader: &fakeLogReader{
+					sizeFunc: func(ctx context.Context) (uint64, error) { return test.integratedSize, nil },
+					readEntryBundle: func(ctx context.Context, index uint64, p uint8) ([]byte, error) {
+						gotRead = true
+						if index != test.wantReadIdx || p != test.wantReadPartial {
+							// Emulate storage, which only has the resource whose shape matches the tree size.
+							return nil, fmt.Errorf("%w: ReadEntryBundle(%d, %d), want (%d, %d)", os.ErrNotExist, index, p, test.wantReadIdx, test.wantReadPartial)
+						}
+						numEntries := int(p)
+						if p == 0 {
+							numEntries = layout.EntryBundleWidth
+						}
+						return marshalEntryBundle(storedEntries[:numEntries]), nil
+					},
+				},
+			}
 
-	drv := &fakeDriver{
-		writer: &fakeMirrorWriter{
-			integrateFunc: func(ctx context.Context, fromBundleIdx uint64, bundles iter.Seq2[*api.EntryBundle, error]) (uint64, []byte, error) {
-				if want := testUploadStart / layout.EntryBundleWidth; fromBundleIdx != want {
-					return 0, nil, fmt.Errorf("got fromBundleIdx %d want %d", fromBundleIdx, want)
-				}
-				for b, err := range bundles {
-					if err != nil {
-						return 0, nil, err
-					}
-					if got, want := uint64(len(b.Entries)), testUploadStart%layout.EntryBundleWidth; got != want {
-						return 0, nil, fmt.Errorf("got %d entries in bundle, want %d", got, want)
-					}
-				}
-				pendingCPRoot, err := base64.StdEncoding.DecodeString(testPendingCPRoot)
-				return testUploadEnd, pendingCPRoot, err
-			},
-			updateCheckpointFunc: func(ctx context.Context, f func(oldCP []byte) (newCP []byte, err error)) error {
-				_, err := f(nil)
-				return err
-			},
-		},
-		reader: &fakeLogReader{
-			sizeFunc: func(ctx context.Context) (uint64, error) { return testIntegratedSize, nil },
-			readEntryBundle: func(ctx context.Context, index uint64, p uint8) ([]byte, error) {
-				readEntryBundleCalled = true
-				if got, want := index, testUploadStart/layout.EntryBundleWidth; got != want {
-					t.Errorf("ReadEntryBundle index: got %d, want %d", got, want)
-				}
-				if got, want := p, uint8(testUploadStart%layout.EntryBundleWidth); got != want {
-					t.Errorf("ReadEntryBundle p: got %d, want %d", got, want)
-				}
-				return padBundleRaw, nil
-			},
-		},
-	}
+			mt, err := NewMirrorTarget(t.Context(), drv, &MirrorOptions{
+				origin:      testPendingCPOrigin,
+				logVerifier: testLogVerifier,
+				signer:      testMirrorSigner,
+				cpSource:    func(ctx context.Context) ([]byte, error) { return []byte(testPendingCP), nil },
+			})
+			if err != nil {
+				t.Fatalf("NewMirrorTarget() failed: %v", err)
+			}
+			// Stub out proof verification since our entries are arbitrary, but record what it was called with.
+			mt.verifySubtreeProof = func(hasher merkle.LogHasher, start, end, size uint64, proof [][]byte, subRoot, root []byte) error {
+				gotVerified, gotStart, gotEnd, gotSize = true, start, end, size
+				return nil
+			}
+			validTicket, err := mt.seal([]byte(testPendingCP))
+			if err != nil {
+				t.Fatalf("seal failed: %v", err)
+			}
 
-	mt, err := NewMirrorTarget(t.Context(), drv, &MirrorOptions{
-		origin:      testPendingCPOrigin,
-		logVerifier: testLogVerifier,
-		signer:      testMirrorSigner,
-		cpSource:    func(ctx context.Context) ([]byte, error) { return []byte(testPendingCP), nil },
-	})
-	if err != nil {
-		t.Fatalf("NewMirrorTarget() failed: %v", err)
-	}
+			firstCall := true
+			if _, _, _, _, err := mt.AddEntries(t.Context(), test.uploadStart, testPendingCPSize, validTicket, func() (*MirrorPackage, error) {
+				if firstCall {
+					firstCall = false
+					return &MirrorPackage{Entries: uploadedEntries, Proof: [][]byte{[]byte("proof")}}, nil
+				}
+				return nil, io.EOF
+			}); err != nil {
+				t.Fatalf("AddEntries: %v", err)
+			}
 
-	validTicket, err := mt.seal([]byte(testPendingCP))
-	if err != nil {
-		t.Fatalf("seal failed: %v", err)
+			if gotRead != test.wantRead {
+				t.Errorf("ReadEntryBundle called: got %t, want %t", gotRead, test.wantRead)
+			}
+			if len(gotBundles) != 1 {
+				t.Fatalf("got %d integrated bundles, want 1", len(gotBundles))
+			}
+			if diff := cmp.Diff(wantEntries, gotBundles[0]); diff != "" {
+				t.Errorf("integrated bundle entries diff (-want +got):\n%s", diff)
+			}
+			if !gotVerified {
+				t.Fatal("verifySubtreeProof was not called")
+			}
+			if gotStart != wantProofStart {
+				t.Errorf("verifySubtreeProof start: got %d, want %d", gotStart, wantProofStart)
+			}
+			if want := wantProofStart + uint64(len(wantEntries)); gotEnd != want {
+				t.Errorf("verifySubtreeProof end: got %d, want %d", gotEnd, want)
+			}
+			if gotSize != testPendingCPSize {
+				t.Errorf("verifySubtreeProof size: got %d, want %d", gotSize, testPendingCPSize)
+			}
+		})
 	}
+}
 
-	_, _, _, _, err = mt.AddEntries(t.Context(), testUploadStart, testUploadEnd, validTicket, func() (*MirrorPackage, error) {
-		return nil, io.EOF
-	})
-	if err != nil {
-		t.Fatalf("unexpected error: %v", err)
-	}
-	if !readEntryBundleCalled {
-		t.Errorf("ReadEntryBundle was not called")
+// marshalEntryBundle serialises the provided entries into entry bundle format.
+func marshalEntryBundle(entries [][]byte) []byte {
+	r := []byte{}
+	for _, e := range entries {
+		r = binary.BigEndian.AppendUint16(r, uint16(len(e)))
+		r = append(r, e...)
 	}
+	return r
 }
 
 func TestMirrorTarget_AddEntries_NoPendingCheckpoint(t *testing.T) {