feat: add jwt middleware

ahttp/context.go

@@ -18,6 +18,18 @@ type Context struct {
 	response http.ResponseWriter
 	query    url.Values
 	handler  HandlerFunc
+	store    map[string]interface{}
+}
+
+func (c *Context) Get(key string) interface{} {
+	return c.store[key]
+}
+
+func (c *Context) Set(key string, val interface{}) {
+	if c.store == nil {
+		c.store = make(map[string]interface{})
+	}
+	c.store[key] = val
 }
 
 func (c *Context) Request() *http.Request {
@@ -88,6 +100,7 @@ func (c *Context) Reset(r *http.Request, w http.ResponseWriter) {
 	c.response = w
 	c.query = nil
 	c.handler = nil
+	c.store = nil
 }
 
 func (c *Context) blob(code int, contentType string, b []byte) error {

ahttp/middleware/jwt.go

@@ -0,0 +1,57 @@
+package middleware
+
+import (
+	"net/http"
+
+	"github.com/golang-jwt/jwt"
+	"github.com/zjregee/anet/ahttp"
+)
+
+var ErrJWTMissing = ahttp.NewHTTPError(http.StatusBadRequest, "missing or malformed jwt")
+var ErrJWTInvalid = ahttp.NewHTTPError(http.StatusUnauthorized, "invalid or expired jwt")
+
+type JWTConfig struct {
+	SigningKey    interface{}
+	SigningMethod string
+	ContextKey    string
+	Claims        jwt.Claims
+	TokenLookup   string
+	AuthScheme    string
+}
+
+var DefaultJWTConfig = JWTConfig{
+	SigningMethod: "HS256",
+	ContextKey:    "user",
+	Claims:        jwt.MapClaims{},
+	TokenLookup:   "header:Authorization",
+	AuthScheme:    "Bearer",
+}
+
+func JWT(key interface{}) ahttp.MiddlewareFunc {
+	if key == nil {
+		panic("shouldn't failed here")
+	}
+	config := DefaultJWTConfig
+	config.SigningKey = key
+	return func(next ahttp.HandlerFunc) ahttp.HandlerFunc {
+		return func(c *ahttp.Context) error {
+			extractor := valuesFromHeader("Authorization", config.AuthScheme)
+			auths, err := extractor(c)
+			if err != nil {
+				return ErrJWTMissing.SetInternal(err)
+			}
+			auth := auths[0]
+			token, err := jwt.Parse(auth, func(token *jwt.Token) (interface{}, error) {
+				return config.SigningKey, nil
+			})
+			if err != nil {
+				return ErrJWTInvalid.SetInternal(err)
+			}
+			if !token.Valid {
+				return ErrJWTInvalid
+			}
+			c.Set(config.ContextKey, token)
+			return next(c)
+		}
+	}
+}

ahttp/middleware/middleware.go

@@ -0,0 +1,36 @@
+package middleware
+
+import (
+	"errors"
+	"net/textproto"
+	"strings"
+
+	"github.com/zjregee/anet/ahttp"
+)
+
+type valuesExtractor func(c *ahttp.Context) ([]string, error)
+
+func valuesFromHeader(header string, valuePrefix string) valuesExtractor {
+	prefixLen := len(valuePrefix)
+	header = textproto.CanonicalMIMEHeaderKey(header)
+	return func(c *ahttp.Context) ([]string, error) {
+		values := c.Request().Header.Values(header)
+		if len(values) == 0 {
+			return nil, errors.New("header not found")
+		}
+		result := make([]string, 0)
+		for _, value := range values {
+			if prefixLen == 0 {
+				result = append(result, value)
+				continue
+			}
+			if len(value) > prefixLen && strings.EqualFold(value[:prefixLen], valuePrefix) {
+				result = append(result, value[prefixLen:])
+			}
+		}
+		if len(result) == 0 {
+			return nil, errors.New("value not found")
+		}
+		return result, nil
+	}
+}

go.mod

@@ -5,6 +5,7 @@ go 1.22.0
 require (
 	github.com/cloudwego/netpoll v0.6.2
 	github.com/ebitengine/purego v0.7.1
+	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/google/uuid v1.6.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/stretchr/testify v1.7.0

go.sum

@@ -7,6 +7,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/ebitengine/purego v0.7.1 h1:6/55d26lG3o9VCZX8lping+bZcmShseiqlh2bnUDiPA=
 github.com/ebitengine/purego v0.7.1/go.mod h1:ah1In8AOtksoNK6yk5z1HTJeUkC1Ez4Wk2idgGslMwQ=
+github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keLg81eXfW3O+oY=
+github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=