OAuth 2.0 Security

// 1. Generate code verifier and challenge
import { randomBytes, createHash } from 'crypto'
 
interface PKCEChallenge {
  codeVerifier: string
  codeChallenge: string
  codeChallengeMethod: 'S256'
}
 
function generatePKCEChallenge(): PKCEChallenge {
  // Generate cryptographically secure random string (43-128 chars)
  const codeVerifier = randomBytes(32)
    .toString('base64url') // URL-safe base64
    .slice(0, 128) // Max 128 characters
 
  // Create SHA-256 hash of verifier
  const hash = createHash('sha256')
    .update(codeVerifier)
    .digest('base64url')
 
  return {
    codeVerifier,
    codeChallenge: hash,
    codeChallengeMethod: 'S256'
  }
}
 
// 2. Generate CSRF protection state
function generateState(): string {
  return randomBytes(32).toString('hex') // 64 hex characters
}

// 3. Complete OAuth flow manager
class OAuth2PKCEManager {
  private codeVerifier: string | null = null
  private state: string | null = null
 
  /**
   * Initiate OAuth flow
   * Returns authorization URL to open in browser
   */
  async initiateFlow(config: OAuthConfig): Promise<string> {
    // Generate PKCE challenge
    const pkce = generatePKCEChallenge()
    this.codeVerifier = pkce.codeVerifier
 
    // Generate CSRF protection state
    this.state = generateState()
 
    // Store in secure temporary storage
    await this.storeFlowState({
      codeVerifier: pkce.codeVerifier,
      state: this.state,
      timestamp: Date.now(),
      provider: config.provider
    })
 
    // Build authorization URL
    const params = new URLSearchParams({
      client_id: config.clientId,
      redirect_uri: config.redirectUri,
      response_type: 'code',
      scope: config.scopes.join(' '),
      state: this.state,
      code_challenge: pkce.codeChallenge,
      code_challenge_method: pkce.codeChallengeMethod,
      // Optional: force reauth, select account, etc.
      prompt: config.prompt || 'consent',
      access_type: 'offline' // Request refresh token
    })
 
    return `${config.authorizationEndpoint}?${params}`
  }
 
  /**
   * Handle OAuth callback
   * Validates state and exchanges code for tokens
   */
  async handleCallback(
    code: string,
    state: string,
    config: OAuthConfig
  ): Promise<TokenResponse> {
    // 1. Retrieve stored flow state
    const flowState = await this.getFlowState()
    if (!flowState) {
      throw new OAuthError('No OAuth flow in progress', 'INVALID_FLOW')
    }
 
    // 2. Validate state (CSRF protection)
    if (state !== flowState.state) {
      throw new OAuthError('Invalid OAuth state parameter', 'INVALID_STATE')
    }
 
    // 3. Check flow expiry (15 minutes)
    const now = Date.now()
    const flowAge = now - flowState.timestamp
    if (flowAge > 15 * 60 * 1000) {
      throw new OAuthError('OAuth flow expired', 'FLOW_EXPIRED')
    }
 
    // 4. Clear flow state
    await this.clearFlowState()
 
    // 5. Exchange authorization code for tokens
    try {
      const tokens = await this.exchangeCode(
        code,
        flowState.codeVerifier,
        config
      )
 
      // 6. Store tokens securely
      await this.storeTokensSecurely(config.provider, tokens)
 
      return tokens
    } catch (error) {
      throw new OAuthError(
        `Token exchange failed: ${error.message}`,
        'TOKEN_EXCHANGE_FAILED',
        error
      )
    }
  }
 
