From 6446628350b86c8b0de927adb0301ca7f2cbc3bf Mon Sep 17 00:00:00 2001
From: Michael Gaffney <mike@gaffney.cc>
Date: Wed, 4 Sep 2019 09:11:41 -0400
Subject: [PATCH] Exit ScanView if context has been cancelled

---
 sdk/logical/storage.go      |  4 ++++
 sdk/logical/storage_test.go | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/sdk/logical/storage.go b/sdk/logical/storage.go
index 477d65a0936f..0802ad01a0f6 100644
--- a/sdk/logical/storage.go
+++ b/sdk/logical/storage.go
@@ -74,6 +74,10 @@ func ScanView(ctx context.Context, view ClearableView, cb func(path string)) err
 
 		// Handle the contents in the directory
 		for _, c := range contents {
+			// Exit if the context has been canceled
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
 			fullPath := current + c
 			if strings.HasSuffix(c, "/") {
 				frontier = append(frontier, fullPath)
diff --git a/sdk/logical/storage_test.go b/sdk/logical/storage_test.go
index aea4e8095d46..c8ba5fc46ce6 100644
--- a/sdk/logical/storage_test.go
+++ b/sdk/logical/storage_test.go
@@ -34,6 +34,24 @@ func TestScanView(t *testing.T) {
 	}
 }
 
+func TestScanView_CancelContext(t *testing.T) {
+	s := prepKeyStorage(t)
+
+	ctx, cancelCtx := context.WithCancel(context.Background())
+	var i int
+	err := ScanView(ctx, s, func(path string) {
+		cancelCtx()
+		i++
+	})
+
+	if err == nil {
+		t.Error("Want context cancel err, got none")
+	}
+	if i != 1 {
+		t.Errorf("Want i==1, got %d", i)
+	}
+}
+
 func TestCollectKeys(t *testing.T) {
 	s := prepKeyStorage(t)