Get token from query String.

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

auth_jwt.go

@@ -281,7 +281,7 @@ func (mw *GinJWTMiddleware) TokenGenerator(userID string) string {
 }
 
 func (mw *GinJWTMiddleware) jwtFromHeader(c *gin.Context, key string) (string, error) {
-	authHeader := c.Request.Header.Get("Authorization")
+	authHeader := c.Request.Header.Get(key)
 
 	if authHeader == "" {
 		return "", errors.New("auth header empty")
@@ -295,6 +295,16 @@ func (mw *GinJWTMiddleware) jwtFromHeader(c *gin.Context, key string) (string, e
 	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) parseToken(c *gin.Context) (*jwt.Token, error) {
 	var token string
 	var err error
@@ -303,8 +313,8 @@ func (mw *GinJWTMiddleware) parseToken(c *gin.Context) (*jwt.Token, error) {
 	switch parts[0] {
 	case "header":
 		token, err = mw.jwtFromHeader(c, parts[1])
-		// case "query":
-		// 	token, err = jwtFromQuery(parts[1])
+	case "query":
+		token, err = mw.jwtFromQuery(c, parts[1])
 		// case "cookie":
 		// 	token, err = jwtFromCookie(parts[1])
 	}

auth_jwt_test.go

@@ -624,3 +624,44 @@ 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)
+		})
+}