OpenShell/crates/openshell-server/src/auth.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Auth endpoint for edge-proxy browser-based login.
//!
//! When the CLI runs `gateway add` or `gateway login`, it opens the user's
//! browser to `GET /auth/connect?callback_port=<port>`. The edge proxy
//! (e.g., Cloudflare Access) intercepts the request and shows its `IdP` login
//! page. After authentication, the proxy sets the `CF_Authorization` cookie
//! and proxies the request to this endpoint.
//!
//! The handler reads the `CF_Authorization` cookie from the request headers
//! (required because the cookie is typically `HttpOnly`) and serves a styled
//! confirmation page. When the user clicks "Connect", JavaScript POSTs
//! the token to `http://127.0.0.1:<callback_port>/callback`, where the
//! CLI's ephemeral localhost server captures and stores the token.

use axum::{
    Router,
    extract::{Query, State},
    http::HeaderMap,
    response::{Html, IntoResponse},
    routing::get,
};
use serde::Deserialize;
use std::sync::Arc;

use crate::ServerState;

#[derive(Deserialize)]
struct ConnectParams {
    callback_port: u16,
    /// One-time confirmation code generated by the CLI. Displayed on the
    /// connect page so the user can verify it matches their terminal before
    /// clicking Connect. Sent back in the XHR POST to the CLI callback
    /// server for cryptographic binding.
    code: String,
}

/// Create the auth router.
pub fn router(state: Arc<ServerState>) -> Router {
    Router::new()
        .route("/auth/connect", get(auth_connect))
        .with_state(state)
}

/// Handle the auth connect request.
///
/// Reads `CF_Authorization` from the cookie header (server-side extraction
/// handles `HttpOnly` cookies) and serves either a waiting page or a
/// styled confirmation page with the token embedded.
async fn auth_connect(
    State(state): State<Arc<ServerState>>,
    Query(params): Query<ConnectParams>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let cf_token = headers
        .get("cookie")
        .and_then(|v| v.to_str().ok())
        .and_then(|cookies| extract_cookie(cookies, "CF_Authorization"));

    // Prefer the Host header (set by Cloudflare Tunnel / reverse proxies)
    // so the page shows the external URL the user actually connected through
    // rather than the internal bind address (e.g. 0.0.0.0:8080).
    let gateway_display = headers
        .get("x-forwarded-host")
        .or_else(|| headers.get("host"))
        .and_then(|v| v.to_str().ok())
        .map_or_else(|| state.config.bind_address.to_string(), String::from);

    match cf_token {
        Some(token) => Html(render_connect_page(
            &gateway_display,
            params.callback_port,
            &token,
            &params.code,
        )),
        None => Html(render_waiting_page(params.callback_port, &params.code)),
    }
}

/// Extract a named cookie value from a `Cookie` header string.
fn extract_cookie(cookies: &str, name: &str) -> Option<String> {
    cookies.split(';').find_map(|c| {
        let mut parts = c.trim().splitn(2, '=');
        let key = parts.next()?.trim();
        let val = parts.next()?.trim();
        if key == name {
            Some(val.to_string())
        } else {
            None
        }
    })
}

