System Design: SSO and Session Management — Authentication at Scale

# Login: server creates a session in Redis
def login(username, password):
    user = db.find_user(username)
    if not verify_password(password, user.password_hash):
        raise AuthError("invalid credentials")

    session_id = generate_random_id()        # e.g. "a3f9c..." (128-bit random)
    session_data = {
        "userId":    user.id,
        "createdAt": now(),
        "expiresAt": now() + timedelta(days=30),
        "ip":        request.remote_addr,
        "userAgent": request.headers["User-Agent"],
    }
    redis.setex(
        "session:" + session_id,
        86400 * 30,           # TTL: 30 days in seconds
        json_encode(session_data)
    )
    return session_id       # stored in browser cookie

# Every request: server validates the session
def authenticate_request(request):
    session_id = request.cookies.get("session_id")
    session = redis.get("session:" + session_id)
    if not session or session["expiresAt"] < now():
        raise AuthError("session expired or invalid")
    return session["userId"]

# Login: server issues a signed JWT
def login_jwt(username, password):
    user = db.find_user(username)
    if not verify_password(password, user.password_hash):
        raise AuthError("invalid credentials")

    payload = {
        "userId":       user.id,
        "email":        user.email,
        "tokenVersion": user.token_version,   # for revocation (section 2)
        "iat":          now_unix(),
        "exp":          now_unix() + 900,      # expires in 15 minutes
    }
    return jwt.encode(payload, SECRET_KEY, algorithm="HS256")

# Every request: server verifies locally — NO Redis lookup
def authenticate_request_jwt(request):
    token = request.headers.get("Authorization").split()[1]
    try:
        claims = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
    except jwt.ExpiredSignatureError:
        raise AuthError("token expired")
    except jwt.InvalidTokenError:
        raise AuthError("invalid token")

    # Optionally verify tokenVersion against DB (section 2)
    return claims["userId"]

# Issuing tokens at login
def login_with_refresh(username, password):
    user = db.authenticate(username, password)

    access_token = jwt.encode({
        "userId": user.id,
        "exp":    now_unix() + 900,     # 15 minutes
    }, SECRET_KEY)

    refresh_token = generate_secure_random(64)
    db.store_refresh_token({
        "token":     sha256(refresh_token),  # store hash, not plaintext
        "userId":    user.id,
        "expiresAt": now() + timedelta(days=30),
        "deviceId":  request.get_device_id(),
    })

    return {"access_token": access_token, "refresh_token": refresh_token}

# Client calls this when access_token expires (HTTP 401)
def refresh_access_token(refresh_token):
    token_hash = sha256(refresh_token)
    record = db.find_refresh_token(token_hash)

    if not record or record["expiresAt"] < now():
        raise AuthError("refresh token invalid or expired")

    # Rotate: old token out, new token in (prevents replay)
    db.delete_refresh_token(token_hash)
    return login_with_refresh.issue_new_pair(record["userId"])

-- Revoke all sessions for a user
UPDATE users
SET    token_version = token_version + 1
WHERE  id = 'user_123';

-- Application check (pseudo-code in SQL style)
-- jwt.tokenVersion must equal users.token_version
SELECT token_version
FROM   users
WHERE  id = jwt_claim_user_id
  AND  token_version = jwt_claim_token_version;

# Auth Server: issue SSO token after successful authentication
def issue_sso_token(user_id, service, return_to):
    token_id = generate_random_id()
    token_data = {
        "userId":    user_id,
        "service":   service,
        "return_to": return_to,
        "createdAt": now_unix(),
        "expiresAt": now_unix() + 60,   # 60-second window
        "used":      False,
    }
    redis.setex("sso_token:" + token_id, 120, json_encode(token_data))
    return token_id

# Service Provider: validate SSO token (server-to-server)
def validate_sso_token(token_id):
    key = "sso_token:" + token_id
    data = redis.get(key)

    if not data:
        raise AuthError("token not found or expired")
    if data["expiresAt"] < now_unix():
        raise AuthError("token expired")
    if data["used"]:
        raise AuthError("token already used — replay attack?")

    # Mark as used atomically to prevent replay
    redis.hset(key, "used", True)

    return data["userId"]

