Add Cookie to KeyAuth middleware's KeyLookup

Author: rohitkg98

middleware/key_auth.go

@@ -2,6 +2,7 @@ package middleware
 
 import (
 	"errors"
+	"fmt"
 	"net/http"
 	"strings"
 
@@ -21,6 +22,7 @@ type (
 		// - "header:<name>"
 		// - "query:<name>"
 		// - "form:<name>"
+		// - "cookie:<name>"
 		KeyLookup string `yaml:"key_lookup"`
 
 		// AuthScheme to be used in the Authorization header.
@@ -91,6 +93,8 @@ func KeyAuthWithConfig(config KeyAuthConfig) echo.MiddlewareFunc {
 		extractor = keyFromQuery(parts[1])
 	case "form":
 		extractor = keyFromForm(parts[1])
+	case "cookie":
+		extractor = keyFromCookie(parts[1])
 	}
 
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -164,3 +168,14 @@ func keyFromForm(param string) keyExtractor {
 		return key, nil
 	}
 }
+
+// keyFromCookie returns a `keyExtractor` that extracts key from the form.
+func keyFromCookie(cookieName string) keyExtractor {
+	return func(c echo.Context) (string, error) {
+		key, err := c.Cookie(cookieName)
+		if err != nil {
+			return "", fmt.Errorf("missing key in cookies: %w", err)
+		}
+		return key.Value, nil
+	}
+}

middleware/key_auth_test.go

@@ -157,6 +157,30 @@ func TestKeyAuthWithConfig(t *testing.T) {
 			expectHandlerCalled: false,
 			expectError:         "code=400, message=missing key in the form",
 		},
+		{
+			name: "ok, custom key lookup, cookie",
+			givenRequest: func(req *http.Request) {
+				req.AddCookie(&http.Cookie{
+					Name:  "key",
+					Value: "valid-key",
+				})
+				q := req.URL.Query()
+				q.Add("key", "valid-key")
+				req.URL.RawQuery = q.Encode()
+			},
+			whenConfig: func(conf *KeyAuthConfig) {
+				conf.KeyLookup = "cookie:key"
+			},
+			expectHandlerCalled: true,
+		},
+		{
+			name: "nok, custom key lookup, missing cookie param",
+			whenConfig: func(conf *KeyAuthConfig) {
+				conf.KeyLookup = "cookie:key"
+			},
+			expectHandlerCalled: false,
+			expectError:         "code=400, message=missing key in cookies: http: named cookie not present",
+		},
 		{
 			name: "nok, custom errorHandler, error from extractor",
 			whenConfig: func(conf *KeyAuthConfig) {