  /**
   * Exchange authorization code for tokens
   * Uses code verifier to prove authenticity
   */
  private async exchangeCode(
    code: string,
    codeVerifier: string,
    config: OAuthConfig
  ): Promise<TokenResponse> {
    const params = new URLSearchParams({
      client_id: config.clientId,
      code,
      code_verifier: codeVerifier,
      grant_type: 'authorization_code',
      redirect_uri: config.redirectUri
    })
 
    // Note: No client_secret needed with PKCE!
    const response = await fetch(config.tokenEndpoint, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
      },
      body: params.toString()
    })
 
    if (!response.ok) {
      const error = await response.json()
      throw new Error(error.error_description || error.error)
    }
 
    const tokens: TokenResponse = await response.json()
 
    return {
      accessToken: tokens.access_token,
      refreshToken: tokens.refresh_token,
      expiresIn: tokens.expires_in,
      tokenType: tokens.token_type,
      scope: tokens.scope,
      expiresAt: Date.now() + (tokens.expires_in * 1000)
    }
  }
 
  /**
   * Store flow state in secure temporary storage
   */
  private async storeFlowState(state: FlowState): Promise<void> {
    // Use encrypted temporary storage (cleared after flow completes)
    await invoke('store_temp_data', {
      key: 'oauth_flow_state',
      data: JSON.stringify(state),
      ttl: 900 // 15 minutes
    })
  }
 
  private async getFlowState(): Promise<FlowState | null> {
    const data = await invoke<string | null>('get_temp_data', {
      key: 'oauth_flow_state'
    })
    return data ? JSON.parse(data) : null
  }
 
  private async clearFlowState(): Promise<void> {
    await invoke('clear_temp_data', { key: 'oauth_flow_state' })
  }
}
 
// TypeScript interfaces
interface OAuthConfig {
  provider: string
  clientId: string
  redirectUri: string
  authorizationEndpoint: string
  tokenEndpoint: string
  scopes: string[]
  prompt?: 'none' | 'consent' | 'select_account'
}
 
interface TokenResponse {
  accessToken: string
  refreshToken?: string
  expiresIn: number
  tokenType: string
  scope?: string
  expiresAt: number
}
 
interface FlowState {
  codeVerifier: string
  state: string
  timestamp: number
  provider: string
}
 
class OAuthError extends Error {
  constructor(
    message: string,
    public code: string,
    public cause?: Error
  ) {
    super(message)
    this.name = 'OAuthError'
  }
}

class HybridRedirectHandler {
  /**
   * Try deep link first, fallback to localhost
   */
  async handleRedirect(authUrl: string): Promise<TokenResponse> {
    // 1. Start localhost server as fallback
    const localPort = await this.startLocalServer()
    const localRedirectUri = `http://127.0.0.1:${localPort}/oauth/callback`
 
    // 2. Try deep link redirect
    const deepLinkUri = 'lokus://oauth/callback'
 
    // Check if deep link is registered
    const hasDeepLink = await this.checkDeepLinkRegistration()
 
    // 3. Update auth URL with appropriate redirect
    const redirectUri = hasDeepLink ? deepLinkUri : localRedirectUri
    const finalAuthUrl = authUrl.replace(
      /redirect_uri=[^&]+/,
      `redirect_uri=${encodeURIComponent(redirectUri)}`
    )
 
    // 4. Open in browser
    await open(finalAuthUrl)
 
    // 5. Wait for callback (either deep link or localhost)
    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.stopLocalServer()
        reject(new Error('OAuth timeout'))
      }, 5 * 60 * 1000) // 5 minutes
 
      // Handle deep link callback
      if (hasDeepLink) {
        this.onDeepLink((code, state) => {
          clearTimeout(timeout)
          this.stopLocalServer()
          resolve(this.completeFlow(code, state))
        })
      }
 
      // Local server handles localhost callback
      // (automatically resolves promise)
    })
  }
 
  private async startLocalServer(): Promise<number> {
    const server = express()
    const port = await getAvailablePort(8000, 8100)
 
    server.get('/oauth/callback', async (req, res) => {
      const { code, state, error, error_description } = req.query
 
      if (error) {
        res.send(`<h1>Authorization failed</h1><p>${error_description}</p>`)
        this.handleError(error as string, error_description as string)
        return
      }
 
      // Handle successful callback
      try {
        await this.completeFlow(code as string, state as string)
        res.send('<h1>Success!</h1><p>You can close this window.</p>')
      } catch (err) {
        res.send(`<h1>Error</h1><p>${err.message}</p>`)
      }
    })
 
    await server.listen(port)
    return port
  }
}

// src-tauri/src/oauth.rs
use base64::{engine::general_purpose, Engine};
use rand::Rng;
use sha2::{Digest, Sha256};
use serde::{Deserialize, Serialize};
 
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PKCEChallenge {
    pub code_verifier: String,
    pub code_challenge: String,
    pub code_challenge_method: String,
}
 
