From 685031a90d9bfa29e30dd843568e484f358604ca Mon Sep 17 00:00:00 2001
From: Ian Shim <100327837+ian-shim@users.noreply.github.com>
Date: Wed, 13 Nov 2024 11:43:42 -0800
Subject: [PATCH] [disperser] Return non-internal error for context errors
 (#887)

---
 api/errors.go                 |  6 +++++-
 disperser/apiserver/server.go | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/api/errors.go b/api/errors.go
index bd4c278842..42632da105 100644
--- a/api/errors.go
+++ b/api/errors.go
@@ -55,7 +55,11 @@ func NewErrorUnimplemented() error {
 
 // HTTP Mapping: 504 Gateway Timeout
 func NewErrorDeadlineExceeded(msg string) error {
-	return newErrorGRPC(codes.DeadlineExceeded, "msg")
+	return newErrorGRPC(codes.DeadlineExceeded, msg)
+}
+
+func NewErrorCanceled(msg string) error {
+	return newErrorGRPC(codes.Canceled, msg)
 }
 
 // ErrorFailover is returned by the disperser-client and eigenda-client to signify
diff --git a/disperser/apiserver/server.go b/disperser/apiserver/server.go
index 001d663ed1..793de14857 100644
--- a/disperser/apiserver/server.go
+++ b/disperser/apiserver/server.go
@@ -300,6 +300,9 @@ func (s *DispersalServer) disperseBlob(ctx context.Context, blob *core.Blob, aut
 
 	requestedAt := uint64(time.Now().UnixNano())
 	metadataKey, err := s.blobStore.StoreBlob(ctx, blob, requestedAt)
+	if ctxErr := contextError(err); ctxErr != nil {
+		return nil, ctxErr
+	}
 	if err != nil {
 		for _, param := range securityParams {
 			s.metrics.HandleBlobStoreFailedRequest(fmt.Sprintf("%d", param.QuorumID), blobSize, apiMethodName)
@@ -1041,3 +1044,15 @@ func (s *DispersalServer) validateRequestAndGetBlob(ctx context.Context, req *pb
 
 	return blob, nil
 }
+
+func contextError(err error) error {
+	if errors.Is(err, context.DeadlineExceeded) {
+		return api.NewErrorDeadlineExceeded(err.Error())
+	}
+
+	if errors.Is(err, context.Canceled) {
+		return api.NewErrorCanceled(err.Error())
+	}
+
+	return nil
+}