move range check to before reading headers from storage 💣 (#596)

Co-authored-by: Kan Zhang <kan@axiomzen.co>

engine/access/rpc/backend/backend_events.go

@@ -40,6 +40,11 @@ func (b *backendEvents) GetEventsForHeightRange(
 		return nil, status.Error(codes.InvalidArgument, "invalid start or end height")
 	}
 
+	rangeSize := endHeight - startHeight + 1 // range is inclusive on both ends
+	if rangeSize > uint64(b.maxHeightRange) {
+		return nil, fmt.Errorf("requested block range (%d) exceeded maximum (%d)", rangeSize, b.maxHeightRange)
+	}
+
 	// get the latest sealed block header
 	head, err := b.state.Sealed().Head()
 	if err != nil {
@@ -79,6 +84,10 @@ func (b *backendEvents) GetEventsForBlockIDs(
 	blockIDs []flow.Identifier,
 ) ([]flow.BlockEvents, error) {
 
+	if uint(len(blockIDs)) > b.maxHeightRange {
+		return nil, fmt.Errorf("requested block range (%d) exceeded maximum (%d)", len(blockIDs), b.maxHeightRange)
+	}
+
 	// find the block headers for all the block IDs
 	blockHeaders := make([]*flow.Header, 0)
 	for _, blockID := range blockIDs {
@@ -110,11 +119,6 @@ func (b *backendEvents) getBlockEventsFromExecutionNode(
 		return []flow.BlockEvents{}, nil
 	}
 
-	// limit height range queries
-	if uint(len(blockIDs)) > b.maxHeightRange {
-		return nil, fmt.Errorf("requested block range (%d) exceeded maximum (%d)", len(blockIDs), b.maxHeightRange)
-	}
-
 	req := execproto.GetEventsForBlockIDsRequest{
 		Type:     eventType,
 		BlockIds: convert.IdentifiersToMessages(blockIDs),

engine/access/rpc/backend/backend_test.go

@@ -900,28 +900,47 @@ func (suite *Suite) TestGetEventsForHeightRange() {
 	suite.state.On("Sealed").Return(suite.snapshot, nil).Maybe()
 
 	ctx := context.Background()
-	var minHeight uint64 = 5
-	var maxHeight uint64 = 10
+	const minHeight uint64 = 5
+	const maxHeight uint64 = 10
 	var headHeight uint64
 	var blockHeaders []*flow.Header
 
+	headersDB := make(map[uint64]*flow.Header) // backend for storage.Headers
+	var head *flow.Header                      // backend for Snapshot.Head
+
+	// mock snapshot to return head backend
+	suite.snapshot.On("Head").Return(
+		func() *flow.Header { return head },
+		func() error { return nil },
+	).Maybe()
+
+	// mock headers to pull from headers backend
+	suite.headers.On("ByHeight", mock.Anything).Return(
+		func(height uint64) *flow.Header {
+			return headersDB[height]
+		},
+		func(height uint64) error {
+			_, ok := headersDB[height]
+			if !ok {
+				return storage.ErrNotFound
+			}
+			return nil
+		}).Maybe()
+
 	setupHeadHeight := func(height uint64) {
 		header := unittest.BlockHeaderFixture() // create a mock header
 		header.Height = height                  // set the header height
-		suite.snapshot.On("Head").Return(&header, nil).Once()
+		head = &header
 	}
 
 	setupStorage := func(min uint64, max uint64) []*flow.Header {
-		headers := make([]*flow.Header, 0)
+		headersDB = make(map[uint64]*flow.Header) // reset backend
 
+		var headers []*flow.Header
 		for i := min; i <= max; i++ {
-			b := unittest.BlockFixture()
-
-			suite.headers.
-				On("ByHeight", i).
-				Return(b.Header, nil).Once()
-
-			headers = append(headers, b.Header)
+			header := unittest.BlockHeaderFixture()
+			headersDB[i] = &header
+			headers = append(headers, &header)
 		}
 
 		return headers