/// Render the styled confirmation page with the edge auth token embedded.
///
/// The page displays the confirmation code so the user can verify it matches
/// their terminal. Clicking "Connect" sends an XHR POST (not a redirect) to
/// the CLI's localhost callback server, keeping the token out of browser
/// history and the URL bar.
fn render_connect_page(
    gateway_addr: &str,
    callback_port: u16,
    cf_token: &str,
    code: &str,
) -> String {
    // Escape the token for safe embedding in a JS string literal.
    let escaped_token = cf_token
        .replace('\\', "\\\\")
        .replace('\'', "\\'")
        .replace('"', "\\\"")
        .replace('<', "\\x3c")
        .replace('>', "\\x3e");

    // Escape the code the same way (it's alphanumeric + dash, but be safe).
    let escaped_code = code
        .replace('\\', "\\\\")
        .replace('\'', "\\'")
        .replace('"', "\\\"")
        .replace('<', "\\x3c")
        .replace('>', "\\x3e");

    let version = openshell_core::VERSION;

    format!(
        r#"<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>OpenShell — Connect to Gateway</title>
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        body {{
            font-family: 'SF Mono', 'Fira Code', 'JetBrains Mono', monospace;
            background: #f5f5f5;
            color: #1a1a1a;
            min-height: 100vh;
            display: flex;
            align-items: center;
            justify-content: center;
        }}
        .card {{
            background: #ffffff;
            border: 1px solid #e0e0e0;
            border-radius: 12px;
            padding: 48px;
            max-width: 480px;
            width: 100%;
            text-align: center;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
        }}
        .logo {{
            font-size: 28px;
            font-weight: 700;
            color: #76b900;
            margin-bottom: 8px;
            letter-spacing: -0.5px;
        }}
        .subtitle {{
            color: #666;
            font-size: 14px;
            margin-bottom: 32px;
        }}
        .info {{
            background: #fafafa;
            border: 1px solid #e0e0e0;
            border-radius: 8px;
            padding: 16px;
            margin-bottom: 24px;
            text-align: left;
        }}
        .info-row {{
            display: flex;
            justify-content: space-between;
            padding: 6px 0;
            font-size: 13px;
        }}
        .info-label {{ color: #888; }}
        .info-value {{
            color: #333;
            font-weight: 500;
            word-break: break-all;
            text-align: right;
            max-width: 60%;
        }}
        .code-box {{
            background: #fafafa;
            border: 2px solid #76b900;
            border-radius: 8px;
            padding: 20px;
            margin-bottom: 24px;
        }}
        .code-label {{
            color: #666;
            font-size: 12px;
            text-transform: uppercase;
            letter-spacing: 1px;
            margin-bottom: 8px;
        }}
        .code-value {{
            font-size: 32px;
            font-weight: 700;
            color: #76b900;
            letter-spacing: 4px;
        }}
        .code-hint {{
            color: #888;
            font-size: 11px;
            margin-top: 8px;
        }}
        .connect-btn {{
            background: #76b900;
            color: #ffffff;
            border: none;
            border-radius: 8px;
            padding: 14px 32px;
            font-size: 15px;
            font-weight: 600;
            font-family: inherit;
            cursor: pointer;
            width: 100%;
            transition: background 0.15s;
        }}
        .connect-btn:hover {{ background: #8ad100; }}
        .connect-btn:active {{ background: #6aa000; }}
        .connect-btn:disabled {{
            background: #ccc;
            color: #888;
            cursor: not-allowed;
        }}
        .hint {{
            color: #888;
            font-size: 12px;
            margin-top: 16px;
        }}
        .status {{
            margin-top: 16px;
            font-size: 13px;
            min-height: 20px;
        }}
        .status-ok {{ color: #76b900; }}
        .status-err {{ color: #d32f2f; }}
    </style>
</head>
<body>
    <div class="card">
        <div class="logo">OpenShell</div>
        <div class="subtitle">Connect to Gateway</div>
        <div class="code-box">
            <div class="code-label">Confirmation Code</div>
            <div class="code-value">{escaped_code}</div>
            <div class="code-hint">Verify this matches the code shown in your terminal</div>
        </div>
        <div class="info">
            <div class="info-row">
                <span class="info-label">Gateway</span>
                <span class="info-value">{gateway_addr}</span>
            </div>
            <div class="info-row">
                <span class="info-label">Version</span>
                <span class="info-value">v{version}</span>
            </div>
        </div>
        <button class="connect-btn" id="connectBtn" onclick="connect()">
            Connect to Gateway
        </button>
        <div class="hint">
            This will authorize the OpenShell CLI to connect to this gateway.
        </div>
        <div class="status" id="status"></div>
    </div>
    <script>
        var token = '{escaped_token}';
        var code = '{escaped_code}';
        var port = {callback_port};
        function connect() {{
            var btn = document.getElementById('connectBtn');
            var status = document.getElementById('status');
            btn.disabled = true;
            btn.textContent = 'Connecting...';
            status.textContent = '';
            status.className = 'status';

            fetch('http://127.0.0.1:' + port + '/callback', {{
                method: 'POST',
                headers: {{ 'Content-Type': 'application/json' }},
                body: JSON.stringify({{ token: token, code: code }})
            }})
            .then(function(resp) {{ return resp.json(); }})
            .then(function(data) {{
                if (data.ok) {{
                    btn.textContent = 'Connected';
                    status.className = 'status status-ok';
                    status.textContent = 'Connected! You can close this tab.';
                }} else {{
                    btn.disabled = false;
                    btn.textContent = 'Connect to Gateway';
                    status.className = 'status status-err';
                    status.textContent = data.error || 'Connection failed. Please try again.';
                }}
            }})
            .catch(function(err) {{
                btn.disabled = false;
                btn.textContent = 'Connect to Gateway';
                status.className = 'status status-err';
                status.textContent = 'Could not reach the CLI. Is openshell still running?';
            }});
        }}
    </script>
</body>
</html>"#,
    )
}

/// Minimal HTML attribute escaping for untrusted values.
///
/// Escapes `&`, `"`, `'`, `<`, and `>` so the value is safe for embedding
/// inside a quoted HTML attribute.
fn html_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        match ch {
            '&' => out.push_str("&amp;"),
            '"' => out.push_str("&quot;"),
            '\'' => out.push_str("&#x27;"),
            '<' => out.push_str("&lt;"),
            '>' => out.push_str("&gt;"),
            c => out.push(c),
        }
    }
    out
}

/// Render a waiting page shown when the edge proxy auth cookie is not yet present.
///
/// This can happen if the edge proxy hasn't completed authentication yet. The
/// page auto-reloads after a short delay to pick up the cookie once login
/// finishes. The `code` parameter is preserved across reloads.
fn render_waiting_page(callback_port: u16, code: &str) -> String {
    let safe_code = html_escape(code);
    format!(
        r#"<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <meta http-equiv="refresh" content="2;url=/auth/connect?callback_port={callback_port}&amp;code={safe_code}">
    <title>OpenShell — Authenticating</title>
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        body {{
            font-family: 'SF Mono', 'Fira Code', 'JetBrains Mono', monospace;
            background: #f5f5f5;
            color: #1a1a1a;
            min-height: 100vh;
            display: flex;
            align-items: center;
            justify-content: center;
        }}
        .card {{
            background: #ffffff;
            border: 1px solid #e0e0e0;
            border-radius: 12px;
            padding: 48px;
            max-width: 480px;
            width: 100%;
            text-align: center;
            box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
        }}
        .logo {{
            font-size: 28px;
            font-weight: 700;
            color: #76b900;
            margin-bottom: 8px;
        }}
        .message {{
            color: #666;
            font-size: 14px;
            margin-top: 16px;
        }}
        .spinner {{
            margin: 24px auto;
            width: 32px;
            height: 32px;
            border: 3px solid #e0e0e0;
            border-top-color: #76b900;
            border-radius: 50%;
            animation: spin 0.8s linear infinite;
        }}
        @keyframes spin {{
            to {{ transform: rotate(360deg); }}
        }}
    </style>
</head>
<body>
    <div class="card">
        <div class="logo">OpenShell</div>
        <div class="spinner"></div>
        <div class="message">Authenticating...</div>
    </div>
</body>
</html>"#,
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn extract_cookie_finds_value() {
        let cookies = "other=foo; CF_Authorization=eyJhbGciOiJSUzI1NiJ9.test; bar=baz";
        assert_eq!(
            extract_cookie(cookies, "CF_Authorization"),
            Some("eyJhbGciOiJSUzI1NiJ9.test".to_string())
        );
    }

    #[test]
    fn extract_cookie_returns_none_when_missing() {
        let cookies = "other=foo; bar=baz";
        assert_eq!(extract_cookie(cookies, "CF_Authorization"), None);
    }

    #[test]
    fn extract_cookie_handles_single_cookie() {
        let cookies = "CF_Authorization=my-token";
        assert_eq!(
            extract_cookie(cookies, "CF_Authorization"),
            Some("my-token".to_string())
        );
    }

    #[test]
    fn render_connect_page_contains_token_and_code() {
        let html = render_connect_page(
            "gateway.example.com:8080",
            12345,
            "test-jwt-token",
            "ABC-1234",
        );
        assert!(html.contains("test-jwt-token"));
        assert!(html.contains("12345"));
        assert!(html.contains("gateway.example.com:8080"));
        assert!(html.contains("OpenShell"));
        assert!(html.contains("ABC-1234"));
        // Should use fetch(), not window.location.href
        assert!(html.contains("fetch("));
        assert!(!html.contains("window.location.href"));
        // Should POST with JSON
        assert!(html.contains("method: 'POST'"));
        assert!(html.contains("JSON.stringify"));
    }

    #[test]
    fn render_connect_page_escapes_special_chars() {
        let html =
            render_connect_page("gw", 1234, "token<script>alert('xss')</script>", "ABC-1234");
        // < and > should be escaped
        assert!(!html.contains("<script>alert"));
        assert!(html.contains("\\x3c"));
    }

    #[test]
    fn render_connect_page_includes_code_in_js_payload() {
        let html = render_connect_page("gw", 1234, "jwt", "XY7-9KLM");
        // The JS should send the code in the JSON payload
        assert!(html.contains("var code = 'XY7-9KLM'"));
        assert!(html.contains("code: code"));
    }

    #[test]
    fn render_waiting_page_has_auto_refresh_with_code() {
        let html = render_waiting_page(12345, "ABC-1234");
        assert!(html.contains("meta http-equiv=\"refresh\""));
        assert!(html.contains("callback_port=12345"));
        assert!(html.contains("code=ABC-1234"));
    }

    #[test]
    fn render_waiting_page_escapes_code() {
        let html = render_waiting_page(1234, "x\"onload=\"alert(1)");
        // The double quote should be escaped so it can't break out of the
        // meta-refresh content attribute.
        assert!(!html.contains("x\"onload"));
        assert!(html.contains("&quot;"));
    }

    #[test]
    fn html_escape_basic() {
        assert_eq!(html_escape("hello"), "hello");
        assert_eq!(html_escape("<b>bold</b>"), "&lt;b&gt;bold&lt;/b&gt;");
        assert_eq!(html_escape("a&b"), "a&amp;b");
        assert_eq!(html_escape("a\"b"), "a&quot;b");
        assert_eq!(html_escape("a'b"), "a&#x27;b");
    }
}