Tighten client nodecache

client/client.go

@@ -462,6 +462,9 @@ func newNodeCache(f TileFetcherFunc, logSize uint64) *nodeCache {
 func (n *nodeCache) GetNode(ctx context.Context, id compact.NodeID) ([]byte, error) {
 	return otel.Trace(ctx, "tessera.client.nodecache.GetNode", tracer, func(ctx context.Context, span trace.Span) ([]byte, error) {
 		span.SetAttributes(indexKey.Int64(otel.Clamp64(id.Index)), levelKey.Int64(int64(id.Level)))
+		if id.Index >= n.logSize>>id.Level {
+			return nil, fmt.Errorf("node %+v out of range for log of size %d", id, n.logSize)
+		}
 		// Fast-path: check to see we have this node in the cache and return it directly if so, otherwise we'll need to fetch it.
 		if e, ok := n.nodes.Get(id); ok {
 			return e, nil
@@ -483,7 +486,7 @@ func (n *nodeCache) GetNode(ctx context.Context, id compact.NodeID) ([]byte, err
 		p := layout.PartialTileSize(tileLevel, tileIndex, n.logSize)
 		nodes, err := n.fetchTileNodes(ctx, tileLevel, tileIndex, p)
 		if err != nil {
-			return nil, fmt.Errorf("failed to fetch and populate node cache: %v", err)
+			return nil, fmt.Errorf("failed to fetch and populate node cache: %w", err)
 		}
 		for k, v := range nodes {
 			n.nodes.Add(k, v)
@@ -520,15 +523,34 @@ func (n *nodeCache) fetchTileNodes(ctx context.Context, tileLevel, tileIndex uin
 	return otel.Trace(ctx, "tessera.client.nodecache.fetchTileNodes", tracer, func(ctx context.Context, span trace.Span) (map[compact.NodeID][]byte, error) {
 		tileRaw, err := n.getTile(ctx, tileLevel, tileIndex, p)
 		if err != nil {
-			return nil, fmt.Errorf("failed to fetch tile: %v", err)
+			return nil, fmt.Errorf("failed to fetch tile: %w", err)
 		}
 
 		var tile api.HashTile
 		if err := tile.UnmarshalText(tileRaw); err != nil {
 			return nil, fmt.Errorf("failed to parse tile: %v", err)
 		}
 
-		ret := make(map[compact.NodeID][]byte, 256*2-1)
+		wantSize := layout.TileWidth
+		if p > 0 {
+			wantSize = int(p)
+		}
+		switch gotLen := len(tile.Nodes); gotLen {
+		case wantSize:
+			// We got the exact size tile we asked for, nothing extra to do.
+		case layout.TileWidth:
+			// Must have asked for a partial tile and got a full tile back.
+			// Trim the full tile down to the size of the requested partial tile.
+			tile.Nodes = tile.Nodes[:wantSize]
+		default:
+			additional := ""
+			if wantSize < layout.TileWidth {
+				additional = fmt.Sprintf(" or %d", layout.TileWidth)
+			}
+			return nil, fmt.Errorf("invalid tile: expected %d%s nodes, got %d", wantSize, additional, gotLen)
+		}
+
+		ret := make(map[compact.NodeID][]byte, wantSize*2-1)
 		// visitFn is a visitor callback which populates the nodes cache.
 		// Used by the calls to compact range below.
 		visitFn := func(intID compact.NodeID, h []byte) {
@@ -546,9 +568,6 @@ func (n *nodeCache) fetchTileNodes(ctx context.Context, tileLevel, tileIndex uin
 				return nil, fmt.Errorf("failed to Append: %v", err)
 			}
 		}
-		if _, err := r.GetRootHash(visitFn); err != nil {
-			return nil, fmt.Errorf("failed to visit all nodes: %v", err)
-		}
 		return ret, nil
 	})
 }

client/client_test.go

@@ -198,28 +198,100 @@ func TestCheckLogStateTracker(t *testing.T) {
 	}
 }
 
-func TestNodeCacheHandlesInvalidRequest(t *testing.T) {
-	ctx := context.Background()
-	wantBytes := []byte("0123456789ABCDEF0123456789ABCDEF")
-	f := func(_ context.Context, _, _ uint64, _ uint8) ([]byte, error) {
-		h := &api.HashTile{
-			Nodes: [][]byte{wantBytes},
-		}
-		return h.MarshalText()
-	}
-
-	// Large tree, but we're emulating skew since f, above, will return a tile which only knows about 1
-	// leaf.
-	nc := newNodeCache(f, 10)
+func TestNodeCache(t *testing.T) {
+	ctx := t.Context()
+	leafHash := []byte("0123456789ABCDEF0123456789ABCDEF")
 
-	if got, err := nc.GetNode(ctx, compact.NewNodeID(0, 0)); err != nil {
-		t.Errorf("got %v, want no error", err)
-	} else if !bytes.Equal(got, wantBytes) {
-		t.Errorf("got %v, want %v", got, wantBytes)
-	}
+	for _, test := range []struct {
+		desc          string
+		logSize       uint64
+		tileNodes     int
+		reqID         compact.NodeID
+		wantErr       bool
+		wantCacheSize int
+	}{
+		{
+			desc:          "valid single leaf",
+			logSize:       1,
+			tileNodes:     1,
+			reqID:         compact.NewNodeID(0, 0),
+			wantCacheSize: 1,
+		},
+		{
+			desc:      "out of range leaf in same tile",
+			logSize:   1,
+			tileNodes: 1,
+			reqID:     compact.NewNodeID(0, 1),
+			wantErr:   true,
+		},
+		{
+			desc:      "out of range leaf in next tile",
+			logSize:   layout.TileWidth,
+			tileNodes: layout.TileWidth,
+			reqID:     compact.NewNodeID(0, layout.TileWidth),
+			wantErr:   true,
+		},
+		{
+			desc:      "out of range high level overflow check",
+			logSize:   layout.TileWidth,
+			tileNodes: layout.TileWidth,
+			reqID:     compact.NewNodeID(64, 0),
+			wantErr:   true,
+		},
+		{
+			desc:      "truncated partial tile",
+			logSize:   10,
+			tileNodes: 1,
+			reqID:     compact.NewNodeID(0, 0),
+			wantErr:   true,
+		},
+		{
+			desc:      "truncated full tile",
+			logSize:   layout.TileWidth,
+			tileNodes: layout.TileWidth - 1,
+			reqID:     compact.NewNodeID(0, 0),
+			wantErr:   true,
+		},
+		{
+			desc:      "oversized tile",
+			logSize:   layout.TileWidth,
+			tileNodes: layout.TileWidth + 1,
+			reqID:     compact.NewNodeID(0, 0),
+			wantErr:   true,
+		},
+		{
+			desc:          "partial tile fallback to full tile truncates and omits ephemeral nodes",
+			logSize:       3,
+			tileNodes:     layout.TileWidth,
+			reqID:         compact.NewNodeID(1, 0),
+			wantCacheSize: 4, // (0,0), (0,1), (1,0), (0,2) - no leaves >= 3 and no ephemeral (2,0)
+		},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+			f := func(_ context.Context, _, _ uint64, _ uint8) ([]byte, error) {
+				nodes := make([][]byte, test.tileNodes)
+				for i := range nodes {
+					nodes[i] = leafHash
+				}
+				return (&api.HashTile{Nodes: nodes}).MarshalText()
+			}
 
-	if _, err := nc.GetNode(ctx, compact.NewNodeID(0, 1)); err == nil {
-		t.Error("got no error, want error because ID is out of range")
+			nc := newNodeCache(f, test.logSize)
+			got, err := nc.GetNode(ctx, test.reqID)
+			if gotErr := err != nil; gotErr != test.wantErr {
+				t.Fatalf("GetNode(%+v) err = %v, wantErr %t", test.reqID, err, test.wantErr)
+			}
+			if test.wantErr {
+				return
+			}
+			if len(got) == 0 {
+				t.Errorf("GetNode(%+v) returned empty hash", test.reqID)
+			}
+			if gotSize := nc.nodes.Len(); gotSize != test.wantCacheSize {
+				t.Errorf("cache size = %d, want %d", gotSize, test.wantCacheSize)
+			}
+		})
 	}
 }