Cyrene_v2/internal/auth/jwt.go

package auth

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// Claims is the verified JWT payload.
type Claims struct {
	Sub      string `json:"sub"`
	Role     Role   `json:"role"`
	Exp      int64  `json:"exp"`
	IssuedAt int64  `json:"iat"`
}

// jwtHeader is the JOSE header for the HS256 tokens this package emits.
type jwtHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
}

// JWTService signs and verifies HS256 JWTs using a shared secret. The secret is
// protected by a RWMutex so RotateSecret can swap it while Sign/Verify read it
// concurrently (the gateway rotates the secret live from the admin panel).
type JWTService struct {
	mu     sync.RWMutex
	secret []byte
	log    *slog.Logger
}

// New validates cfg and returns a JWTService. It panics when JWTSecret is
// shorter than 32 characters so a misconfiguration fails fast at startup.
func New(cfg config.AuthConfig) *JWTService {
	if err := Validate(cfg); err != nil {
		panic(err)
	}
	return &JWTService{
		secret: []byte(cfg.JWTSecret),
		log:    logger.WithSource(slog.Default(), source),
	}
}

// Validate checks that cfg satisfies the auth package's configuration
// requirements, returning ErrInvalidSecret when JWTSecret is too short.
func Validate(cfg config.AuthConfig) error {
	if len(cfg.JWTSecret) < minSecretLen {
		return ErrInvalidSecret
	}
	return nil
}

// Sign issues an HS256 token for sub with the given role and lifetime.
func (s *JWTService) Sign(sub string, role Role, ttl time.Duration) (string, error) {
	if !role.Valid() {
		return "", fmt.Errorf("%w: %q", ErrInvalidRole, role)
	}
	now := time.Now().Unix()
	claims := Claims{
		Sub:      sub,
		Role:     role,
		IssuedAt: now,
		Exp:      now + int64(ttl.Seconds()),
	}
	headerB64, err := marshalB64(jwtHeader{Alg: "HS256", Typ: "JWT"})
	if err != nil {
		return "", fmt.Errorf("auth: encode header: %w", err)
	}
	payloadB64, err := marshalB64(claims)
	if err != nil {
		return "", fmt.Errorf("auth: encode payload: %w", err)
	}
	signingInput := headerB64 + "." + payloadB64
	s.mu.RLock()
	sig := hmacSign(s.secret, []byte(signingInput))
	s.mu.RUnlock()
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(sig), nil
}

// RotateSecret replaces the signing secret with newSecret, invalidating every
// token issued under the previous secret. It returns ErrInvalidSecret when
// newSecret is shorter than the 32-character minimum. The new secret is never
// logged or echoed back to the caller.
func (s *JWTService) RotateSecret(newSecret string) error {
	if len(newSecret) < minSecretLen {
		return fmt.Errorf("%w: got %d chars", ErrInvalidSecret, len(newSecret))
	}
	s.mu.Lock()
	s.secret = []byte(newSecret)
	s.mu.Unlock()
	s.log.Info("jwt secret rotated; previously issued tokens are now invalid")
	return nil
}

// Verify validates the token signature, algorithm, role, and expiry, returning
// the parsed Claims on success.
func (s *JWTService) Verify(token string) (*Claims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, fmt.Errorf("%w: expected 3 segments", ErrMalformedToken)
	}
	headerB64, payloadB64, sigB64 := parts[0], parts[1], parts[2]
	signingInput := headerB64 + "." + payloadB64

	provided, err := base64.RawURLEncoding.DecodeString(sigB64)
	if err != nil {
		return nil, fmt.Errorf("%w: decode signature: %v", ErrMalformedToken, err)
	}
	s.mu.RLock()
	expected := hmacSign(s.secret, []byte(signingInput))
	s.mu.RUnlock()
	if !hmac.Equal(provided, expected) {
		return nil, ErrInvalidSignature
	}

	headerJSON, err := base64.RawURLEncoding.DecodeString(headerB64)
	if err != nil {
		return nil, fmt.Errorf("%w: decode header: %v", ErrMalformedToken, err)
	}
	var hdr jwtHeader
	if err := json.Unmarshal(headerJSON, &hdr); err != nil {
		return nil, fmt.Errorf("%w: decode header: %v", ErrMalformedToken, err)
	}
	if hdr.Alg != "HS256" {
		return nil, fmt.Errorf("%w: %q", ErrUnsupportedAlg, hdr.Alg)
	}

	payloadJSON, err := base64.RawURLEncoding.DecodeString(payloadB64)
	if err != nil {
		return nil, fmt.Errorf("%w: decode payload: %v", ErrMalformedToken, err)
	}
	var claims Claims
	if err := json.Unmarshal(payloadJSON, &claims); err != nil {
		return nil, fmt.Errorf("%w: decode payload: %v", ErrMalformedToken, err)
	}
	if !claims.Role.Valid() {
		return nil, fmt.Errorf("%w: %q", ErrInvalidRole, claims.Role)
	}
	if claims.Exp > 0 && time.Now().Unix() >= claims.Exp {
		return nil, ErrTokenExpired
	}
	return &claims, nil
}

// marshalB64 JSON-encodes v and base64url-encodes the result without padding.
func marshalB64(v any) (string, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

// hmacSign returns the HS256 HMAC of input under secret.
func hmacSign(secret, input []byte) []byte {
	h := hmac.New(sha256.New, secret)
	h.Write(input)
	return h.Sum(nil)
}