ai-gateway-go/internal/identity/session.go

package identity

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
)

var ErrInvalidSession = errors.New("invalid or expired session")

type Principal struct {
	Kind        Kind   `json:"kind"`
	SubjectID   string `json:"subject_id"`
	Login       string `json:"login"`
	DisplayName string `json:"display_name"`
	Role        string `json:"role,omitempty"`
	Purpose     string `json:"purpose"`
	IssuedAt    int64  `json:"issued_at"`
	// AuthVersion 是签发会话时账号的凭据版本;凭据变更(改密/2FA 变更)会
	// 递增该版本,旧版本会话在 Authenticate 时被拒绝,被盗会话无法在
	// 凭据轮换后继续存活。
	AuthVersion int64 `json:"auth_version,omitempty"`
}

// authVersionTTL 必须严格大于会话 TTL(上限 7 天,由 config 校验保证):
// 版本键过期时所有旧会话已自然过期,凭据变更后旧会话不会复活。
const authVersionTTL = 14 * 24 * time.Hour

func authVersionKey(kind Kind, subjectID string) string {
	return "gateway:auth-version:" + string(kind) + ":" + subjectID
}

// AuthVersion 返回账号当前凭据版本;从未变更过则为 0。
func (s *SessionStore) AuthVersion(ctx context.Context, kind Kind, subjectID string) int64 {
	if s.client == nil {
		return 0
	}
	value, err := s.client.Get(ctx, authVersionKey(kind, subjectID)).Int64()
	if err != nil {
		return 0
	}
	return value
}

// BumpAuthVersion 使账号的全部既有会话失效(改密、2FA 启用/停用等凭据变更后调用)。
// Redis 不可用时静默失败:会话仍按 TTL 自然过期,凭据变更的即时失效降级为延迟生效。
func (s *SessionStore) BumpAuthVersion(ctx context.Context, kind Kind, subjectID string) {
	if s.client == nil {
		return
	}
	key := authVersionKey(kind, subjectID)
	_ = s.client.Incr(ctx, key).Err()
	_ = s.client.Expire(ctx, key, authVersionTTL).Err()
}

type SessionStore struct {
	client *redis.Client
	ttl    time.Duration
}

func NewSessionStore(client *redis.Client, ttl time.Duration) *SessionStore {
	return &SessionStore{client: client, ttl: ttl}
}

func (s *SessionStore) Create(ctx context.Context, principal Principal) (string, error) {
	principal.Purpose = "session"
	return s.create(ctx, principal, s.ttl)
}

func (s *SessionStore) CreatePending(ctx context.Context, principal Principal, ttl time.Duration) (string, error) {
	principal.Purpose = "totp_pending"
	return s.create(ctx, principal, ttl)
}

func (s *SessionStore) create(ctx context.Context, principal Principal, ttl time.Duration) (string, error) {
	if s.client == nil {
		return "", ErrUnavailable
	}
	random := make([]byte, 32)
	if _, err := rand.Read(random); err != nil {
		return "", err
	}
	token := base64.RawURLEncoding.EncodeToString(random)
	principal.IssuedAt = time.Now().Unix()
	principal.AuthVersion = s.AuthVersion(ctx, principal.Kind, principal.SubjectID)
	payload, err := json.Marshal(principal)
	if err != nil {
		return "", err
	}
	if err := s.client.Set(ctx, sessionKey(token), payload, ttl).Err(); err != nil {
		return "", fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return token, nil
}

func (s *SessionStore) Authenticate(ctx context.Context, authorization string, expected Kind) (Principal, error) {
	if s.client == nil {
		return Principal{}, ErrUnavailable
	}
	token, ok := bearerToken(authorization)
	if !ok {
		return Principal{}, ErrInvalidSession
	}
	payload, err := s.client.Get(ctx, sessionKey(token)).Bytes()
	if errors.Is(err, redis.Nil) {
		return Principal{}, ErrInvalidSession
	}
	if err != nil {
		return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	var principal Principal
	if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "session" {
		return Principal{}, ErrInvalidSession
	}
	// 凭据版本不匹配:改密/2FA 变更后旧会话一律失效。
	if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
		return Principal{}, ErrInvalidSession
	}
	return principal, nil
}

// AuthenticatePending 只读取(不消费)挑战令牌,供 CompleteTOTPLogin 在
// 验证前解析 principal;验证码输错时令牌保留可重试。
func (s *SessionStore) AuthenticatePending(ctx context.Context, token string, expected Kind) (Principal, error) {
	if s.client == nil {
		return Principal{}, ErrUnavailable
	}
	payload, err := s.client.Get(ctx, sessionKey(strings.TrimSpace(token))).Bytes()
	if errors.Is(err, redis.Nil) {
		return Principal{}, ErrInvalidSession
	}
	if err != nil {
		return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	var principal Principal
	if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
		return Principal{}, ErrInvalidSession
	}
	if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
		return Principal{}, ErrInvalidSession
	}
	return principal, nil
}

