crawler/oauth_session.go

package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

const (
	sessionCookieName = "1440_session"
	sessionTTL        = 24 * time.Hour
)

// OAuthSession stores the OAuth session state for a user
type OAuthSession struct {
	ID        string    `json:"id"`
	DID       string    `json:"did"`
	Handle    string    `json:"handle"`
	CreatedAt time.Time `json:"created_at"`
	ExpiresAt time.Time `json:"expires_at"`

	// OAuth tokens (stored server-side only)
	AccessToken  string    `json:"access_token"`
	RefreshToken string    `json:"refresh_token"`
	TokenExpiry  time.Time `json:"token_expiry"`

	// DPoP state
	DpopPrivateJWK      string `json:"dpop_private_jwk"`
	DpopAuthserverNonce string `json:"dpop_authserver_nonce"`
	DpopPdsNonce        string `json:"dpop_pds_nonce"`

	// Auth server info
	PdsURL        string `json:"pds_url"`
	AuthserverIss string `json:"authserver_iss"`
}

// PendingAuth stores state during the OAuth flow (before callback)
type PendingAuth struct {
	State          string    `json:"state"`
	PkceVerifier   string    `json:"pkce_verifier"`
	DpopPrivateJWK string    `json:"dpop_private_jwk"`
	DpopNonce      string    `json:"dpop_nonce"`
	DID            string    `json:"did"`
	PdsURL         string    `json:"pds_url"`
	AuthserverIss  string    `json:"authserver_iss"`
	CreatedAt      time.Time `json:"created_at"`
}

// SessionStore manages sessions in memory
type SessionStore struct {
	sessions    map[string]*OAuthSession
	pending     map[string]*PendingAuth // keyed by state
	mu          sync.RWMutex
	cleanupOnce sync.Once
}

// NewSessionStore creates a new session store
func NewSessionStore() *SessionStore {
	s := &SessionStore{
		sessions: make(map[string]*OAuthSession),
		pending:  make(map[string]*PendingAuth),
	}
	s.startCleanup()
	return s
}

// startCleanup starts a background goroutine to clean up expired sessions
func (s *SessionStore) startCleanup() {
	s.cleanupOnce.Do(func() {
		go func() {
			ticker := time.NewTicker(5 * time.Minute)
			defer ticker.Stop()
			for range ticker.C {
				s.cleanup()
			}
		}()
	})
}

// cleanup removes expired sessions and pending auths
func (s *SessionStore) cleanup() {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()

	// Clean up expired sessions
	for id, sess := range s.sessions {
		if now.After(sess.ExpiresAt) {
			delete(s.sessions, id)
		}
	}

	// Clean up old pending auths (10 minute timeout)
	for state, pending := range s.pending {
		if now.Sub(pending.CreatedAt) > 10*time.Minute {
			delete(s.pending, state)
		}
	}
}

// CreateSession creates a new session and returns it
func (s *SessionStore) CreateSession(did, handle string) (*OAuthSession, error) {
	id, err := generateRandomID()
	if err != nil {
		return nil, err
	}

	now := time.Now()
	session := &OAuthSession{
		ID:        id,
		DID:       did,
		Handle:    handle,
		CreatedAt: now,
		ExpiresAt: now.Add(sessionTTL),
	}

	s.mu.Lock()
	s.sessions[id] = session
	s.mu.Unlock()

	return session, nil
}

// GetSession retrieves a session by ID
func (s *SessionStore) GetSession(id string) *OAuthSession {
	s.mu.RLock()
	defer s.mu.RUnlock()

	session, ok := s.sessions[id]
	if !ok || time.Now().After(session.ExpiresAt) {
		return nil
	}
	return session
}

// UpdateSession updates a session
func (s *SessionStore) UpdateSession(session *OAuthSession) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[session.ID] = session
}

// DeleteSession removes a session
func (s *SessionStore) DeleteSession(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.sessions, id)
}

// SavePending saves pending OAuth state
func (s *SessionStore) SavePending(state string, pending *PendingAuth) {
	s.mu.Lock()
	defer s.mu.Unlock()
	pending.CreatedAt = time.Now()
	s.pending[state] = pending
}

// GetPending retrieves and removes pending OAuth state
func (s *SessionStore) GetPending(state string) *PendingAuth {
	s.mu.Lock()
	defer s.mu.Unlock()

	pending, ok := s.pending[state]
	if ok {
		delete(s.pending, state)
	}
	return pending
}

// generateRandomID generates a random session ID
func generateRandomID() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

// encryptSessionID encrypts a session ID using AES-256-GCM
func encryptSessionID(sessionID string, key []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	ciphertext := gcm.Seal(nonce, nonce, []byte(sessionID), nil)
	return base64.URLEncoding.EncodeToString(ciphertext), nil
}

// decryptSessionID decrypts a session ID using AES-256-GCM
func decryptSessionID(encrypted string, key []byte) (string, error) {
	ciphertext, err := base64.URLEncoding.DecodeString(encrypted)
	if err != nil {
		return "", err
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	if len(ciphertext) < gcm.NonceSize() {
		return "", fmt.Errorf("ciphertext too short")
	}

	nonce, ciphertext := ciphertext[:gcm.NonceSize()], ciphertext[gcm.NonceSize():]
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

// SetSessionCookie sets an encrypted session cookie
func (m *OAuthManager) SetSessionCookie(w http.ResponseWriter, sessionID string) error {
	encrypted, err := encryptSessionID(sessionID, m.cookieSecret)
	if err != nil {
		return err
	}

	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookieName,
		Value:    encrypted,
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   int(sessionTTL.Seconds()),
	})

	return nil
}

// GetSessionFromCookie retrieves the session from the request cookie
func (m *OAuthManager) GetSessionFromCookie(r *http.Request) *OAuthSession {
	cookie, err := r.Cookie(sessionCookieName)
	if err != nil {
		return nil
	}

	sessionID, err := decryptSessionID(cookie.Value, m.cookieSecret)
	if err != nil {
		return nil
	}

	return m.sessions.GetSession(sessionID)
}

// ClearSessionCookie removes the session cookie
func (m *OAuthManager) ClearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     sessionCookieName,
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		Secure:   true,
		SameSite: http.SameSiteLaxMode,
		MaxAge:   -1,
	})
}

// SessionInfo is the public session info returned to the client
type SessionInfo struct {
	DID       string    `json:"did"`
	Handle    string    `json:"handle"`
	ExpiresAt time.Time `json:"expires_at"`
}

// MarshalJSON converts SessionInfo to JSON
func (s *SessionInfo) MarshalJSON() ([]byte, error) {
	type Alias SessionInfo
	return json.Marshal((*Alias)(s))
}