forked from razonyang/clevergo-examples
-
Notifications
You must be signed in to change notification settings - Fork 2
/
jwt.go
101 lines (85 loc) · 2.22 KB
/
jwt.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
// JSON WEB TOKEN MIDDLEWARE
package jwtmiddleware
import (
"github.com/clevergo/jwt"
"github.com/headwindfly/clevergo"
"github.com/valyala/fasthttp"
)
const (
jwtKey = "_jwt"
)
var (
errHandler = func(ctx *clevergo.Context) {
ctx.SetStatusCode(fasthttp.StatusUnauthorized)
}
successHandler = func(ctx *clevergo.Context, token *jwt.Token) {}
)
// JWTMiddleware JSON WEB TOKEN Middleware.
type JWTMiddleware struct {
jwt *jwt.JWT
key string
errorHandler clevergo.HandlerFunc
successHandler func(ctx *clevergo.Context, token *jwt.Token)
}
// NewJWTMiddleware returns JWT Middleware instance.
func NewJWTMiddleware(jwt *jwt.JWT) JWTMiddleware {
return JWTMiddleware{
jwt: jwt,
key: jwtKey,
errorHandler: errHandler,
successHandler: successHandler,
}
}
// Key returns key.
func (m JWTMiddleware) Key() string {
return m.key
}
// SetKey for setting key.
func (m *JWTMiddleware) SetKey(key string) {
m.key = key
}
// SetErrorHandler for setting error handler.
func (m *JWTMiddleware) SetErrorHandler(handler clevergo.HandlerFunc) {
m.errorHandler = handler
}
// SetSuccessHandler for setting success handler.
func (m *JWTMiddleware) SetSuccessHandler(handler func(ctx *clevergo.Context, token *jwt.Token)) {
m.successHandler = handler
}
// Handle implemented Middleware Interface.
func (m JWTMiddleware) Handle(next clevergo.Handler) clevergo.Handler {
return clevergo.HandlerFunc(func(ctx *clevergo.Context) {
// Try to get JWT raw token from request form.
rawToken := ctx.FormValue(m.key)
if len(rawToken) == 0 {
// Try to get JWT raw token from Header.
if ah := ctx.Request.Header.Peek("Authorization"); len(ah) > 0 {
// Should be a bearer token
if len(ah) > 6 {
rawToken = ah[7:]
}
}
}
// Check raw token is valid.
if len(rawToken) == 0 {
m.errorHandler(ctx)
return
}
// Get JWT by raw token.
var err error
token, err := m.jwt.NewTokenByRaw(string(rawToken))
if err != nil {
m.errorHandler(ctx)
return
}
// Validate Token.
if err := token.Validate(); err != nil {
m.errorHandler(ctx)
return
}
// Invoke the success handler
m.successHandler(ctx, token)
// Validate successfully.
next.Handle(ctx)
})
}