Merge pull request #44 from appleboy/patch-3

Support fetch token from cookie, query or header

Author: appleboy

.drone.yml

@@ -1,11 +1,16 @@
 pipeline:
   build:
-    image: appleboy/golang-testing
+    image: golang:${GO_VERSION}
     environment:
       - GOPATH=/go
     commands:
-      - make install
-      - coverage all
+      - go get -d -t -v
+      - go test -v
 
 workspace:
   path: /go/src/github.com/appleboy/gin-jwt
+
+matrix:
+  GO_VERSION:
+    - 1.6.3
+    - 1.7.1

Makefile

@@ -9,7 +9,7 @@ update:
 	glide up
 
 test:
-	go test -v -cover .
+	go test -v -cover -coverprofile=.cover/coverage.txt
 
 html:
 	go tool cover -html=.cover/coverage.txt

README.md

@@ -93,6 +93,16 @@ func main() {
         "message": message,
       })
     },
+    // TokenLookup is a string in the form of "<source>:<name>" that is used
+    // to extract token from the request.
+    // Optional. Default value "header:Authorization".
+    // Possible values:
+    // - "header:<name>"
+    // - "query:<name>"
+    // - "cookie:<name>"
+    TokenLookup: "header:Authorization",
+    // TokenLookup: "query:token",
+    // TokenLookup: "cookie:token",
   }
 
   r.POST("/login", authMiddleware.LoginHandler)

auth_jwt.go

@@ -54,6 +54,15 @@ type GinJWTMiddleware struct {
 
 	// User can define own Unauthorized func.
 	Unauthorized func(*gin.Context, int, string)
+
+	// TokenLookup is a string in the form of "<source>:<name>" that is used
+	// to extract token from the request.
+	// Optional. Default value "header:Authorization".
+	// Possible values:
+	// - "header:<name>"
+	// - "query:<name>"
+	// - "cookie:<name>"
+	TokenLookup string
 }
 
 // Login form structure.
