优化对ws的支持

auth.go

@@ -39,6 +39,7 @@ func SetAuthToken(uniqueUser, privateKeyPath string, expire time.Duration) (toke
 	return tokenString, nil
 }
 
+// ParseFromRequest 兼容 http,ws
 func ParseFromRequest(req *http.Request, publicKeyPath string) (uniqueUser, signature string, err error) {
 	token, err := request.ParseFromRequest(req, request.AuthorizationHeaderExtractor, func(t *jwt.Token) (interface{}, error) {
 		publicKey, innErr := os.ReadFile(publicKeyPath)
@@ -52,51 +53,37 @@ func ParseFromRequest(req *http.Request, publicKeyPath string) (uniqueUser, sign
 		return "", "", err
 	}
 
-	wsp := req.Header.Get("Sec-Websocket-Protocol")
-
-	if !token.Valid && len(wsp) > 0 {
-		return "", "", jwt.ErrTokenSignatureInvalid
+	if token.Valid {
+		return parseToken(token)
 	}
 
-	claims, ok := token.Claims.(*jwt.RegisteredClaims)
-	if !ok {
-		return "", "", jwt.ErrTokenInvalidClaims
-	}
-
-	return claims.Issuer, token.Signature, nil
-}
-
-func ParseWithClaims(req *http.Request, publicKeyPath string) (uniqueUser, signature string, err error) {
-	wsp := req.Header.Get("Sec-Websocket-Protocol")
-	if len(wsp) > 0 {
-		token, er := jwt.ParseWithClaims(wsp, &jwt.RegisteredClaims{}, func(t *jwt.Token) (interface{}, error) {
+	// 兼容 ws
+	if wsp := req.Header.Get("Sec-Websocket-Protocol"); len(wsp) > 0 {
+		token, err = jwt.ParseWithClaims(wsp, &jwt.RegisteredClaims{}, func(t *jwt.Token) (interface{}, error) {
 			publicKey, innErr := os.ReadFile(publicKeyPath)
 			if innErr != nil {
 				return "", innErr
 			}
 
 			return jwt.ParseRSAPublicKeyFromPEM(publicKey)
 		})
-		if er != nil {
-			return "", "", er
+		if err != nil {
+			return "", "", err
 		}
 
-		if !token.Valid && len(wsp) > 0 {
-			return "", "", jwt.ErrTokenSignatureInvalid
-		}
-
-		claims, ok := token.Claims.(*jwt.RegisteredClaims)
-		if !ok {
-			return "", "", jwt.ErrTokenInvalidClaims
+		if token.Valid {
+			return parseToken(token)
 		}
+	}
 
-		//issuer, er := claims.GetIssuer()
-		//if er != nil {
-		//	return "", "", er
-		//}
+	return "", "", jwt.ErrTokenSignatureInvalid
+}
 
-		return claims.Issuer, token.Signature, nil
+func parseToken(token *jwt.Token) (uniqueUser, signature string, err error) {
+	claims, ok := token.Claims.(*jwt.RegisteredClaims)
+	if !ok {
+		return "", "", jwt.ErrTokenInvalidClaims
 	}
 
-	return "", "", nil
+	return claims.Issuer, token.Signature, nil
 }