add flexible env for auth url

Author: evanliu048

crates/chat-cli/src/auth/portal.rs

@@ -395,3 +395,387 @@ async fn bind_allowed_port(ports: &[u16]) -> Result<TcpListener, AuthError> {
 fn get_auth_portal_url() -> String {
     env::var("KIRO_AUTH_PORTAL_URL").unwrap_or_else(|_| DEFAULT_AUTH_PORTAL_URL.to_string())
 }
+//! Unified auth portal integration for streamlined authentication
+//! Handles callbacks from https://app.kiro.dev/signin
+
+use std::time::Duration;
+
+use bytes::Bytes;
+use http_body_util::Full;
+use hyper::body::Incoming;
+use hyper::server::conn::http1;
+use hyper::service::Service;
+use hyper::{
+    Request,
+    Response,
+};
+use hyper_util::rt::TokioIo;
+use rand::Rng;
+use tokio::net::TcpListener;
+use tracing::{
+    debug,
+    error,
+    info,
+    warn,
+};
+
+use crate::auth::AuthError;
+use crate::auth::pkce::{
+    generate_code_challenge,
+    generate_code_verifier,
+};
+use crate::auth::social::{
+    CALLBACK_PORTS,
+    SocialProvider,
+    SocialToken,
+};
+use crate::database::Database;
+use crate::util::system_info::is_mwinit_available;
+
+const AUTH_PORTAL_URL: &str = "https://app.kiro.dev/signin";
+const DEFAULT_AUTHORIZATION_TIMEOUT: Duration = Duration::from_secs(600);
+
+#[derive(Debug, Clone)]
+struct AuthPortalCallback {
+    login_option: String,
+    code: Option<String>,
+    issuer_url: Option<String>,
+    sso_region: Option<String>,
+    state: String,
+    path: String,
+    error: Option<String>,
+    error_description: Option<String>,
+}
+
+pub enum PortalResult {
+    Social(SocialProvider),
+    BuilderId {
+        issuer_url: String,
+        idc_region: String,
+    },
+    AwsIdc {
+        issuer_url: String,
+        idc_region: String,
+    },
+    /// Internal amazon user
+    Internal {
+        issuer_url: String,
+        idc_region: String,
+    },
+}
+
+/// Local-only: open unified portal and handle single callback
+pub async fn start_unified_auth(db: &mut Database) -> Result<PortalResult, AuthError> {
+    info!("Starting unified auth portal flow");
+
+    // PKCE params for portal + social token exchange
+    let verifier = generate_code_verifier();
+    let challenge = generate_code_challenge(&verifier);
+    let state = rand::rng()
+        .sample_iter(rand::distr::Alphanumeric)
+        .take(10)
+        .collect::<Vec<_>>();
+    let state = String::from_utf8(state).unwrap_or("state".to_string());
+
+    let listener = bind_allowed_port(CALLBACK_PORTS).await?;
+    let port = listener.local_addr()?.port();
+
+    let redirect_base = format!("http://localhost:{port}");
+    info!(%port, %redirect_base, "Unified auth portal listening for callback");
+
+    let auth_url = build_auth_url(&redirect_base, &state, &challenge);
+
+    crate::util::open::open_url_async(&auth_url)
+        .await
+        .map_err(|e| AuthError::OAuthCustomError(format!("Failed to open browser: {e}")))?;
+
+    let callback = wait_for_auth_callback(listener, state.clone()).await?;
+
+    if let Some(error) = &callback.error {
+        let friendly_msg =
+            format_user_friendly_error(error, callback.error_description.as_deref(), &callback.login_option);
+
+        warn!(
+            "OAuth error for {}: {} - {}",
+            callback.login_option, error, friendly_msg
+        );
+
+        return Err(match callback.login_option.as_str() {
+            "google" | "github" => AuthError::SocialAuthProviderFailure(friendly_msg),
+            _ => AuthError::OAuthCustomError(friendly_msg),
+        });
+    }
+
+    process_portal_callback(db, callback, port, &verifier).await
+}
+
+fn format_user_friendly_error(error_code: &str, description: Option<&str>, provider: &str) -> String {
+    let cleaned_description = description.map(|d| {
+        let first_part = d.split(';').next().unwrap_or(d);
+        // Replace + with spaces (URL encoding)
+        first_part.replace('+', " ").trim().to_string()
+    });
+
+    match error_code {
+        "access_denied" => {
+            format!("{provider} denied access to Kiro. Please ensure you grant all required permissions.")
+        },
+        "invalid_request" => "Authentication failed due to an invalid request. Please try again.".to_string(),
+        "unauthorized_client" => "The application is not authorized. Please contact support.".to_string(),
+        "server_error" => {
+            format!("{provider} login is temporarily unavailable. Please try again later.")
+        },
+        "invalid_scope" => "The requested permissions are invalid. Please contact support.".to_string(),
+        _ => {
+            // For unknown errors, use cleaned description or a generic message
+            cleaned_description.unwrap_or_else(|| format!("Authentication failed: {error_code}. Please try again."))
+        },
+    }
+}
+
+/// Build the authorization URL with all required parameters
+fn build_auth_url(redirect_base: &str, state: &str, challenge: &str) -> String {
+    let is_internal = is_mwinit_available();
+    let internal_param = if is_internal { "&from_amazon_internal=true" } else { "" };
+
+    format!(
+        "{}?state={}&code_challenge={}&code_challenge_method=S256&redirect_uri={}{}&redirect_from=kirocli",
+        AUTH_PORTAL_URL,
+        state,
+        challenge,
+        urlencoding::encode(redirect_base),
+        internal_param
+    )
+}
+
+async fn process_portal_callback(
+    db: &mut Database,
+    callback: AuthPortalCallback,
+    port: u16,
+    verifier: &str,
+) -> Result<PortalResult, AuthError> {
+    match callback.login_option.as_str() {
+        "google" | "github" => handle_social_callback(db, callback, port, verifier).await,
+        "internal" => {
+            let (issuer_url, sso_region) = extract_sso_params(&callback, "internal")?;
+            Ok(PortalResult::Internal {
+                issuer_url,
+                idc_region: sso_region,
+            })
+        },
+        "awsidc" => {
+            let (issuer_url, sso_region) = extract_sso_params(&callback, "awsIdc")?;
+            Ok(PortalResult::AwsIdc {
+                issuer_url,
+                idc_region: sso_region,
+            })
+        },
+        "builderid" => {
+            let (issuer_url, sso_region) = extract_sso_params(&callback, "builderId")?;
+            Ok(PortalResult::BuilderId {
+                issuer_url,
+                idc_region: sso_region,
+            })
+        },
+        other => Err(AuthError::OAuthCustomError(format!("Unknown login_option: {other}"))),
+    }
+}
+
+/// Handle social provider callback (Google/GitHub)
+async fn handle_social_callback(
+    db: &mut Database,
+    callback: AuthPortalCallback,
+    port: u16,
+    verifier: &str,
+) -> Result<PortalResult, AuthError> {
+    let provider = match callback.login_option.as_str() {
+        "google" => SocialProvider::Google,
+        "github" => SocialProvider::Github,
+        _ => unreachable!(),
+    };
+
+    let code = callback.code.ok_or(AuthError::OAuthMissingCode)?;
+    let redirect_uri = format!(
+        "http://localhost:{}{}?login_option={}",
+        port,
+        callback.path,
+        urlencoding::encode(&callback.login_option)
+    );
+
+    SocialToken::exchange_social_token(db, provider, verifier, &code, &redirect_uri).await?;
+    Ok(PortalResult::Social(provider))
+}
+
+/// Extract issuer_url and sso_region from callback, returning descriptive error if missing
+fn extract_sso_params(callback: &AuthPortalCallback, auth_type: &str) -> Result<(String, String), AuthError> {
+    let issuer_url = callback
+        .issuer_url
+        .clone()
+        .ok_or_else(|| AuthError::OAuthCustomError(format!("Missing issuer_url for {auth_type} auth")))?;
+
+    let sso_region = callback
+        .sso_region
+        .clone()
+        .ok_or_else(|| AuthError::OAuthCustomError(format!("Missing sso_region for {auth_type} auth")))?;
+
+    Ok((issuer_url, sso_region))
+}
+
+async fn wait_for_auth_callback(
+    listener: TcpListener,
+    expected_state: String,
+) -> Result<AuthPortalCallback, AuthError> {
+    let (tx, mut rx) = tokio::sync::mpsc::channel::<AuthPortalCallback>(1);
+
+    let server_handle = tokio::spawn(async move {
+        const MAX_CONNECTIONS: usize = 3;
+        let mut count = 0;
+
+        loop {
+            if count >= MAX_CONNECTIONS {
+                warn!("Reached max connections ({})", MAX_CONNECTIONS);
+                break;
+            }
+
+            match listener.accept().await {
+                Ok((stream, _)) => {
+                    count += 1;
+                    debug!("Connection {}/{}", count, MAX_CONNECTIONS);
+
+                    let io = TokioIo::new(stream);
+                    let service = AuthCallbackService { tx: tx.clone() };
+
+                    tokio::spawn(async move {
+                        let _ = http1::Builder::new().serve_connection(io, service).await;
+                    });
+                },
+                Err(e) => {
+                    error!("Accept failed: {}", e);
+                    break;
+                },
+            }
+        }
+    });
+
+    let callback = tokio::select! {
+        result = rx.recv() => {
+            result.ok_or(AuthError::OAuthCustomError("Failed to receive callback".into()))?
+        },
+        _ = tokio::time::sleep(DEFAULT_AUTHORIZATION_TIMEOUT) => {
+            return Err(AuthError::OAuthTimeout);
+        }
+    };
+
+    server_handle.abort();
+
+    if callback.state != expected_state {
+        return Err(AuthError::OAuthStateMismatch {
+            actual: callback.state,
+            expected: expected_state,
+        });
+    }
+
+    Ok(callback)
+}
+
+#[derive(Clone)]
+struct AuthCallbackService {
+    tx: tokio::sync::mpsc::Sender<AuthPortalCallback>,
+}
+
+impl Service<Request<Incoming>> for AuthCallbackService {
+    type Error = AuthError;
+    type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+    type Response = Response<Full<Bytes>>;
+
+    fn call(&self, req: Request<Incoming>) -> Self::Future {
+        let tx = self.tx.clone();
+
+        Box::pin(async move {
+            let uri = req.uri();
+            let path = uri.path();
+
+            if path == "/oauth/callback" || path == "/signin/callback" {
+                handle_valid_callback(uri, path, tx).await
+            } else {
+                handle_invalid_callback(path).await
+            }
+        })
+    }
+}
+
+/// Handle valid callback paths
+async fn handle_valid_callback(
+    uri: &hyper::Uri,
+    path: &str,
+    tx: tokio::sync::mpsc::Sender<AuthPortalCallback>,
+) -> Result<Response<Full<Bytes>>, AuthError> {
+    let query_params = uri
+        .query()
+        .map(|query| {
+            query
+                .split('&')
+                .filter_map(|kv| {
+                    kv.split_once('=')
+                        .map(|(k, v)| (k.to_string(), urlencoding::decode(v).unwrap_or_default().to_string()))
+                })
+                .collect::<std::collections::HashMap<String, String>>() // 
+        })
+        .ok_or(AuthError::OAuthCustomError("query parameters are missing".into()))?;
+
+    let callback = AuthPortalCallback {
+        login_option: query_params.get("login_option").cloned().unwrap_or_default(),
+        code: query_params.get("code").cloned(),
+        issuer_url: query_params.get("issuer_url").cloned(),
+        sso_region: query_params.get("idc_region").cloned(),
+        state: query_params.get("state").cloned().unwrap_or_default(),
+        path: path.to_string(),
+        error: query_params.get("error").cloned(),
+        error_description: query_params.get("error_description").cloned(),
+    };
+
+    let _ = tx.send(callback.clone()).await;
+
+    if let Some(error) = &callback.error {
+        let error_msg = callback.error_description.as_deref().unwrap_or(error.as_str());
+        build_redirect_response("error", Some(error_msg))
+    } else {
+        build_redirect_response("success", None)
+    }
+}
+
+async fn handle_invalid_callback(path: &str) -> Result<Response<Full<Bytes>>, AuthError> {
+    info!(%path, "Invalid callback path: {}, redirecting to portal", path);
+    build_redirect_response("error", Some("Invalid callback path"))
+}
+
+/// Build a redirect response to the auth portal
+fn build_redirect_response(status: &str, error_message: Option<&str>) -> Result<Response<Full<Bytes>>, AuthError> {
+    let mut redirect_url = format!("{AUTH_PORTAL_URL}?auth_status={status}&redirect_from=kirocli");
+
+    if let Some(msg) = error_message {
+        redirect_url.push_str(&format!("&error_message={}", urlencoding::encode(msg)));
+    }
+
+    Ok(Response::builder()
+        .status(302)
+        .header("Location", redirect_url)
+        .header("Cache-Control", "no-store")
+        .body(Full::new(Bytes::from("")))
+        .expect("valid response"))
+}
+
+async fn bind_allowed_port(ports: &[u16]) -> Result<TcpListener, AuthError> {
+    for port in ports {
+        match TcpListener::bind(("127.0.0.1", *port)).await {
+            Ok(listener) => return Ok(listener),
+            Err(e) => {
+                debug!("Failed to bind to port {}: {}", port, e);
+            },
+        }
+    }
+
+    Err(AuthError::OAuthCustomError(
+        "All callback ports are in use. Please close some applications and try again.".into(),
+    ))
+}