@@ -65,6 +74,10 @@ type Login struct {
 // MiddlewareInit initialize jwt configs.
 func (mw *GinJWTMiddleware) MiddlewareInit() error {
 
+	if mw.TokenLookup == "" {
+		mw.TokenLookup = "header:Authorization"
+	}
+
 	if mw.SigningAlgorithm == "" {
 		mw.SigningAlgorithm = "HS256"
 	}
@@ -267,19 +280,60 @@ func (mw *GinJWTMiddleware) TokenGenerator(userID string) string {
 	return tokenString
 }
 
-func (mw *GinJWTMiddleware) parseToken(c *gin.Context) (*jwt.Token, error) {
-	authHeader := c.Request.Header.Get("Authorization")
+func (mw *GinJWTMiddleware) jwtFromHeader(c *gin.Context, key string) (string, error) {
+	authHeader := c.Request.Header.Get(key)
 
 	if authHeader == "" {
-		return nil, errors.New("auth header empty")
+		return "", errors.New("auth header empty")
 	}
 
 	parts := strings.SplitN(authHeader, " ", 2)
 	if !(len(parts) == 2 && parts[0] == "Bearer") {
-		return nil, errors.New("invalid auth header")
+		return "", errors.New("invalid auth header")
+	}
+
+	return parts[1], nil
+}
+
+func (mw *GinJWTMiddleware) jwtFromQuery(c *gin.Context, key string) (string, error) {
+	token := c.Query(key)
+
+	if token == "" {
+		return "", errors.New("Query token empty")
+	}
+
+	return token, nil
+}
+
+func (mw *GinJWTMiddleware) jwtFromCookie(c *gin.Context, key string) (string, error) {
+	cookie, _ := c.Cookie(key)
+
+	if cookie == "" {
+		return "", errors.New("Cookie token empty")
+	}
+
+	return cookie, nil
+}
+
+func (mw *GinJWTMiddleware) parseToken(c *gin.Context) (*jwt.Token, error) {
+	var token string
+	var err error
+
+	parts := strings.Split(mw.TokenLookup, ":")
+	switch parts[0] {
+	case "header":
+		token, err = mw.jwtFromHeader(c, parts[1])
+	case "query":
+		token, err = mw.jwtFromQuery(c, parts[1])
+	case "cookie":
+		token, err = mw.jwtFromCookie(c, parts[1])
+	}
+
+	if err != nil {
+		return nil, err
 	}
 
-	return jwt.Parse(parts[1], func(token *jwt.Token) (interface{}, error) {
+	return jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
 		if jwt.GetSigningMethod(mw.SigningAlgorithm) != token.Method {
 			return nil, errors.New("invalid signing algorithm")
 		}

auth_jwt_test.go

@@ -108,6 +108,25 @@ func TestMissingTimeOut(t *testing.T) {
 	assert.Equal(t, time.Hour, authMiddleware.Timeout)
 }
 
+func TestMissingTokenLookup(t *testing.T) {
+
+	authMiddleware := &GinJWTMiddleware{
+		Realm: "test zone",
+		Key:   key,
+		Authenticator: func(userId string, password string, c *gin.Context) (string, bool) {
+			if userId == "admin" && password == "admin" {
+				return "", true
+			}
+
+			return "", false
+		},
+	}
+
+	authMiddleware.MiddlewareInit()
+
+	assert.Equal(t, "header:Authorization", authMiddleware.TokenLookup)
+}
+
 func helloHandler(c *gin.Context) {
 	c.JSON(200, gin.H{
 		"text": "Hello World.",
@@ -605,3 +624,85 @@ func TestTokenExpire(t *testing.T) {
 			assert.Equal(t, http.StatusUnauthorized, r.Code)
 		})
 }
+
+func TestTokenFromQueryString(t *testing.T) {
+	// the middleware to test
+	authMiddleware := &GinJWTMiddleware{
+		Realm:   "test zone",
+		Key:     key,
+		Timeout: time.Hour,
+		Authenticator: func(userId string, password string, c *gin.Context) (string, bool) {
+			if userId == "admin" && password == "admin" {
+				return userId, true
+			}
+			return userId, false
+		},
+		Unauthorized: func(c *gin.Context, code int, message string) {
+			c.String(code, message)
+		},
+		TokenLookup: "query:token",
+	}
+
+	handler := ginHandler(authMiddleware)
+
+	r := gofight.New()
+
+	userToken := authMiddleware.TokenGenerator("admin")
+
+	r.GET("/auth/refresh_token").
+		SetHeader(gofight.H{
+			"Authorization": "Bearer " + userToken,
+		}).
+		Run(handler, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+			assert.Equal(t, http.StatusUnauthorized, r.Code)
+		})
+
+	r.GET("/auth/refresh_token?token="+userToken).
+		SetHeader(gofight.H{
+			"Authorization": "Bearer " + userToken,
+		}).
+		Run(handler, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+			assert.Equal(t, http.StatusOK, r.Code)
+		})
+}
+
+func TestTokenFromCookieString(t *testing.T) {
+	// the middleware to test
+	authMiddleware := &GinJWTMiddleware{
+		Realm:   "test zone",
+		Key:     key,
+		Timeout: time.Hour,
+		Authenticator: func(userId string, password string, c *gin.Context) (string, bool) {
+			if userId == "admin" && password == "admin" {
+				return userId, true
+			}
+			return userId, false
+		},
+		Unauthorized: func(c *gin.Context, code int, message string) {
+			c.String(code, message)
+		},
+		TokenLookup: "cookie:token",
+	}
+
+	handler := ginHandler(authMiddleware)
+
+	r := gofight.New()
+
+	userToken := authMiddleware.TokenGenerator("admin")
+
+	r.GET("/auth/refresh_token").
+		SetHeader(gofight.H{
+			"Authorization": "Bearer " + userToken,
+		}).
+		Run(handler, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+			assert.Equal(t, http.StatusUnauthorized, r.Code)
+		})
+
+	r.GET("/auth/refresh_token").
+		SetCookie(gofight.H{
+			"token": userToken,
+		}).
+		Run(handler, func(r gofight.HTTPResponse, rq gofight.HTTPRequest) {
+			assert.Equal(t, http.StatusOK, r.Code)
+		})
+}

example/server.go

@@ -50,6 +50,16 @@ func main() {
 				"message": message,
 			})
 		},
+		// TokenLookup is a string in the form of "<source>:<name>" that is used
+		// to extract token from the request.
+		// Optional. Default value "header:Authorization".
+		// Possible values:
+		// - "header:<name>"
+		// - "query:<name>"
+		// - "cookie:<name>"
+		TokenLookup: "header:Authorization",
+		// TokenLookup: "query:token",
+		// TokenLookup: "cookie:token",
 	}
 
 	r.POST("/login", authMiddleware.LoginHandler)