/// Generate PKCE challenge pair
pub fn generate_pkce_challenge() -> PKCEChallenge {
    // Generate 32 random bytes
    let mut rng = rand::thread_rng();
    let random_bytes: Vec<u8> = (0..32).map(|_| rng.gen()).collect();
 
    // Base64URL encode (code verifier)
    let code_verifier = general_purpose::URL_SAFE_NO_PAD
        .encode(&random_bytes);
 
    // SHA-256 hash the verifier
    let mut hasher = Sha256::new();
    hasher.update(code_verifier.as_bytes());
    let hash_result = hasher.finalize();
 
    // Base64URL encode the hash (code challenge)
    let code_challenge = general_purpose::URL_SAFE_NO_PAD
        .encode(&hash_result);
 
    PKCEChallenge {
        code_verifier,
        code_challenge,
        code_challenge_method: "S256".to_string(),
    }
}
 
#[derive(Debug, Serialize, Deserialize)]
pub struct OAuthTokens {
    pub access_token: String,
    pub refresh_token: Option<String>,
    pub expires_in: u64,
    pub token_type: String,
    pub scope: Option<String>,
}
 
/// Exchange authorization code for tokens
pub async fn exchange_code(
    code: &str,
    code_verifier: &str,
    config: &OAuthConfig,
) -> Result<OAuthTokens, OAuthError> {
    let client = reqwest::Client::new();
 
    let params = [
        ("client_id", config.client_id.as_str()),
        ("code", code),
        ("code_verifier", code_verifier),
        ("grant_type", "authorization_code"),
        ("redirect_uri", config.redirect_uri.as_str()),
    ];
 
    let response = client
        .post(&config.token_endpoint)
        .form(&params)
        .send()
        .await?;
 
    if !response.status().is_success() {
        let error_body = response.text().await?;
        return Err(OAuthError::TokenExchangeFailed(error_body));
    }
 
    let tokens: OAuthTokens = response.json().await?;
    Ok(tokens)
}
 
#[derive(Debug, Serialize, Deserialize)]
pub struct OAuthConfig {
    pub client_id: String,
    pub redirect_uri: String,
    pub authorization_endpoint: String,
    pub token_endpoint: String,
    pub scopes: Vec<String>,
}
 
#[derive(Debug, thiserror::Error)]
pub enum OAuthError {
    #[error("Token exchange failed: {0}")]
    TokenExchangeFailed(String),
 
    #[error("Invalid state parameter")]
    InvalidState,
 
    #[error("Flow expired")]
    FlowExpired,
 
    #[error("Network error: {0}")]
    NetworkError(#[from] reqwest::Error),
}

use keyring::Keyring;
 
fn store_token(service: &str, account: &str, token: &str) -> Result<(), String> {
    let keyring = Keyring::new(service, account);
    keyring.set_password(token)
        .map_err(|e| e.to_string())
}
 
fn retrieve_token(service: &str, account: &str) -> Result<String, String> {
    let keyring = Keyring::new(service, account);
    keyring.get_password()
        .map_err(|e| e.to_string())
}

use windows::Security::Credentials::PasswordVault;
 
fn store_token_windows(resource: &str, username: &str, password: &str) {
    let vault = PasswordVault::new().unwrap();
    let credential = PasswordCredential::CreatePasswordCredential(
        resource,
        username,
        password
    ).unwrap();
    vault.Add(&credential).unwrap();
}

class TokenManager {
  private accessToken: string | null = null
  private refreshToken: string | null = null
  private expiresAt: number = 0
  private refreshPromise: Promise<void> | null = null
  private config: OAuthConfig
 
  constructor(config: OAuthConfig) {
    this.config = config
    this.loadTokens()
  }
 
  /**
   * Get valid access token (refreshes if expired)
   */
  async getAccessToken(): Promise<string> {
    // Check if token exists
    if (!this.accessToken) {
      throw new Error('Not authenticated')
    }
 
    // Check if token is expired (with 5 minute buffer)
    const buffer = 5 * 60 * 1000 // 5 minutes
    if (Date.now() >= this.expiresAt - buffer) {
      await this.refreshAccessToken()
    }
 
    return this.accessToken
  }
 
  /**
   * Refresh access token using refresh token
   * Thread-safe: multiple calls wait for single refresh
   */
  async refreshAccessToken(): Promise<void> {
    // If refresh is already in progress, wait for it
    if (this.refreshPromise) {
      return this.refreshPromise
    }
 
    // Start refresh
    this.refreshPromise = this._performRefresh()
 
    try {
      await this.refreshPromise
    } finally {
      this.refreshPromise = null
    }
  }
 
