stats-gateway/internal/auth/jwt.go

package auth

import (
	"bufio"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"
)

// Claims holds the JWT payload fields used by this service.
type Claims struct {
	Subject  string `json:"sub"`
	Scope    string `json:"scope"`
	IssuedAt int64  `json:"iat"`
	Expires  int64  `json:"exp"`
	JTI      string `json:"jti"`
}

// Verifier validates HMAC-HS256 tokens against a secret and optional revocation list.
type Verifier struct {
	secret  []byte
	revoked map[string]bool
}

// NewVerifier creates a Verifier. If revoked is nil an empty map is used.
func NewVerifier(secret []byte, revoked map[string]bool) *Verifier {
	if revoked == nil {
		revoked = map[string]bool{}
	}
	return &Verifier{secret: secret, revoked: revoked}
}

// Verify checks signature, expiry, scope=="tunnel", and revocation list.
// Returns parsed Claims on success.
func (v *Verifier) Verify(token string) (*Claims, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, errors.New("malformed token")
	}

	signingInput := parts[0] + "." + parts[1]
	expectedSig, err := base64URLDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("decode sig: %w", err)
	}
	mac := hmac.New(sha256.New, v.secret)
	mac.Write([]byte(signingInput))
	if !hmac.Equal(mac.Sum(nil), expectedSig) {
		return nil, errors.New("bad signature")
	}

	headerJSON, err := base64URLDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("decode header: %w", err)
	}
	var header struct {
		Alg string `json:"alg"`
		Typ string `json:"typ"`
	}
	if err := json.Unmarshal(headerJSON, &header); err != nil {
		return nil, fmt.Errorf("header: %w", err)
	}
	if header.Alg != "HS256" {
		return nil, fmt.Errorf("alg %q not allowed", header.Alg)
	}

	payload, err := base64URLDecode(parts[1])
	if err != nil {
		return nil, fmt.Errorf("decode payload: %w", err)
	}
	var c Claims
	if err := json.Unmarshal(payload, &c); err != nil {
		return nil, fmt.Errorf("payload: %w", err)
	}

	now := time.Now().Unix()
	if c.Expires < now {
		return nil, errors.New("token expired")
	}
	if c.Scope != "tunnel" {
		return nil, fmt.Errorf("scope %q not allowed", c.Scope)
	}
	if v.revoked[c.JTI] {
		return nil, errors.New("token revoked")
	}
	return &c, nil
}

// Issue creates a new HS256 token. ttl <= 0 produces an already-expired token (useful for tests).
func Issue(secret []byte, subject, scope string, ttl time.Duration) (string, error) {
	jtiBytes := make([]byte, 8)
	if _, err := rand.Read(jtiBytes); err != nil {
		return "", err
	}
	return IssueWithJTI(secret, subject, scope, ttl, fmt.Sprintf("%x", jtiBytes))
}

// IssueWithJTI creates a token with a caller-supplied jti (useful for revocation tests).
func IssueWithJTI(secret []byte, subject, scope string, ttl time.Duration, jti string) (string, error) {
	now := time.Now().Unix()
	c := Claims{
		Subject:  subject,
		Scope:    scope,
		IssuedAt: now,
		Expires:  now + int64(ttl.Seconds()),
		JTI:      jti,
	}
	header := `{"alg":"HS256","typ":"JWT"}`
	payload, err := json.Marshal(c)
	if err != nil {
		return "", err
	}

	h := base64URLEncode([]byte(header))
	p := base64URLEncode(payload)
	signingInput := h + "." + p
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(signingInput))
	sig := base64URLEncode(mac.Sum(nil))
	return signingInput + "." + sig, nil
}

// LoadRevokedFile reads one jti per line from path. Lines starting with '#'
// or empty are ignored. Missing file → empty map (no error).
func LoadRevokedFile(path string) (map[string]bool, error) {
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return map[string]bool{}, nil
		}
		return nil, err
	}
	defer f.Close()
	out := map[string]bool{}
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		out[line] = true
	}
	return out, sc.Err()
}

func base64URLEncode(b []byte) string          { return base64.RawURLEncoding.EncodeToString(b) }
func base64URLDecode(s string) ([]byte, error) { return base64.RawURLEncoding.DecodeString(s) }