refactor: use metadata store to check for blob existence (#1380)

hopeyen · web-flow · commit 54a44b3fbc90 · 2025-03-13T19:00:48.000-05:00
diff --git a/common/aws/dynamodb/client.go b/common/aws/dynamodb/client.go
@@ -55,6 +55,7 @@ type Client interface {
 	UpdateItemWithCondition(ctx context.Context, tableName string, key Key, item Item, condition expression.ConditionBuilder) (Item, error)
 	IncrementBy(ctx context.Context, tableName string, key Key, attr string, value uint64) (Item, error)
 	GetItem(ctx context.Context, tableName string, key Key) (Item, error)
+	GetItemWithInput(ctx context.Context, input *dynamodb.GetItemInput) (Item, error)
 	GetItems(ctx context.Context, tableName string, keys []Key, consistentRead bool) ([]Item, error)
 	QueryIndex(ctx context.Context, tableName string, indexName string, keyCondition string, expAttributeValues ExpressionValues) ([]Item, error)
 	Query(ctx context.Context, tableName string, keyCondition string, expAttributeValues ExpressionValues) ([]Item, error)
@@ -269,6 +270,16 @@ func (c *client) GetItem(ctx context.Context, tableName string, key Key) (Item,
 	return resp.Item, nil
 }
 
+// GetItemWithInput is a wrapper for the GetItem function that allows for a custom GetItemInput
+func (c *client) GetItemWithInput(ctx context.Context, input *dynamodb.GetItemInput) (Item, error) {
+	resp, err := c.dynamoClient.GetItem(ctx, input)
+	if err != nil {
+		return nil, err
+	}
+
+	return resp.Item, nil
+}
+
 // GetItems returns the items for the given keys
 // Note: ordering of items is not guaranteed
 func (c *client) GetItems(ctx context.Context, tableName string, keys []Key, consistentRead bool) ([]Item, error) {
diff --git a/common/aws/mock/dynamodb_client.go b/common/aws/mock/dynamodb_client.go
@@ -59,6 +59,11 @@ func (c *MockDynamoDBClient) GetItem(ctx context.Context, tableName string, key
 	return args.Get(0).(dynamodb.Item), args.Error(1)
 }
 
+func (c *MockDynamoDBClient) GetItemWithInput(ctx context.Context, input *awsdynamodb.GetItemInput) (dynamodb.Item, error) {
+	args := c.Called()
+	return args.Get(0).(dynamodb.Item), args.Error(1)
+}
+
 func (c *MockDynamoDBClient) GetItems(ctx context.Context, tableName string, keys []dynamodb.Key, consistentRead bool) ([]dynamodb.Item, error) {
 	args := c.Called()
 	return args.Get(0).([]dynamodb.Item), args.Error(1)
diff --git a/disperser/apiserver/disperse_blob_v2.go b/disperser/apiserver/disperse_blob_v2.go
@@ -29,10 +29,15 @@ func (s *DispersalServerV2) DisperseBlob(ctx context.Context, req *pb.DisperseBl
 	if onchainState == nil {
 		return nil, api.NewErrorInternal("onchain state is nil")
 	}
-	if err := s.validateDispersalRequest(ctx, req, onchainState); err != nil {
+	blobHeader, err := s.validateDispersalRequest(req, onchainState)
+	if err != nil {
 		return nil, api.NewErrorInvalidArg(fmt.Sprintf("failed to validate the request: %v", err))
 	}
 
+	if err := s.checkBlobExistence(ctx, blobHeader); err != nil {
+		return nil, err
+	}
+
 	// Check against payment meter to make sure there is quota remaining
 	if err := s.checkPaymentMeter(ctx, req, start); err != nil {
 		return nil, err
@@ -43,10 +48,6 @@ func (s *DispersalServerV2) DisperseBlob(ctx context.Context, req *pb.DisperseBl
 
 	blob := req.GetBlob()
 	s.metrics.reportDisperseBlobSize(len(blob))
-	blobHeader, err := corev2.BlobHeaderFromProtobuf(req.GetBlobHeader())
-	if err != nil {
-		return nil, api.NewErrorInvalidArg(fmt.Sprintf("failed to parse the blob header proto: %v", err))
-	}
 	s.logger.Debug("received a new blob dispersal request", "blobSizeBytes", len(blob), "quorums", req.GetBlobHeader().GetQuorumNumbers())
 
 	blobKey, err := s.StoreBlob(ctx, blob, blobHeader, req.GetSignature(), time.Now(), onchainState.TTL)
@@ -133,100 +134,110 @@ func (s *DispersalServerV2) checkPaymentMeter(ctx context.Context, req *pb.Dispe
 }
 
 func (s *DispersalServerV2) validateDispersalRequest(
-	ctx context.Context,
 	req *pb.DisperseBlobRequest,
-	onchainState *OnchainState) error {
+	onchainState *OnchainState) (*corev2.BlobHeader, error) {
 
 	signature := req.GetSignature()
 	if len(signature) != 65 {
-		return fmt.Errorf("signature is expected to be 65 bytes, but got %d bytes", len(signature))
+		return nil, fmt.Errorf("signature is expected to be 65 bytes, but got %d bytes", len(signature))
 	}
 	blob := req.GetBlob()
 	blobSize := len(blob)
 	if blobSize == 0 {
-		return errors.New("blob size must be greater than 0")
+		return nil, errors.New("blob size must be greater than 0")
 	}
 	blobLength := encoding.GetBlobLengthPowerOf2(uint(blobSize))
 	if blobLength > uint(s.maxNumSymbolsPerBlob) {
-		return errors.New("blob size too big")
+		return nil, errors.New("blob size too big")
 	}
 
 	blobHeaderProto := req.GetBlobHeader()
 	if blobHeaderProto.GetCommitment() == nil {
-		return errors.New("blob header must contain commitments")
+		return nil, errors.New("blob header must contain commitments")
 	}
 
 	if blobHeaderProto.GetCommitment() == nil {
-		return errors.New("blob header must contain a commitment")
+		return nil, errors.New("blob header must contain a commitment")
 	}
 	commitedBlobLength := blobHeaderProto.GetCommitment().GetLength()
 	if commitedBlobLength == 0 || commitedBlobLength != encoding.NextPowerOf2(commitedBlobLength) {
-		return errors.New("invalid commitment length, must be a power of 2")
+		return nil, errors.New("invalid commitment length, must be a power of 2")
 	}
 	lengthPowerOf2 := encoding.GetBlobLengthPowerOf2(uint(blobSize))
 	if lengthPowerOf2 > uint(commitedBlobLength) {
-		return fmt.Errorf("commitment length %d is less than blob length %d", commitedBlobLength, lengthPowerOf2)
+		return nil, fmt.Errorf("commitment length %d is less than blob length %d", commitedBlobLength, lengthPowerOf2)
 	}
 
 	blobHeader, err := corev2.BlobHeaderFromProtobuf(blobHeaderProto)
 	if err != nil {
-		return fmt.Errorf("invalid blob header: %w", err)
+		return nil, fmt.Errorf("invalid blob header: %w", err)
 	}
 
 	if blobHeader.PaymentMetadata == (core.PaymentMetadata{}) {
-		return errors.New("payment metadata is required")
+		return nil, errors.New("payment metadata is required")
 	}
 
 	timestampIsNegative := blobHeader.PaymentMetadata.Timestamp < 0
 	paymentIsNegative := blobHeader.PaymentMetadata.CumulativePayment.Cmp(big.NewInt(0)) == -1
 	timestampIsZeroAndPaymentIsZero := blobHeader.PaymentMetadata.Timestamp == 0 && blobHeader.PaymentMetadata.CumulativePayment.Cmp(big.NewInt(0)) == 0
 	if timestampIsNegative || paymentIsNegative || timestampIsZeroAndPaymentIsZero {
-		return errors.New("invalid payment metadata")
+		return nil, errors.New("invalid payment metadata")
 	}
 
 	if len(blobHeaderProto.GetQuorumNumbers()) == 0 {
-		return errors.New("blob header must contain at least one quorum number")
+		return nil, errors.New("blob header must contain at least one quorum number")
 	}
 
 	if len(blobHeaderProto.GetQuorumNumbers()) > int(onchainState.QuorumCount) {
-		return fmt.Errorf("too many quorum numbers specified: maximum is %d", onchainState.QuorumCount)
+		return nil, fmt.Errorf("too many quorum numbers specified: maximum is %d", onchainState.QuorumCount)
 	}
 
 	for _, quorum := range blobHeaderProto.GetQuorumNumbers() {
 		if quorum > corev2.MaxQuorumID || uint8(quorum) >= onchainState.QuorumCount {
-			return fmt.Errorf("invalid quorum number %d; maximum is %d", quorum, onchainState.QuorumCount)
+			return nil, fmt.Errorf("invalid quorum number %d; maximum is %d", quorum, onchainState.QuorumCount)
 		}
 	}
 
 	// validate every 32 bytes is a valid field element
 	_, err = rs.ToFrArray(blob)
 	if err != nil {
 		s.logger.Error("failed to convert a 32bytes as a field element", "err", err)
-		return errors.New("encountered an error to convert a 32-bytes into a valid field element, please use the correct format where every 32bytes(big-endian) is less than 21888242871839275222246405745257275088548364400416034343698204186575808495617")
+		return nil, errors.New("encountered an error to convert a 32-bytes into a valid field element, please use the correct format where every 32bytes(big-endian) is less than 21888242871839275222246405745257275088548364400416034343698204186575808495617")
 	}
 
 	if _, ok := onchainState.BlobVersionParameters.Get(corev2.BlobVersion(blobHeaderProto.GetVersion())); !ok {
-		return fmt.Errorf("invalid blob version %d; valid blob versions are: %v", blobHeaderProto.GetVersion(), onchainState.BlobVersionParameters.Keys())
+		return nil, fmt.Errorf("invalid blob version %d; valid blob versions are: %v", blobHeaderProto.GetVersion(), onchainState.BlobVersionParameters.Keys())
 	}
 
 	if err = s.authenticator.AuthenticateBlobRequest(blobHeader, signature); err != nil {
-		return fmt.Errorf("authentication failed: %w", err)
+		return nil, fmt.Errorf("authentication failed: %w", err)
 	}
 
 	commitments, err := s.prover.GetCommitmentsForPaddedLength(blob)
 	if err != nil {
-		return fmt.Errorf("failed to get commitments: %w", err)
+		return nil, fmt.Errorf("failed to get commitments: %w", err)
 	}
 	if !commitments.Equal(&blobHeader.BlobCommitments) {
-		return errors.New("invalid blob commitment")
+		return nil, errors.New("invalid blob commitment")
 	}
 
+	return blobHeader, nil
+}
+
+func (s *DispersalServerV2) checkBlobExistence(ctx context.Context, blobHeader *corev2.BlobHeader) error {
 	blobKey, err := blobHeader.BlobKey()
 	if err != nil {
-		return fmt.Errorf("failed to get blob key: %w", err)
+		return api.NewErrorInvalidArg(fmt.Sprintf("failed to parse the blob header: %v", err))
+	}
+
+	// check if blob already exists
+	exists, err := s.blobMetadataStore.CheckBlobExists(ctx, blobKey)
+	if err != nil {
+		return api.NewErrorInternal(fmt.Sprintf("failed to check blob existence: %v", err))
 	}
-	if s.blobStore.CheckBlobExists(ctx, blobKey) {
-		return fmt.Errorf("blob already exists: %s", blobKey.Hex())
+
+	if exists {
+		return api.NewErrorAlreadyExists(fmt.Sprintf("blob already exists: %s", blobKey.Hex()))
 	}
 
 	return nil
diff --git a/disperser/common/v2/blobstore/dynamo_metadata_store.go b/disperser/common/v2/blobstore/dynamo_metadata_store.go
@@ -265,6 +265,30 @@ func (s *BlobMetadataStore) GetBlobMetadata(ctx context.Context, blobKey corev2.
 	return metadata, nil
 }
 
+// CheckBlobExists checks if a blob exists without fetching the entire metadata.
+func (s *BlobMetadataStore) CheckBlobExists(ctx context.Context, blobKey corev2.BlobKey) (bool, error) {
+	input := &dynamodb.GetItemInput{
+		TableName: aws.String(s.tableName),
+		Key: map[string]types.AttributeValue{
+			"PK": &types.AttributeValueMemberS{
+				Value: blobKeyPrefix + blobKey.Hex(),
+			},
+			"SK": &types.AttributeValueMemberS{
+				Value: blobMetadataSK,
+			},
+		},
+		ProjectionExpression: aws.String("PK"), // Only fetch the PK attribute
+	}
+
+	item, err := s.dynamoDBClient.GetItemWithInput(ctx, input)
+	if err != nil {
+		return false, fmt.Errorf("failed to check blob existence: %w", err)
+	}
+
+	// If the item is not nil, the blob exists
+	return item != nil, nil
+}
+
 // GetBlobMetadataByStatus returns all the metadata with the given status that were updated after lastUpdatedAt
 // Because this function scans the entire index, it should only be used for status with a limited number of items.
 // Results are ordered by UpdatedAt in ascending order.
diff --git a/disperser/common/v2/blobstore/dynamo_metadata_store_test.go b/disperser/common/v2/blobstore/dynamo_metadata_store_test.go
@@ -1934,3 +1934,53 @@ func newBlob(t *testing.T) (corev2.BlobKey, *corev2.BlobHeader) {
 	require.NoError(t, err)
 	return bk, bh
 }
+
+func TestCheckBlobExists(t *testing.T) {
+	ctx := context.Background()
+	// Create a test blob
+	blobKey, blobHeader := newBlob(t)
+
+	// Check that the blob does not exist initially
+	exists, err := blobMetadataStore.CheckBlobExists(ctx, blobKey)
+	require.NoError(t, err)
+	require.False(t, exists, "Blob should not exist before being added")
+
+	// Create blob metadata
+	blobMetadata := &v2.BlobMetadata{
+		BlobHeader:  blobHeader,
+		Signature:   []byte("test-signature"),
+		BlobStatus:  v2.Queued,
+		Expiry:      uint64(time.Now().Add(time.Hour).Unix()),
+		NumRetries:  0,
+		BlobSize:    1024,
+		RequestedAt: uint64(time.Now().UnixNano()),
+		UpdatedAt:   uint64(time.Now().UnixNano()),
+	}
+
+	// Store the blob metadata
+	err = blobMetadataStore.PutBlobMetadata(ctx, blobMetadata)
+	require.NoError(t, err)
+
+	// Check that the blob now exists
+	exists, err = blobMetadataStore.CheckBlobExists(ctx, blobKey)
+	require.NoError(t, err)
+	require.True(t, exists, "Blob should exist after being added")
+
+	// Delete the blob metadata
+	err = blobMetadataStore.DeleteBlobMetadata(ctx, blobKey)
+	require.NoError(t, err)
+
+	// Check that the blob no longer exists
+	exists, err = blobMetadataStore.CheckBlobExists(ctx, blobKey)
+	require.NoError(t, err)
+	require.False(t, exists, "Blob should not exist after being deleted")
+
+	// Test with non-existent blob key
+	randomKey := corev2.BlobKey{}
+	_, err = rand.Read(randomKey[:])
+	require.NoError(t, err)
+
+	exists, err = blobMetadataStore.CheckBlobExists(ctx, randomKey)
+	require.NoError(t, err)
+	require.False(t, exists, "Random blob key should not exist")
+}
diff --git a/disperser/common/v2/blobstore/s3_blob_store.go b/disperser/common/v2/blobstore/s3_blob_store.go
@@ -26,25 +26,20 @@ func NewBlobStore(s3BucketName string, s3Client s3.Client, logger logging.Logger
 
 // StoreBlob adds a blob to the blob store
 func (b *BlobStore) StoreBlob(ctx context.Context, key corev2.BlobKey, data []byte) error {
-	if b.CheckBlobExists(ctx, key) {
+	_, err := b.s3Client.HeadObject(ctx, b.bucketName, s3.ScopedBlobKey(key))
+	if err == nil {
 		b.logger.Warnf("blob already exists in bucket %s: %s", b.bucketName, key)
 		return common.ErrAlreadyExists
 	}
 
-	err := b.s3Client.UploadObject(ctx, b.bucketName, s3.ScopedBlobKey(key), data)
+	err = b.s3Client.UploadObject(ctx, b.bucketName, s3.ScopedBlobKey(key), data)
 	if err != nil {
 		b.logger.Errorf("failed to upload blob in bucket %s: %v", b.bucketName, err)
 		return err
 	}
 	return nil
 }
 
-// CheckBlobExists checks if a blob exists in the blob store
-func (b *BlobStore) CheckBlobExists(ctx context.Context, key corev2.BlobKey) bool {
-	_, err := b.s3Client.HeadObject(ctx, b.bucketName, s3.ScopedBlobKey(key))
-	return err == nil
-}
-
 // GetBlob retrieves a blob from the blob store
 func (b *BlobStore) GetBlob(ctx context.Context, key corev2.BlobKey) ([]byte, error) {
 	data, err := b.s3Client.DownloadObject(ctx, b.bucketName, s3.ScopedBlobKey(key))
