adding value or default context helper

projectdiscovery · Jul 21, 2023 · 04c5c27 · 04c5c27
1 parent 5efadaf
commit 04c5c27
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/context/context.go b/context/context.go
@@ -7,8 +7,12 @@ import (
 
 var ErrIncorrectNumberOfItems = errors.New("number of items is not even")
 
+var DefaultContext = context.TODO()
+
+type ContextArg string
+
 // WithValues combines multiple key-value into an existing context
-func WithValues(ctx context.Context, keyValue ...string) (context.Context, error) {
+func WithValues(ctx context.Context, keyValue ...ContextArg) (context.Context, error) {
 	if len(keyValue)%2 != 0 {
 		return ctx, ErrIncorrectNumberOfItems
 	}
@@ -18,3 +22,12 @@ func WithValues(ctx context.Context, keyValue ...string) (context.Context, error
 	}
 	return ctx, nil
 }
+
+// ValueOrDefault returns default context if given is nil (using interface to avoid static check reporting)
+func ValueOrDefault(value interface{}) context.Context {
+	if ctx, ok := value.(context.Context); ok && ctx != nil {
+		return ctx
+	}
+
+	return DefaultContext
+}
diff --git a/context/context_test.go b/context/context_test.go
@@ -10,27 +10,27 @@ import (
 func TestWithValues(t *testing.T) {
 	type testCase struct {
 		name          string
-		keyValue      []string
+		keyValue      []ContextArg
 		expectedError error
 		expectedValue map[string]string
 	}
 
 	var testCases = []testCase{
 		{
 			name:          "even number of key-value pairs",
-			keyValue:      []string{"key1", "value1", "key2", "value2"},
+			keyValue:      []ContextArg{"key1", "value1", "key2", "value2"},
 			expectedError: nil,
 			expectedValue: map[string]string{"key1": "value1", "key2": "value2"},
 		},
 		{
 			name:          "odd number of key-value pairs",
-			keyValue:      []string{"key1", "value1", "key2"},
+			keyValue:      []ContextArg{"key1", "value1", "key2"},
 			expectedError: ErrIncorrectNumberOfItems,
 			expectedValue: map[string]string{},
 		},
 		{
 			name:          "overwriting values",
-			keyValue:      []string{"key1", "value1", "key1", "newValue"},
+			keyValue:      []ContextArg{"key1", "value1", "key1", "newValue"},
 			expectedError: nil,
 			expectedValue: map[string]string{"key1": "newValue"},
 		},