1
0
forked from wrenn/wrenn
Files
wrenn-releases/internal/auth/jwt.go
pptx704 9bf67aa7f7 Implement host registration, JWT refresh tokens, and multi-host scheduling
Replaces the hardcoded CP_HOST_AGENT_ADDR single-agent setup with a
DB-driven registration system supporting multiple host agents (BYOC).

Key changes:
- Host agents register via one-time token, receive a 7-day JWT + 60-day
  refresh token; heartbeat loop auto-refreshes on 401/403 and pauses all
  sandboxes if refresh fails
- HostClientPool: lazy Connect RPC client cache keyed by host ID, replacing
  the single static agent client throughout the API and service layers
- RoundRobinScheduler: picks an online host for each new sandbox via
  ListActiveHosts; extensible for future scheduling strategies
- HostMonitor (replaces Reconciler): passive heartbeat staleness check marks
  hosts unreachable and sandboxes missing after 90s; active reconciliation
  per online host restores missing-but-alive sandboxes and stops orphans
- Graceful host delete: returns 409 with affected sandbox list without
  ?force=true; force-delete destroys sandboxes then evicts pool client
- Snapshot delete broadcasts to all online hosts (templates have no host_id)
- sandbox.Manager.PauseAll: pauses all running VMs on CP connectivity loss
- New migration: host_refresh_tokens table with token rotation (issue-then-
  revoke ordering to prevent lockout on mid-rotation crash)
- New sandbox status 'missing' (reversible, unlike 'stopped') and host
  status 'unreachable'; both reflected in OpenAPI spec
- Fix: refresh token auth failure now returns 401 (was 400 via generic
  'invalid' substring match in serviceErrToHTTP)
2026-03-24 18:32:05 +06:00

108 lines
3.3 KiB
Go

package auth
import (
"fmt"
"time"
"github.com/golang-jwt/jwt/v5"
)
const jwtExpiry = 6 * time.Hour
const hostJWTExpiry = 7 * 24 * time.Hour // 7 days; host refreshes via refresh token
const HostRefreshTokenExpiry = 60 * 24 * time.Hour // 60 days; exported for service layer
// Claims are the JWT payload for user tokens.
type Claims struct {
Type string `json:"typ,omitempty"` // empty for user tokens; used to reject host tokens
TeamID string `json:"team_id"`
Role string `json:"role"` // owner, admin, or member within TeamID
Email string `json:"email"`
Name string `json:"name"`
jwt.RegisteredClaims
}
// SignJWT signs a new 6-hour JWT for the given user.
func SignJWT(secret []byte, userID, teamID, email, name, role string) (string, error) {
now := time.Now()
claims := Claims{
TeamID: teamID,
Role: role,
Email: email,
Name: name,
RegisteredClaims: jwt.RegisteredClaims{
Subject: userID,
IssuedAt: jwt.NewNumericDate(now),
ExpiresAt: jwt.NewNumericDate(now.Add(jwtExpiry)),
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
return token.SignedString(secret)
}
// VerifyJWT parses and validates a user JWT, returning the claims on success.
// Rejects host JWTs (which carry a "typ" claim) to prevent cross-token confusion.
func VerifyJWT(secret []byte, tokenStr string) (Claims, error) {
token, err := jwt.ParseWithClaims(tokenStr, &Claims{}, func(t *jwt.Token) (any, error) {
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
}
return secret, nil
})
if err != nil {
return Claims{}, fmt.Errorf("invalid token: %w", err)
}
c, ok := token.Claims.(*Claims)
if !ok || !token.Valid {
return Claims{}, fmt.Errorf("invalid token claims")
}
if c.Type == "host" {
return Claims{}, fmt.Errorf("invalid token: host token cannot be used as user token")
}
return *c, nil
}
// HostClaims are the JWT payload for host agent tokens.
type HostClaims struct {
Type string `json:"typ"` // always "host"
HostID string `json:"host_id"`
jwt.RegisteredClaims
}
// SignHostJWT signs a long-lived (1 year) JWT for a registered host agent.
func SignHostJWT(secret []byte, hostID string) (string, error) {
now := time.Now()
claims := HostClaims{
Type: "host",
HostID: hostID,
RegisteredClaims: jwt.RegisteredClaims{
Subject: hostID,
IssuedAt: jwt.NewNumericDate(now),
ExpiresAt: jwt.NewNumericDate(now.Add(hostJWTExpiry)),
},
}
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
return token.SignedString(secret)
}
// VerifyHostJWT parses and validates a host JWT, returning the claims on success.
// It rejects user JWTs by checking the "typ" claim.
func VerifyHostJWT(secret []byte, tokenStr string) (HostClaims, error) {
token, err := jwt.ParseWithClaims(tokenStr, &HostClaims{}, func(t *jwt.Token) (any, error) {
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
}
return secret, nil
})
if err != nil {
return HostClaims{}, fmt.Errorf("invalid token: %w", err)
}
c, ok := token.Claims.(*HostClaims)
if !ok || !token.Valid {
return HostClaims{}, fmt.Errorf("invalid token claims")
}
if c.Type != "host" {
return HostClaims{}, fmt.Errorf("invalid token type: expected host")
}
return *c, nil
}