feat: 实现 OIDC 服务端会话与请求刷新
使用 AES-256-GCM 保存认证中心令牌，并以随机 Cookie 哈希关联 PostgreSQL 会话。加入闲置与绝对期限、Refresh Token 轮换以及多实例刷新锁。
This commit is contained in:
@@ -0,0 +1,326 @@
|
||||
package oidcsession
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrSessionInvalid = errors.New("OIDC session is invalid")
|
||||
ErrSessionExpired = errors.New("OIDC session has expired")
|
||||
ErrSessionRefreshUnavailable = errors.New("OIDC session refresh is unavailable")
|
||||
ErrSessionStoreUnavailable = errors.New("OIDC session store is unavailable")
|
||||
ErrGatewayUserDisabled = errors.New("gateway user is disabled")
|
||||
)
|
||||
|
||||
type Repository interface {
|
||||
CreateOIDCSession(context.Context, store.CreateOIDCSessionInput) (store.OIDCSession, error)
|
||||
FindOIDCSessionByHash(context.Context, []byte) (store.OIDCSession, error)
|
||||
TouchOIDCSession(context.Context, string, time.Time, time.Time) error
|
||||
AcquireOIDCSessionRefreshLock(context.Context, string, int64, string, time.Time, time.Time) (bool, error)
|
||||
CompleteOIDCSessionRefresh(context.Context, string, int64, string, []byte, time.Time) error
|
||||
DeleteOIDCSessionByHash(context.Context, []byte) error
|
||||
DeleteOIDCSessionByID(context.Context, string) error
|
||||
CleanupExpiredOIDCSessions(context.Context, time.Time) (int64, error)
|
||||
}
|
||||
|
||||
type TokenVerifier interface {
|
||||
Verify(context.Context, string) (*auth.User, error)
|
||||
}
|
||||
|
||||
type PublicClient interface {
|
||||
Refresh(context.Context, string) (auth.OIDCTokenResponse, error)
|
||||
RevokeRefreshToken(context.Context, string) error
|
||||
EndSessionURL(context.Context, string) (string, error)
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
IdleTTL time.Duration
|
||||
AbsoluteTTL time.Duration
|
||||
RefreshBefore time.Duration
|
||||
RefreshLease time.Duration
|
||||
RefreshWait time.Duration
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
repository Repository
|
||||
cipher *Cipher
|
||||
verifier TokenVerifier
|
||||
client PublicClient
|
||||
config Config
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
func NewService(repository Repository, cipher *Cipher, verifier TokenVerifier, client PublicClient, config Config) (*Service, error) {
|
||||
if repository == nil || cipher == nil || verifier == nil || client == nil {
|
||||
return nil, errors.New("OIDC session dependencies are required")
|
||||
}
|
||||
if config.IdleTTL <= 0 || config.AbsoluteTTL <= config.IdleTTL || config.RefreshBefore <= 0 {
|
||||
return nil, errors.New("OIDC session TTL configuration is invalid")
|
||||
}
|
||||
if config.RefreshLease <= 0 {
|
||||
config.RefreshLease = 5 * time.Second
|
||||
}
|
||||
if config.RefreshWait <= 0 {
|
||||
config.RefreshWait = 2 * time.Second
|
||||
}
|
||||
return &Service{repository: repository, cipher: cipher, verifier: verifier, client: client, config: config, now: time.Now}, nil
|
||||
}
|
||||
|
||||
func (s *Service) Create(ctx context.Context, bundle TokenBundle, localUser *auth.User) (string, error) {
|
||||
if localUser == nil || localUser.GatewayUserID == "" || localUser.GatewayTenantID == "" || bundle.RefreshToken == "" {
|
||||
return "", ErrSessionInvalid
|
||||
}
|
||||
verified, err := s.verifier.Verify(ctx, bundle.AccessToken)
|
||||
if err != nil || verified == nil || verified.Source != "oidc" || verified.ID == "" || verified.ID != localUser.ID {
|
||||
return "", ErrSessionInvalid
|
||||
}
|
||||
now := s.now()
|
||||
if !verified.TokenExpiresAt.After(now) {
|
||||
return "", ErrSessionExpired
|
||||
}
|
||||
raw, hash, err := newSessionToken()
|
||||
if err != nil {
|
||||
return "", ErrSessionStoreUnavailable
|
||||
}
|
||||
aadSessionID := hex.EncodeToString(hash)
|
||||
ciphertext, err := s.cipher.EncryptBundle(bundle, aadSessionID, localUser.GatewayUserID)
|
||||
if err != nil {
|
||||
return "", ErrSessionStoreUnavailable
|
||||
}
|
||||
_, err = s.repository.CreateOIDCSession(ctx, store.CreateOIDCSessionInput{
|
||||
SessionTokenHash: hash, GatewayUserID: localUser.GatewayUserID, GatewayTenantID: localUser.GatewayTenantID,
|
||||
TokenCiphertext: ciphertext, AccessTokenExpiresAt: verified.TokenExpiresAt,
|
||||
LastSeenAt: now, IdleExpiresAt: now.Add(s.config.IdleTTL), AbsoluteExpiresAt: now.Add(s.config.AbsoluteTTL),
|
||||
})
|
||||
if err != nil {
|
||||
return "", ErrSessionStoreUnavailable
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
func (s *Service) Resolve(ctx context.Context, raw string) (*auth.User, error) {
|
||||
hash, err := sessionTokenHash(raw)
|
||||
if err != nil {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
|
||||
if errors.Is(err, store.ErrOIDCSessionNotFound) {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
if err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
return s.resolveRecord(ctx, hash, record)
|
||||
}
|
||||
|
||||
func (s *Service) resolveRecord(ctx context.Context, hash []byte, record store.OIDCSession) (*auth.User, error) {
|
||||
now := s.now()
|
||||
if record.UserDeleted || record.UserStatus != "active" {
|
||||
return nil, ErrGatewayUserDisabled
|
||||
}
|
||||
if !record.IdleExpiresAt.After(now) || !record.AbsoluteExpiresAt.After(now) {
|
||||
_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
|
||||
return nil, ErrSessionExpired
|
||||
}
|
||||
bundle, err := s.cipher.DecryptBundle(record.TokenCiphertext, hex.EncodeToString(hash), record.GatewayUserID)
|
||||
if err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
if !record.AccessTokenExpiresAt.After(now.Add(s.config.RefreshBefore)) {
|
||||
return s.refresh(ctx, hash, record, bundle)
|
||||
}
|
||||
user, err := s.verifySessionUser(ctx, record, bundle.AccessToken)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.touch(ctx, record, now); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
func (s *Service) refresh(ctx context.Context, hash []byte, record store.OIDCSession, bundle TokenBundle) (*auth.User, error) {
|
||||
now := s.now()
|
||||
lockID, err := newUUID()
|
||||
if err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
acquired, err := s.repository.AcquireOIDCSessionRefreshLock(ctx, record.ID, record.RefreshVersion, lockID, now.Add(s.config.RefreshLease), now)
|
||||
if err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
if !acquired {
|
||||
return s.waitForRefresh(ctx, hash, record, bundle)
|
||||
}
|
||||
|
||||
refreshed, err := s.client.Refresh(ctx, bundle.RefreshToken)
|
||||
if errors.Is(err, auth.ErrOIDCInvalidGrant) {
|
||||
_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
|
||||
return nil, ErrSessionExpired
|
||||
}
|
||||
if err != nil {
|
||||
// Keep the lease until it expires: an indeterminate refresh response must not
|
||||
// cause the same rotating refresh token to be replayed immediately.
|
||||
if record.AccessTokenExpiresAt.After(now) {
|
||||
user, verifyErr := s.verifySessionUser(ctx, record, bundle.AccessToken)
|
||||
if verifyErr == nil {
|
||||
if touchErr := s.touch(ctx, record, now); touchErr != nil {
|
||||
return nil, touchErr
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
}
|
||||
return nil, ErrSessionRefreshUnavailable
|
||||
}
|
||||
user, err := s.verifySessionUser(ctx, record, refreshed.AccessToken)
|
||||
if err != nil {
|
||||
_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
|
||||
return nil, ErrSessionExpired
|
||||
}
|
||||
if refreshed.RefreshToken == "" {
|
||||
// OAuth token responses may omit refresh_token when the issuer keeps the
|
||||
// existing token valid. Persist the old value unless a rotated one exists.
|
||||
refreshed.RefreshToken = bundle.RefreshToken
|
||||
}
|
||||
if refreshed.IDToken == "" {
|
||||
refreshed.IDToken = bundle.IDToken
|
||||
}
|
||||
newBundle := TokenBundle{AccessToken: refreshed.AccessToken, RefreshToken: refreshed.RefreshToken, IDToken: refreshed.IDToken}
|
||||
ciphertext, err := s.cipher.EncryptBundle(newBundle, hex.EncodeToString(hash), record.GatewayUserID)
|
||||
if err != nil {
|
||||
// The remote refresh succeeded, so the previous rotating refresh token may
|
||||
// already be invalid. Destroy the stale local session instead of replaying it.
|
||||
_ = s.repository.DeleteOIDCSessionByID(ctx, record.ID)
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
if err := s.repository.CompleteOIDCSessionRefresh(ctx, record.ID, record.RefreshVersion, lockID, ciphertext, user.TokenExpiresAt); err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
record.AccessTokenExpiresAt = user.TokenExpiresAt
|
||||
if err := s.touch(ctx, record, now); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
func (s *Service) waitForRefresh(ctx context.Context, hash []byte, original store.OIDCSession, oldBundle TokenBundle) (*auth.User, error) {
|
||||
deadline := time.Now().Add(s.config.RefreshWait)
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ErrSessionRefreshUnavailable
|
||||
case <-time.After(25 * time.Millisecond):
|
||||
}
|
||||
record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
|
||||
if errors.Is(err, store.ErrOIDCSessionNotFound) {
|
||||
return nil, ErrSessionExpired
|
||||
}
|
||||
if err != nil {
|
||||
return nil, ErrSessionStoreUnavailable
|
||||
}
|
||||
if record.RefreshVersion > original.RefreshVersion {
|
||||
return s.resolveRecord(ctx, hash, record)
|
||||
}
|
||||
}
|
||||
now := s.now()
|
||||
if original.AccessTokenExpiresAt.After(now) {
|
||||
user, err := s.verifySessionUser(ctx, original, oldBundle.AccessToken)
|
||||
if err == nil {
|
||||
if touchErr := s.touch(ctx, original, now); touchErr != nil {
|
||||
return nil, touchErr
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
}
|
||||
return nil, ErrSessionRefreshUnavailable
|
||||
}
|
||||
|
||||
func (s *Service) verifySessionUser(ctx context.Context, record store.OIDCSession, accessToken string) (*auth.User, error) {
|
||||
user, err := s.verifier.Verify(ctx, accessToken)
|
||||
if err != nil || user == nil || user.Source != "oidc" || user.ID != record.ExternalUserID {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
return user, nil
|
||||
}
|
||||
|
||||
func (s *Service) touch(ctx context.Context, record store.OIDCSession, now time.Time) error {
|
||||
if err := s.repository.TouchOIDCSession(ctx, record.ID, now, now.Add(s.config.IdleTTL)); err != nil {
|
||||
if errors.Is(err, store.ErrOIDCSessionNotFound) {
|
||||
return ErrSessionExpired
|
||||
}
|
||||
return ErrSessionStoreUnavailable
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) Delete(ctx context.Context, raw string) (TokenBundle, error) {
|
||||
hash, err := sessionTokenHash(raw)
|
||||
if err != nil {
|
||||
return TokenBundle{}, nil
|
||||
}
|
||||
record, err := s.repository.FindOIDCSessionByHash(ctx, hash)
|
||||
if errors.Is(err, store.ErrOIDCSessionNotFound) {
|
||||
return TokenBundle{}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return TokenBundle{}, ErrSessionStoreUnavailable
|
||||
}
|
||||
bundle, decryptErr := s.cipher.DecryptBundle(record.TokenCiphertext, hex.EncodeToString(hash), record.GatewayUserID)
|
||||
if err := s.repository.DeleteOIDCSessionByHash(ctx, hash); err != nil {
|
||||
return TokenBundle{}, ErrSessionStoreUnavailable
|
||||
}
|
||||
if decryptErr != nil {
|
||||
// The session row is already gone. Treat logout as successful even though
|
||||
// the unusable refresh token could not be revoked at the issuer.
|
||||
return TokenBundle{}, nil
|
||||
}
|
||||
return bundle, nil
|
||||
}
|
||||
|
||||
func (s *Service) Cleanup(ctx context.Context) (int64, error) {
|
||||
count, err := s.repository.CleanupExpiredOIDCSessions(ctx, s.now())
|
||||
if err != nil {
|
||||
return 0, ErrSessionStoreUnavailable
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func newSessionToken() (string, []byte, error) {
|
||||
raw := make([]byte, 32)
|
||||
if _, err := rand.Read(raw); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
encoded := base64.RawURLEncoding.EncodeToString(raw)
|
||||
hash := sha256.Sum256([]byte(encoded))
|
||||
return encoded, hash[:], nil
|
||||
}
|
||||
|
||||
func sessionTokenHash(raw string) ([]byte, error) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(raw)
|
||||
if err != nil || len(decoded) != 32 {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
hash := sha256.Sum256([]byte(raw))
|
||||
return hash[:], nil
|
||||
}
|
||||
|
||||
func newUUID() (string, error) {
|
||||
value := make([]byte, 16)
|
||||
if _, err := rand.Read(value); err != nil {
|
||||
return "", err
|
||||
}
|
||||
value[6] = value[6]&0x0f | 0x40
|
||||
value[8] = value[8]&0x3f | 0x80
|
||||
return hex.EncodeToString(value[0:4]) + "-" + hex.EncodeToString(value[4:6]) + "-" + hex.EncodeToString(value[6:8]) + "-" + hex.EncodeToString(value[8:10]) + "-" + hex.EncodeToString(value[10:16]), nil
|
||||
}
|
||||
Reference in New Issue
Block a user