From e6503c425b3c1e02a7642de8dd1c88f2e4264c55 Mon Sep 17 00:00:00 2001
From: Lukas Malkmus <lukasmalkmus@users.noreply.github.com>
Date: Tue, 10 Oct 2017 12:29:25 +0200
Subject: [PATCH] Implement custom search functions (#20)

* Implement custom search functions

* Improve function and parameter names

* Preserve original search order
---
 jwtauth.go | 70 ++++++++++++++++++++++++++++++------------------------
 1 file changed, 39 insertions(+), 31 deletions(-)

diff --git a/jwtauth.go b/jwtauth.go
index a9d1b2f..57aa592 100644
--- a/jwtauth.go
+++ b/jwtauth.go
@@ -21,6 +21,34 @@ var (
 	ErrExpired      = errors.New("jwtauth: token is expired")
 )
 
+var (
+	// TokenFromCookie tries to retreive the token string from a cookie named
+	// "jwt".
+	TokenFromCookie = func(r *http.Request) string {
+		cookie, err := r.Cookie("jwt")
+		if err != nil {
+			return ""
+		}
+		return cookie.Value
+	}
+	// TokenFromHeader tries to retreive the token string from the
+	// "Authorization" reqeust header: "Authorization: BEARER T".
+	TokenFromHeader = func(r *http.Request) string {
+		// Get token from authorization header.
+		bearer := r.Header.Get("Authorization")
+		if len(bearer) > 7 && strings.ToUpper(bearer[0:6]) == "BEARER" {
+			return bearer[7:]
+		}
+		return ""
+	}
+	// TokenFromQuery tries to retreive the token string from the "jwt" URI
+	// query parameter.
+	TokenFromQuery = func(r *http.Request) string {
+		// Get token from query param named "jwt".
+		return r.URL.Query().Get("jwt")
+	}
+)
+
 type JwtAuth struct {
 	signKey   interface{}
 	verifyKey interface{}
@@ -68,15 +96,15 @@ func NewWithParser(alg string, parser *jwt.Parser, signKey interface{}, verifyKe
 // http response.
 func Verifier(ja *JwtAuth) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
-		return Verify(ja, "")(next)
+		return Verify(ja, TokenFromQuery, TokenFromHeader, TokenFromCookie)(next)
 	}
 }
 
-func Verify(ja *JwtAuth, paramAliases ...string) func(http.Handler) http.Handler {
+func Verify(ja *JwtAuth, findTokenFns ...func(r *http.Request) string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		hfn := func(w http.ResponseWriter, r *http.Request) {
 			ctx := r.Context()
-			token, err := VerifyRequest(ja, r, paramAliases...)
+			token, err := VerifyRequest(ja, r, findTokenFns...)
 			ctx = NewContext(ctx, token, err)
 			next.ServeHTTP(w, r.WithContext(ctx))
 		}
@@ -84,37 +112,17 @@ func Verify(ja *JwtAuth, paramAliases ...string) func(http.Handler) http.Handler
 	}
 }
 
-func VerifyRequest(ja *JwtAuth, r *http.Request, paramAliases ...string) (*jwt.Token, error) {
+func VerifyRequest(ja *JwtAuth, r *http.Request, findTokenFns ...func(r *http.Request) string) (*jwt.Token, error) {
 	var tokenStr string
 	var err error
 
-	// Get token from query params
-	tokenStr = r.URL.Query().Get("jwt")
-
-	// Get token from other param aliases
-	if tokenStr == "" && paramAliases != nil && len(paramAliases) > 0 {
-		for _, p := range paramAliases {
-			tokenStr = r.URL.Query().Get(p)
-			if tokenStr != "" {
-				break
-			}
-		}
-	}
-
-	// Get token from authorization header
-	if tokenStr == "" {
-		bearer := r.Header.Get("Authorization")
-		if len(bearer) > 7 && strings.ToUpper(bearer[0:6]) == "BEARER" {
-			tokenStr = bearer[7:]
-		}
-	}
-
-	// Get token from cookie
-	if tokenStr == "" {
-		// TODO: paramAliases should apply to cookies too..
-		cookie, err := r.Cookie("jwt")
-		if err == nil {
-			tokenStr = cookie.Value
+	// Extract token string from the request by calling token find functions in
+	// the order they where provided. Further extraction stops if a function
+	// returns a non-empty string.
+	for _, fn := range findTokenFns {
+		tokenStr = fn(r)
+		if tokenStr != "" {
+			break
 		}
 	}