Add get from header function.

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

Author: appleboy

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

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,40 @@ func (mw *GinJWTMiddleware) TokenGenerator(userID string) string {
 	return tokenString
 }
 
-func (mw *GinJWTMiddleware) parseToken(c *gin.Context) (*jwt.Token, error) {
+func (mw *GinJWTMiddleware) jwtFromHeader(c *gin.Context, key string) (string, error) {
 	authHeader := c.Request.Header.Get("Authorization")
 
 	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) 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 = jwtFromQuery(parts[1])
+		// case "cookie":
+		// 	token, err = jwtFromCookie(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.",