fix(share/shrex): hotfixes for shrex getter respecting context (cele…

…stiaorg#2115) Closes celestiaorg#2110, related to celestiaorg#2109. Only adds a hotfix for 2109 --------- Co-authored-by: rene <[email protected]> Co-authored-by: Vlad <[email protected]> Co-authored-by: Vlad <[email protected]>
tirom14 · Apr 26, 2023 · fc691ea · fc691ea
1 parent 2d5b385
commit fc691ea
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/share/getters/shrex.go b/share/getters/shrex.go
@@ -67,6 +67,9 @@ func (sg *ShrexGetter) GetEDS(ctx context.Context, root *share.Root) (*rsmt2d.Ex
 		err     error
 	)
 	for {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 		attempt++
 		start := time.Now()
 		peer, setStatus, getErr := sg.peerManager.Peer(ctx, root.Hash())
@@ -120,6 +123,9 @@ func (sg *ShrexGetter) GetSharesByNamespace(
 		err     error
 	)
 	for {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
 		attempt++
 		start := time.Now()
 		peer, setStatus, getErr := sg.peerManager.Peer(ctx, root.Hash())

diff --git a/share/getters/shrex_test.go b/share/getters/shrex_test.go
@@ -104,6 +104,21 @@ func TestShrexGetter(t *testing.T) {
 		require.Equal(t, eds.Flattened(), got.Flattened())
 	})
 
+	t.Run("EDS_ctx_deadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(ctx, time.Second)
+
+		// generate test data
+		_, dah, _ := generateTestEDS(t)
+		peerManager.Validate(ctx, srvHost.ID(), shrexsub.Notification{
+			DataHash: dah.Hash(),
+			Height:   1,
+		})
+
+		cancel()
+		_, err := getter.GetEDS(ctx, &dah)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
 	t.Run("EDS_err_not_found", func(t *testing.T) {
 		ctx, cancel := context.WithTimeout(ctx, time.Second)
 		t.Cleanup(cancel)