// PKCE: generating the code_verifier and code_challenge
async function generatePKCE() {
    // 1. Generate a cryptographically random verifier
    const array = new Uint8Array(64);
    crypto.getRandomValues(array);
    const verifier = base64URLEncode(array);

    // 2. Hash it: challenge = BASE64URL(SHA256(verifier))
    const data = new TextEncoder().encode(verifier);
    const hashBuffer = await crypto.subtle.digest("SHA-256", data);
    const challenge = base64URLEncode(new Uint8Array(hashBuffer));

    return { verifier, challenge };
}

function base64URLEncode(buffer) {
    return btoa(String.fromCharCode(...buffer))
        .replace(/\+/g, "-")
        .replace(/\//g, "_")
        .replace(/=/g, "");
}

class TieredSessionStore:

    def get(self, session_id):
        # 1. Check hot tier (Redis) first
        session = self.redis_hot.get("sess:" + session_id)
        if session:
            self.redis_hot.expire("sess:" + session_id, 604800)  # refresh TTL
            return decode(session)

        # 2. Check warm tier
        session = self.redis_warm.get("sess:" + session_id)
        if session:
            self._promote_to_hot(session_id, session)
            return decode(session)

        # 3. Check cold tier (Cassandra)
        row = self.cassandra.execute(
            "SELECT * FROM sessions WHERE session_id = ?",
            [session_id]
        ).one()
        if row:
            self._promote_to_hot(session_id, encode(row))
            return row

        return None  # session not found anywhere

    def _promote_to_hot(self, session_id, data):
        self.redis_hot.setex("sess:" + session_id, 604800, data)
        # Optionally delete from warm/cold to avoid duplication

import hmac, hashlib, struct, time, base64

def generate_totp(secret_base32, digits=6, period=30):
    # 1. Decode the shared secret (stored in user DB, displayed as QR code)
    secret = base64.b32decode(secret_base32.upper())

    # 2. Compute time counter: 30-second windows since Unix epoch
    counter = int(time.time()) // period

    # 3. HMAC-SHA1 of the 8-byte big-endian counter
    msg = struct.pack(">Q", counter)
    h = hmac.new(secret, msg, hashlib.sha1).digest()

    # 4. Dynamic truncation: take 4 bytes at offset indicated by last nibble
    offset = h[-1] & 0x0F
    code = struct.unpack(">I", h[offset:offset+4])[0] & 0x7FFFFFFF

    # 5. Modulo to get N-digit code
    return str(code % (10 ** digits)).zfill(digits)

def verify_totp(secret, provided_code, window=1):
    # Accept current window and ±1 (clock skew tolerance)
    for drift in range(-window, window + 1):
        expected = generate_totp(secret, period=30)
        if hmac.compare_digest(expected, provided_code):
            return True
    return False

def login_step1(username, password):
    user = db.authenticate(username, password)
    if not user.mfa_enabled:
        return issue_full_session(user)  # no MFA, done

    # Issue a short-lived challenge token (not a full session!)
    challenge = {
        "userId":    user.id,
        "mfaNeeded": True,
        "exp":       now_unix() + 60,  # 60-second window to enter MFA code
    }
    challenge_token = jwt.encode(challenge, MFA_KEY, algorithm="HS256")
    return {"mfa_required": True, "challenge_token": challenge_token}

def login_step2(challenge_token, totp_code):
    claims = jwt.decode(challenge_token, MFA_KEY, algorithms=["HS256"])
    user = db.find_user(claims["userId"])

    if not verify_totp(user.mfa_secret, totp_code):
        raise AuthError("invalid TOTP code")

    return issue_full_session(user)  # MFA passed, issue real session

def use_refresh_token(token):
    record = db.find_refresh_token(sha256(token))

    if not record:
        # Token not found — either expired or already used.
        # Check if this token was recently rotated (possible replay attack)
        rotated = db.find_rotated_token(sha256(token))
        if rotated:
            # Replay detected: revoke the entire token family
            db.revoke_token_family(rotated.family_id)
            raise SecurityAlert("refresh token reuse detected")
        raise AuthError("token invalid")

    # Valid: rotate (issue new, invalidate old)
    new_token = generate_secure_random(64)
    db.rotate_refresh_token(
        old_token_hash=sha256(token),
        new_token_hash=sha256(new_token),
        family_id=record.family_id
    )
    new_access = issue_access_token(record.user_id)
    return {"access_token": new_access, "refresh_token": new_token}

# Secure cookie settings
# Set-Cookie: refresh_token=XYZ; HttpOnly; Secure; SameSite=Strict; Path=/auth/refresh
# Path=/auth/refresh means the cookie is ONLY sent to the refresh endpoint