// ConsumePending 原子消费挑战令牌(GetDel):验证通过后调用,并发请求用同一
// 令牌时只有一个能成功,防止一次 2FA 挑战铸出两个会话。
func (s *SessionStore) ConsumePending(ctx context.Context, token string, expected Kind) (Principal, error) {
	if s.client == nil {
		return Principal{}, ErrUnavailable
	}
	payload, err := s.client.GetDel(ctx, sessionKey(strings.TrimSpace(token))).Bytes()
	if errors.Is(err, redis.Nil) {
		return Principal{}, ErrInvalidSession
	}
	if err != nil {
		return Principal{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	var principal Principal
	if err := json.Unmarshal(payload, &principal); err != nil || principal.Kind != expected || principal.SubjectID == "" || principal.Purpose != "totp_pending" {
		return Principal{}, ErrInvalidSession
	}
	if principal.AuthVersion != s.AuthVersion(ctx, principal.Kind, principal.SubjectID) {
		return Principal{}, ErrInvalidSession
	}
	return principal, nil
}

func (s *SessionStore) DeleteToken(ctx context.Context, token string) error {
	if s.client == nil {
		return ErrUnavailable
	}
	if err := s.client.Del(ctx, sessionKey(strings.TrimSpace(token))).Err(); err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return nil
}

func (s *SessionStore) Delete(ctx context.Context, authorization string) error {
	if s.client == nil {
		return ErrUnavailable
	}
	token, ok := bearerToken(authorization)
	if !ok {
		return nil
	}
	if err := s.client.Del(ctx, sessionKey(token)).Err(); err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return nil
}

func (s *SessionStore) StoreOneTime(ctx context.Context, namespace string, value any, ttl time.Duration) (string, error) {
	if s.client == nil {
		return "", ErrUnavailable
	}
	random := make([]byte, 32)
	if _, err := rand.Read(random); err != nil {
		return "", err
	}
	token := base64.RawURLEncoding.EncodeToString(random)
	payload, err := json.Marshal(value)
	if err != nil {
		return "", err
	}
	if err := s.client.Set(ctx, oneTimeKey(namespace, token), payload, ttl).Err(); err != nil {
		return "", fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return token, nil
}

func (s *SessionStore) ConsumeOneTime(ctx context.Context, namespace, token string, target any) error {
	if s.client == nil {
		return ErrUnavailable
	}
	payload, err := s.client.GetDel(ctx, oneTimeKey(namespace, strings.TrimSpace(token))).Bytes()
	if errors.Is(err, redis.Nil) {
		return ErrInvalidSession
	}
	if err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	if err := json.Unmarshal(payload, target); err != nil {
		return ErrInvalidSession
	}
	return nil
}

// ClaimIdentifier atomically claims a caller-provided replay identifier for
// the TTL. It is used for signed protocol message IDs, not bearer secrets.
func (s *SessionStore) ClaimIdentifier(ctx context.Context, namespace, identifier string, ttl time.Duration) (bool, error) {
	if s.client == nil {
		return false, ErrUnavailable
	}
	identifier = strings.TrimSpace(identifier)
	if identifier == "" || len(identifier) > 512 {
		return false, ErrInvalidSession
	}
	claimed, err := s.client.SetNX(ctx, oneTimeKey(namespace, identifier), "1", ttl).Result()
	if err != nil {
		return false, fmt.Errorf("%w: %v", ErrUnavailable, err)
	}
	return claimed, nil
}

func bearerToken(authorization string) (string, bool) {
	authorization = strings.TrimSpace(authorization)
	if len(authorization) <= len("Bearer ") || !strings.EqualFold(authorization[:len("Bearer ")], "Bearer ") {
		return "", false
	}
	token := strings.TrimSpace(authorization[len("Bearer "):])
	return token, token != ""
}

func sessionKey(token string) string {
	digest := sha256.Sum256([]byte(token))
	return "gateway:session:v1:" + hex.EncodeToString(digest[:])
}

func oneTimeKey(namespace, token string) string {
	digest := sha256.Sum256([]byte(token))
	return "gateway:one-time:" + namespace + ":" + hex.EncodeToString(digest[:])
}