p2p-llm / auth /jwt.go
arpinfidel's picture
temp
48511d8
package auth
import (
"crypto/rsa"
"fmt"
"net/http"
"strings"
"github.com/golang-jwt/jwt"
)
type JWTMiddleware struct {
publicKey *rsa.PublicKey
}
func NewJWTMiddleware(publicKey *rsa.PublicKey) *JWTMiddleware {
return &JWTMiddleware{publicKey: publicKey}
}
func (m *JWTMiddleware) Middleware(next http.HandlerFunc) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Skip auth endpoint
if r.URL.Path == "/auth" {
next(w, r)
return
}
authHeader := r.Header.Get("Authorization")
if authHeader == "" {
http.Error(w, "Authorization header required", http.StatusUnauthorized)
return
}
tokenString := strings.TrimPrefix(authHeader, "Bearer ")
token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
}
return m.publicKey, nil
})
if err != nil {
http.Error(w, "Invalid token", http.StatusUnauthorized)
return
}
if !token.Valid {
http.Error(w, "Invalid token", http.StatusUnauthorized)
return
}
next(w, r)
})
}