Make client code fallback to requesting full tiles

client/client.go

@@ -163,13 +163,23 @@ func GetEntryBundle(ctx context.Context, f EntryBundleFetcherFunc, i, logSize ui
 	span.SetAttributes(indexKey.Int64(otel.Clamp64(i)), logSizeKey.Int64(otel.Clamp64(logSize)))
 
 	bundle := api.EntryBundle{}
-	sRaw, err := f(ctx, i, layout.PartialTileSize(0, i, logSize))
-	if err != nil {
-		if errors.Is(err, os.ErrNotExist) {
-			return bundle, fmt.Errorf("leaf bundle at index %d not found: %v", i, err)
+	p := layout.PartialTileSize(0, i, logSize)
+	sRaw, err := f(ctx, i, p)
+	switch {
+	case errors.Is(err, os.ErrNotExist) && p == 0:
+		return bundle, fmt.Errorf("full leaf bundle at index %d not found: %v", i, err)
+	case errors.Is(err, os.ErrNotExist) && p > 0:
+		// It could be that the partial bundle was removed as the tree has grown and a full bundle is now present, so try
+		// falling back to that.
+		sRaw, err = f(ctx, i, 0)
+		if err != nil {
+			return bundle, fmt.Errorf("partial bundle at %[1]d.p/%[2]d and full bundle at %[1]d both not found: %[3]w", i, p, err)
 		}
+	case err != nil:
 		return bundle, fmt.Errorf("failed to fetch leaf bundle at index %d: %v", i, err)
+	default:
 	}
+
 	if err := bundle.UnmarshalText(sRaw); err != nil {
 		return bundle, fmt.Errorf("failed to parse EntryBundle at index %d: %v", i, err)
 	}
@@ -403,7 +413,8 @@ func (n *nodeCache) GetNode(ctx context.Context, id compact.NodeID) ([]byte, err
 	t, ok := n.tiles[tKey]
 	if !ok {
 		span.AddEvent("cache miss")
-		tileRaw, err := n.getTile(ctx, tileLevel, tileIndex, layout.PartialTileSize(tileLevel, tileIndex, n.logSize))
+		p := layout.PartialTileSize(tileLevel, tileIndex, n.logSize)
+		tileRaw, err := fetchPartialOrFullTile(ctx, n.getTile, tileLevel, tileIndex, p)
 		if err != nil {
 			return nil, fmt.Errorf("failed to fetch tile: %v", err)
 		}
@@ -430,3 +441,26 @@ func (n *nodeCache) GetNode(ctx context.Context, id compact.NodeID) ([]byte, err
 	}
 	return r.GetRootHash(nil)
 }
+
+// fetchPartialOrFullTile attempts to fetch the tile at the provided coordinates.
+// If no tile is found, and the coordinates refer to a partial tile, fallback to trying the corresponding
+// full tile.
+func fetchPartialOrFullTile(ctx context.Context, f TileFetcherFunc, l, i uint64, p uint8) ([]byte, error) {
+	sRaw, err := f(ctx, l, i, p)
+	switch {
+	case errors.Is(err, os.ErrNotExist) && p == 0:
+		return sRaw, fmt.Errorf("full tile at index %d not found: %w", i, err)
+	case errors.Is(err, os.ErrNotExist) && p > 0:
+		// It could be that the partial tile was removed as the tree has grown and a full tile is now present, so try
+		// falling back to that.
+		sRaw, err = f(ctx, l, i, 0)
+		if err != nil {
+			return sRaw, fmt.Errorf("partial tile at %[1]d/%[2]d.p/%[3]d and full bundle at %[1]d/%[2]d both not found: %[4]w", l, i, p, err)
+		}
+		return sRaw, nil
+	case err != nil:
+		return sRaw, fmt.Errorf("failed to fetch tile at %d/%d(.p/%d]): %v", l, i, p, err)
+	default:
+		return sRaw, nil
+	}
+}

client/client_test.go

@@ -17,6 +17,7 @@ package client
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"errors"
 	"fmt"
 	"os"
@@ -25,6 +26,7 @@ import (
 
 	"github.com/transparency-dev/formats/log"
 	"github.com/transparency-dev/merkle/compact"
+	"github.com/transparency-dev/merkle/proof"
 	"github.com/transparency-dev/tessera/api"
 	"github.com/transparency-dev/tessera/api/layout"
 	"golang.org/x/mod/sumdb/note"
@@ -237,31 +239,46 @@ func TestHandleZeroRoot(t *testing.T) {
 func TestGetEntryBundleAddressing(t *testing.T) {
 	for _, test := range []struct {
 		name                string
-		idx, logSize        uint64
+		idx                 uint64
+		clientLogSize       uint64
+		actualLogSize       uint64
 		wantPartialTileSize uint8
 	}{
 		{
 			name:                "works - partial tile",
 			idx:                 0,
-			logSize:             34,
+			clientLogSize:       34,
+			actualLogSize:       34,
 			wantPartialTileSize: 34,
 		},
 		{
 			name:                "works - full tile",
 			idx:                 1,
-			logSize:             layout.TileWidth*2 + 45,
+			clientLogSize:       layout.TileWidth*2 + 45,
+			actualLogSize:       layout.TileWidth*2 + 45,
 			wantPartialTileSize: 0,
 		},
+		{
+			name:                "works - request partial but fallback to full tile",
+			idx:                 3,                       // Request the partial bundle at the end of the log
+			clientLogSize:       layout.TileWidth*2 + 45, // bundle 3 is partial according to client's PoV
+			actualLogSize:       layout.TileWidth * 3,    // but the log has grown and bundle 3 is now full.
+			wantPartialTileSize: 0,                       // so we expect the last call to the fetcher to be for a full bundle.
+		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			gotIdx := uint64(0)
 			gotTileSize := uint8(0)
 			f := func(_ context.Context, i uint64, sz uint8) ([]byte, error) {
 				gotIdx = i
 				gotTileSize = sz
+				p := layout.PartialTileSize(0, i, test.actualLogSize)
+				if p != sz {
+					return nil, os.ErrNotExist
+				}
 				return []byte{}, nil
 			}
-			_, err := GetEntryBundle(context.Background(), f, test.idx, test.logSize)
+			_, err := GetEntryBundle(context.Background(), f, test.idx, test.clientLogSize)
 			if err != nil {
 				t.Fatalf("GetEntryBundle: %v", err)
 			}
@@ -274,3 +291,76 @@ func TestGetEntryBundleAddressing(t *testing.T) {
 		})
 	}
 }
+
+func TestNodeFetcherAddressing(t *testing.T) {
+	for _, test := range []struct {
+		name                string
+		nodeLevel           uint
+		nodeIdx             uint64
+		clientLogSize       uint64
+		actualLogSize       uint64
+		wantPartialTileSize uint8
+	}{
+		{
+			name:                "works - partial tile",
+			nodeIdx:             0,
+			clientLogSize:       34,
+			actualLogSize:       34,
+			wantPartialTileSize: 34,
+		},
+		{
+			name:                "works - full tile",
+			nodeIdx:             56,
+			clientLogSize:       layout.TileWidth*2 + 45,
+			actualLogSize:       layout.TileWidth*2 + 45,
+			wantPartialTileSize: 0,
+		},
+		{
+			name:                "works - request partial but fallback to full tile",
+			nodeIdx:             3*layout.TileWidth + 23, // Request node from the partial tile at the end of the log
+			clientLogSize:       layout.TileWidth*2 + 45, // tile 3 is partial according to client's PoV
+			actualLogSize:       layout.TileWidth * 3,    // but the log has grown and tile 3 is now full.
+			wantPartialTileSize: 0,                       // so we expect the last call to the fetcher to be for a full tile.
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			gotLevel, gotIdx, gotTileSize := uint(0), uint64(0), uint8(0)
+			f := func(_ context.Context, l, i uint64, sz uint8) ([]byte, error) {
+				gotLevel = uint(l)
+				gotIdx = i
+				gotTileSize = sz
+				p := layout.PartialTileSize(l, i, test.actualLogSize)
+				if p != sz {
+					return nil, os.ErrNotExist
+				}
+				r := api.HashTile{}
+				s := int(sz)
+				if s == 0 {
+					s = layout.TileWidth
+				}
+				for x := range s {
+					h := sha256.Sum256(fmt.Appendf(nil, "node at %d/%d", l, i+uint64(x)))
+					r.Nodes = append(r.Nodes, h[:])
+				}
+				return r.MarshalText()
+			}
+			pb, err := NewProofBuilder(t.Context(), test.clientLogSize, f)
+			if err != nil {
+				t.Fatalf("NewProofBuilder: %v", err)
+			}
+			_, err = pb.fetchNodes(t.Context(), proof.Nodes{IDs: []compact.NodeID{compact.NewNodeID(test.nodeLevel, test.nodeIdx)}})
+			if err != nil {
+				t.Fatalf("fetchNodes: %v", err)
+			}
+			if wantLevel := test.nodeLevel >> layout.TileHeight; gotLevel != wantLevel {
+				t.Errorf("f got level %d, want %d", gotLevel, wantLevel)
+			}
+			if wantIdx := test.nodeIdx >> layout.TileHeight; gotIdx != wantIdx {
+				t.Errorf("f got idx %d, want %d", gotIdx, wantIdx)
+			}
+			if gotTileSize != test.wantPartialTileSize {
+				t.Errorf("f got tileSize %d, want %d", gotTileSize, test.wantPartialTileSize)
+			}
+		})
+	}
+}