From a8b48af60843c288391bb115ba75d459f8fb29a5 Mon Sep 17 00:00:00 2001
From: Przemek Komosa <przemek.komosa@gmail.com>
Date: Sat, 2 Jul 2016 11:26:05 +0200
Subject: [PATCH] Read JWT from cookie

Note: writing token to cookie is out of scope of this middleware
(exactly as in other reading methods).
---
 middleware/jwt.go      | 19 ++++++++++++++++++-
 middleware/jwt_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/middleware/jwt.go b/middleware/jwt.go
index 6f09bab1..be85535a 100644
--- a/middleware/jwt.go
+++ b/middleware/jwt.go
@@ -31,6 +31,7 @@ type (
 		// Possible values:
 		// - "header:<name>"
 		// - "query:<name>"
+		// - "cookie:<name>"
 		TokenLookup string `json:"token_lookup"`
 
 		// HandleEmptyToken is handler executed when there is no token.
@@ -71,6 +72,7 @@ var (
 // For empty or invalid `Authorization` header, it sends "400 - Bad Request".
 //
 // See: https://jwt.io/introduction
+// See `JWTConfig.TokenLookup`
 func JWT(key []byte) echo.MiddlewareFunc {
 	c := DefaultJWTConfig
 	c.SigningKey = key
@@ -99,10 +101,14 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 
 	// Initialize
 	parts := strings.Split(config.TokenLookup, ":")
-	extractor := jwtFromHeader(parts[1])
+	var extractor jwtExtractor
 	switch parts[0] {
 	case "query":
 		extractor = jwtFromQuery(parts[1])
+	case "cookie":
+		extractor = jwtFromCookie(parts[1])
+	default:
+		extractor = jwtFromHeader(parts[1])
 	}
 
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -161,3 +167,14 @@ func jwtFromQuery(param string) jwtExtractor {
 		return token, err
 	}
 }
+
+// jwtFromCookie returns a `jwtExtractor` that extracts token from named cookie.
+func jwtFromCookie(name string) jwtExtractor {
+	return func(c echo.Context) (string, error) {
+		cookie, err := c.Cookie(name)
+		if err != nil {
+			return "", errJWTEmptyToken
+		}
+		return cookie.Value(), nil
+	}
+}
diff --git a/middleware/jwt_test.go b/middleware/jwt_test.go
index 181b507f..8a5c5398 100644
--- a/middleware/jwt_test.go
+++ b/middleware/jwt_test.go
@@ -29,6 +29,7 @@ func TestJWT(t *testing.T) {
 		config     JWTConfig
 		reqURL     string // "/" if empty
 		hdrAuth    string
+		hdrCookie  string // test.Request doesn't provide SetCookie(); use name=val
 		info       string
 	}{
 		{expPanic: true, info: "No signing key provided"},
@@ -117,6 +118,42 @@ func TestJWT(t *testing.T) {
 			expErrCode: http.StatusFound,
 			info:       "Empty query with redirect",
 		},
+		{
+			config: JWTConfig{
+				SigningKey:  validKey,
+				TokenLookup: "cookie:jwt",
+			},
+			hdrCookie: "jwt=" + token,
+			info:      "Valid cookie method",
+		},
+		{
+			config: JWTConfig{
+				SigningKey:  validKey,
+				TokenLookup: "cookie:jwt",
+			},
+			expErrCode: http.StatusUnauthorized,
+			hdrCookie:  "jwt=invalid",
+			info:       "Invalid token with cookie method",
+		},
+		{
+			config: JWTConfig{
+				SigningKey:  validKey,
+				TokenLookup: "cookie:jwt",
+			},
+			expErrCode: http.StatusBadRequest,
+			hdrCookie:  "",
+			info:       "Empty cookie",
+		},
+		{
+			config: JWTConfig{
+				SigningKey:       validKey,
+				TokenLookup:      "cookie:jwt",
+				HandleEmptyToken: redirect,
+			},
+			expErrCode: http.StatusFound,
+			hdrCookie:  "",
+			info:       "Empty cookie with redirect",
+		},
 	} {
 
 		if tc.reqURL == "" {
@@ -126,6 +163,7 @@ func TestJWT(t *testing.T) {
 		req := test.NewRequest(echo.GET, tc.reqURL, nil)
 		res := test.NewResponseRecorder()
 		req.Header().Set(echo.HeaderAuthorization, tc.hdrAuth)
+		req.Header().Set(echo.HeaderCookie, tc.hdrCookie)
 		c := e.NewContext(req, res)
 
 		if tc.expPanic {