Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package middlewares
18

19
import (
20
	"context"
21
	"net/http"
22
	"time"
23

24
	jwt "github.com/appleboy/gin-jwt/v2"
25
	"github.com/gin-gonic/gin"
26

27
	"d7y.io/dragonfly/v2/manager/config"
28
	"d7y.io/dragonfly/v2/manager/models"
29
	"d7y.io/dragonfly/v2/manager/service"
30
	"d7y.io/dragonfly/v2/manager/types"
31
)
32

33
const (
34
	// defaultIdentityKey is default of the identity key.
35
	defaultIdentityKey = "id"
36
)
37

38
func Jwt(cfg config.JWTConfig, service service.Service) (*jwt.GinJWTMiddleware, error) {
39
	authMiddleware, err := jwt.New(&jwt.GinJWTMiddleware{
40
		Realm:       cfg.Realm,
41
		Key:         []byte(cfg.Key),
42
		Timeout:     cfg.Timeout,
43
		MaxRefresh:  cfg.MaxRefresh,
44
		IdentityKey: defaultIdentityKey,
45

46
		IdentityHandler: func(c *gin.Context) any {
47
			claims := jwt.ExtractClaims(c)
48

49
			id, ok := claims[defaultIdentityKey]
50
			if !ok {
51
				c.JSON(http.StatusUnauthorized, gin.H{
52
					"message": "Unavailable token: require user id",
53
				})
54
				c.Abort()
55
				return nil
56
			}
57

58
			c.Set("id", id)
59
			return id
60
		},
61

62
		Authenticator: func(c *gin.Context) (any, error) {
63
			// Oauth2 signin
64
			if rawUser, ok := c.Get("user"); ok {
65
				user, ok := rawUser.(*models.User)
66
				if !ok {
67
					return "", jwt.ErrFailedAuthentication
68
				}
69
				return user, nil
70
			}
71

72
			// Normal signin
73
			var json types.SignInRequest
74
			if err := c.ShouldBindJSON(&json); err != nil {
75
				return "", jwt.ErrMissingLoginValues
76
			}
77

78
			user, err := service.SignIn(context.TODO(), json)
79
			if err != nil {
80
				return "", jwt.ErrFailedAuthentication
81
			}
82

83
			return user, nil
84
		},
85

86
		PayloadFunc: func(data any) jwt.MapClaims {
87
			if user, ok := data.(*models.User); ok {
88
				return jwt.MapClaims{
89
					defaultIdentityKey: user.ID,
90
				}
91
			}
92

93
			return jwt.MapClaims{}
94
		},
95

96
		Unauthorized: func(c *gin.Context, code int, message string) {
97
			c.JSON(code, gin.H{
98
				"message": http.StatusText(code),
99
			})
100
		},
101

102
		LoginResponse: func(c *gin.Context, code int, token string, expire time.Time) {
103
			// Oauth2 signin
104
			if _, ok := c.Get("user"); ok {
105
				c.Redirect(http.StatusFound, "/")
106
				return
107
			}
108

109
			// Normal signin
110
			c.JSON(code, gin.H{
111
				"token":  token,
112
				"expire": expire.Format(time.RFC3339),
113
			})
114
		},
115

116
		LogoutResponse: func(c *gin.Context, code int) {
117
			c.Status(code)
118
		},
119

120
		RefreshResponse: func(c *gin.Context, code int, token string, expire time.Time) {
121
			c.JSON(code, gin.H{
122
				"token":  token,
123
				"expire": expire.Format(time.RFC3339),
124
			})
125
		},
126

127
		TokenLookup:    "cookie: jwt, header: Authorization, query: token",
128
		TokenHeadName:  "Bearer",
129
		TimeFunc:       time.Now,
130
		SendCookie:     true,
131
		CookieHTTPOnly: false,
132
	})
133
	if err != nil {
134
		return nil, err
135
	}
136

137
	return authMiddleware, nil
138
}
139