341 lines
12 KiB
Go
341 lines
12 KiB
Go
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")
|
|
ErrSecurityStateUnavailable = errors.New("OIDC security event state is unavailable")
|
|
)
|
|
|
|
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
|
|
}
|
|
if errors.Is(verifyErr, ErrSecurityStateUnavailable) {
|
|
return nil, verifyErr
|
|
}
|
|
}
|
|
return nil, ErrSessionRefreshUnavailable
|
|
}
|
|
user, err := s.verifySessionUser(ctx, record, refreshed.AccessToken)
|
|
if err != nil {
|
|
if errors.Is(err, ErrSecurityStateUnavailable) {
|
|
return nil, err
|
|
}
|
|
_ = 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
|
|
}
|
|
if errors.Is(err, ErrSecurityStateUnavailable) {
|
|
return nil, err
|
|
}
|
|
}
|
|
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)
|
|
var requestError *auth.RequestAuthError
|
|
if errors.As(err, &requestError) && requestError.Status == 503 {
|
|
return nil, ErrSecurityStateUnavailable
|
|
}
|
|
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
|
|
}
|