hspguard/internal/auth/jwt.go

package auth

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"os"

	"gitea.local/admin/hspguard/internal/types"
	"gitea.local/admin/hspguard/internal/web"
	"github.com/golang-jwt/jwt/v5"
)

func WriteJWKS(w http.ResponseWriter, r *http.Request) {
	pubKey, err := parseBase64PublicKey("JWT_PUBLIC_KEY")
	if err != nil {
		web.Error(w, "failed to parse public key", http.StatusInternalServerError)
	}

	n := base64.RawURLEncoding.EncodeToString(pubKey.N.Bytes())
	e := base64.RawURLEncoding.EncodeToString([]byte{1, 0, 1}) // 65537 = 0x010001

	jwks := map[string]interface{}{
		"keys": []map[string]string{
			{
				"kty": "RSA",
				"kid": "my-rsa-key-1",
				"use": "sig",
				"alg": "RS256",
				"n":   n,
				"e":   e,
			},
		},
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(jwks)
}

func OpenIdConfiguration(w http.ResponseWriter, r *http.Request) {
	type Response struct {
		TokenEndpoint         string `json:"token_endpoint"`
		AuthorizationEndpoint string `json:"authorization_endpoint"`
		JwksURI               string `json:"jwks_uri"`
		Issuer                string `json:"issuer"`
		EndSessionEndpoint    string `json:"end_session_endpoint"`
	}

	w.Header().Set("Content-Type", "application/json")

	encoder := json.NewEncoder(w)
	if err := encoder.Encode(Response{
		TokenEndpoint:         "https://cb5f-2a00-10-5b00-c801-e955-5c68-63d0-b777.ngrok-free.app/api/v1/oauth/token",
		AuthorizationEndpoint: "https://cb5f-2a00-10-5b00-c801-e955-5c68-63d0-b777.ngrok-free.app/authorize",
		JwksURI:               "https://cb5f-2a00-10-5b00-c801-e955-5c68-63d0-b777.ngrok-free.app/.well-known/jwks.json",
		Issuer:                "https://cb5f-2a00-10-5b00-c801-e955-5c68-63d0-b777.ngrok-free.app",
		EndSessionEndpoint:    "https://cb5f-2a00-10-5b00-c801-e955-5c68-63d0-b777.ngrok-free.app/api/v1/oauth/logout",
	}); err != nil {
		web.Error(w, "failed to encode response", http.StatusInternalServerError)
	}
}

func parseBase64PrivateKey(envVar string) (*rsa.PrivateKey, error) {
	b64 := os.Getenv(envVar)
	if b64 == "" {
		return nil, fmt.Errorf("env var %s is empty", envVar)
	}

	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode base64 key: %v", err)
	}

	key, err := x509.ParsePKCS8PrivateKey(decoded)
	return key.(*rsa.PrivateKey), err
}

func parseBase64PublicKey(envVar string) (*rsa.PublicKey, error) {
	b64 := os.Getenv(envVar)
	if b64 == "" {
		return nil, fmt.Errorf("env var %s is empty", envVar)
	}

	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode base64 key: %v", err)
	}

	pubInterface, err := x509.ParsePKIXPublicKey(decoded)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %v", err)
	}

	pubKey, ok := pubInterface.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA public key")
	}

	return pubKey, nil
}

func SignJwtToken(claims jwt.Claims) (string, error) {
	privateKey, err := parseBase64PrivateKey("JWT_PRIVATE_KEY")
	if err != nil {
		return "", err
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)

	token.Header["kid"] = "my-rsa-key-1"

	s, err := token.SignedString(privateKey)
	if err != nil {
		return "", err
	}

	return s, nil
}

func VerifyToken(token string) (*jwt.Token, *types.UserClaims, error) {
	publicKey, err := parseBase64PublicKey("JWT_PUBLIC_KEY")
	if err != nil {
		return nil, nil, err
	}

	claims := &types.UserClaims{}
	parsed, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		return publicKey, nil
	})

	if err != nil {
		return nil, nil, fmt.Errorf("invalid token: %w", err)
	}

	if !parsed.Valid {
		return nil, nil, fmt.Errorf("token is not valid")
	}

	return parsed, claims, nil
}