From 71ab20ee47b64d485d49f6422b9b22778568853e Mon Sep 17 00:00:00 2001
From: Cody Littley <cody@eigenlabs.org>
Date: Wed, 13 Nov 2024 15:39:04 -0600
Subject: [PATCH] Made suggested changes.

Signed-off-by: Cody Littley <cody@eigenlabs.org>
---
 disperser/apiserver/server.go |  2 +-
 relay/server.go               | 24 ++++++++++++++++++++----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/disperser/apiserver/server.go b/disperser/apiserver/server.go
index 001d663ed1..f7eae5ccff 100644
--- a/disperser/apiserver/server.go
+++ b/disperser/apiserver/server.go
@@ -16,7 +16,7 @@ import (
 	commonpb "github.com/Layr-Labs/eigenda/api/grpc/common"
 	pb "github.com/Layr-Labs/eigenda/api/grpc/disperser"
 	"github.com/Layr-Labs/eigenda/common"
-	healthcheck "github.com/Layr-Labs/eigenda/common/healthcheck"
+	"github.com/Layr-Labs/eigenda/common/healthcheck"
 	"github.com/Layr-Labs/eigenda/core"
 	"github.com/Layr-Labs/eigenda/core/auth"
 	"github.com/Layr-Labs/eigenda/core/meterer"
diff --git a/relay/server.go b/relay/server.go
index 9cd83982a9..e64609c4b7 100644
--- a/relay/server.go
+++ b/relay/server.go
@@ -83,7 +83,12 @@ func (s *Server) GetBlob(ctx context.Context, request *pb.GetBlobRequest) (*pb.G
 	//  - per-connection throttle
 	//  - timeouts
 
-	keys := []v2.BlobKey{v2.BlobKey(request.BlobKey)}
+	key, err := v2.BytesToBlobKey(request.BlobKey)
+	if err != nil {
+		return nil, fmt.Errorf("invalid blob key: %w", err)
+	}
+
+	keys := []v2.BlobKey{key}
 	mMap, err := s.metadataServer.GetMetadataForBlobs(keys)
 	if err != nil {
 		return nil, fmt.Errorf(
@@ -94,7 +99,6 @@ func (s *Server) GetBlob(ctx context.Context, request *pb.GetBlobRequest) (*pb.G
 		return nil, fmt.Errorf("blob not found")
 	}
 
-	key := v2.BlobKey(request.BlobKey)
 	data, err := s.blobServer.GetBlob(key)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching blob %s: %w", key.Hex(), err)
@@ -116,14 +120,26 @@ func (s *Server) GetChunks(ctx context.Context, request *pb.GetChunksRequest) (*
 	//  - per-connection throttle
 	//  - timeouts
 
+	if len(request.ChunkRequests) <= 0 {
+		return nil, fmt.Errorf("no chunk requests provided")
+	}
+
 	keys := make([]v2.BlobKey, 0, len(request.ChunkRequests))
 
 	for _, chunkRequest := range request.ChunkRequests {
 		var key v2.BlobKey
 		if chunkRequest.GetByIndex() != nil {
-			key = v2.BlobKey(chunkRequest.GetByIndex().GetBlobKey())
+			var err error
+			key, err = v2.BytesToBlobKey(chunkRequest.GetByIndex().GetBlobKey())
+			if err != nil {
+				return nil, fmt.Errorf("invalid blob key: %w", err)
+			}
 		} else {
-			key = v2.BlobKey(chunkRequest.GetByRange().GetBlobKey())
+			var err error
+			key, err = v2.BytesToBlobKey(chunkRequest.GetByRange().GetBlobKey())
+			if err != nil {
+				return nil, fmt.Errorf("invalid blob key: %w", err)
+			}
 		}
 		keys = append(keys, key)
 	}