  private async _performRefresh(): Promise<void> {
    if (!this.refreshToken) {
      throw new Error('No refresh token available')
    }
 
    const maxRetries = 3
    let lastError: Error | null = null
 
    // Exponential backoff retry
    for (let attempt = 0; attempt < maxRetries; attempt++) {
      try {
        const newTokens = await this._exchangeRefreshToken()
 
        // Update tokens
        this.accessToken = newTokens.accessToken
        if (newTokens.refreshToken) {
          // Some providers rotate refresh tokens
          this.refreshToken = newTokens.refreshToken
        }
        this.expiresAt = newTokens.expiresAt
 
        // Store new tokens securely
        await this.storeTokens({
          accessToken: this.accessToken,
          refreshToken: this.refreshToken,
          expiresAt: this.expiresAt
        })
 
        return // Success!
      } catch (error) {
        lastError = error
 
        // Don't retry on certain errors
        if (this.isNonRetryableError(error)) {
          throw error
        }
 
        // Exponential backoff: 1s, 2s, 4s
        const delay = Math.pow(2, attempt) * 1000
        await this.sleep(delay)
      }
    }
 
    // All retries failed
    throw new Error(`Token refresh failed after ${maxRetries} attempts: ${lastError}`)
  }
 
  private async _exchangeRefreshToken(): Promise<TokenResponse> {
    const params = new URLSearchParams({
      client_id: this.config.clientId,
      refresh_token: this.refreshToken!,
      grant_type: 'refresh_token'
    })
 
    const response = await fetch(this.config.tokenEndpoint, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
      },
      body: params.toString()
    })
 
    if (!response.ok) {
      const error = await response.json()
      throw new TokenRefreshError(
        error.error_description || error.error,
        response.status,
        error.error
      )
    }
 
    const tokens = await response.json()
 
    return {
      accessToken: tokens.access_token,
      refreshToken: tokens.refresh_token, // May be null
      expiresIn: tokens.expires_in,
      tokenType: tokens.token_type,
      expiresAt: Date.now() + (tokens.expires_in * 1000)
    }
  }
 
  private isNonRetryableError(error: any): boolean {
    // Don't retry on these errors
    const nonRetryableCodes = [
      'invalid_grant', // Refresh token invalid/expired
      'unauthorized_client',
      'invalid_client'
    ]
 
    if (error instanceof TokenRefreshError) {
      return (
        nonRetryableCodes.includes(error.errorCode) ||
        error.httpStatus === 400
      )
    }
 
    return false
  }
 
  /**
   * Revoke tokens (sign out)
   */
  async revokeTokens(): Promise<void> {
    if (!this.accessToken) return
 
    try {
      // Revoke with provider (if supported)
      if (this.config.revocationEndpoint) {
        await fetch(this.config.revocationEndpoint, {
          method: 'POST',
          headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
          body: new URLSearchParams({
            token: this.accessToken,
            token_type_hint: 'access_token'
          }).toString()
        })
      }
    } catch (error) {
      console.error('Failed to revoke tokens:', error)
    } finally {
      // Clear local tokens regardless of revocation result
      await this.clearTokens()
    }
  }
 
  private sleep(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms))
  }
}
 
class TokenRefreshError extends Error {
  constructor(
    message: string,
    public httpStatus: number,
    public errorCode: string
  ) {
    super(message)
    this.name = 'TokenRefreshError'
  }
}
 
// Usage example
const tokenManager = new TokenManager({
  provider: 'gmail',
  clientId: GMAIL_CLIENT_ID,
  tokenEndpoint: 'https://oauth2.googleapis.com/token',
  revocationEndpoint: 'https://oauth2.googleapis.com/revoke'
})
 
// Get valid token (auto-refreshes if needed)
const accessToken = await tokenManager.getAccessToken()
 
// Use token
const response = await fetch('https://gmail.googleapis.com/gmail/v1/users/me/messages', {
  headers: {
    'Authorization': `Bearer ${accessToken}`
  }
})

function ensureHttps(url: string): string {
  if (url.startsWith('http://')) {
    return url.replace('http://', 'https://');
  }
  return url;
}
 
// In configuration
{
  "security": {
    "enforceHttps": true,
    "allowInsecureConnections": false
  }
}

const ALLOWED_ORIGINS = [
  'https://api.example.com',
  'https://mail.google.com'
];
 
function validateOrigin(origin: string): boolean {
  return ALLOWED_ORIGINS.includes(origin);
}