atman_runtime/
oauth_server.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::net::{Ipv4Addr, SocketAddr};
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::Duration;
7
8use anyhow::{Context, Result};
9use axum::extract::Query;
10use axum::response::Html;
11use axum::routing::get;
12use tokio::sync::Mutex;
13
14use crate::oauth::callback_page;
15
16type ExchangeFuture = Pin<Box<dyn Future<Output = std::result::Result<(), String>> + Send>>;
17type ExchangeFn = Box<dyn FnOnce(String) -> ExchangeFuture + Send>;
18
19pub async fn capture_oauth_callback(
20    port: u16,
21    expected_state: String,
22    exchange_fn: ExchangeFn,
23    timeout: Duration,
24) -> Result<()> {
25    let exchange: Arc<Mutex<Option<ExchangeFn>>> = Arc::new(Mutex::new(Some(exchange_fn)));
26    let result: Arc<Mutex<Option<Result<()>>>> = Arc::new(Mutex::new(None));
27
28    let expected = expected_state;
29    let exchange_for_handler = exchange.clone();
30    let result_for_handler = result.clone();
31
32    let app = axum::Router::new().route(
33        "/auth/callback",
34        get(move |Query(params): Query<HashMap<String, String>>| {
35            let exchange = exchange_for_handler.clone();
36            let result = result_for_handler.clone();
37            let expected = expected.clone();
38            async move {
39                let (outcome, page) = if params.get("state").map(|s| s != &expected).unwrap_or(true)
40                {
41                    (
42                        Err(anyhow::anyhow!("state mismatch")),
43                        callback_page(
44                            false,
45                            "State Mismatch",
46                            "The OAuth state parameter did not match.",
47                        ),
48                    )
49                } else if let Some(err) = params.get("error").cloned() {
50                    (
51                        Err(anyhow::anyhow!("oauth error: {err}")),
52                        callback_page(false, "授权被拒绝", &format!("授权服务器返回: {err}")),
53                    )
54                } else if let Some(code) = params.get("code").cloned() {
55                    let exchange_fn = exchange.lock().await.take();
56                    match exchange_fn {
57                        Some(f) => match f(code).await {
58                            Ok(()) => (
59                                Ok(()),
60                                callback_page(true, "认证成功", "已接入账户。您可以关闭此页面。"),
61                            ),
62                            Err(msg) => (
63                                Err(anyhow::anyhow!("token exchange failed: {msg}")),
64                                callback_page(false, "登录失败", &msg),
65                            ),
66                        },
67                        None => (
68                            Err(anyhow::anyhow!("duplicate callback")),
69                            callback_page(false, "请求无效", "重复的回调请求。"),
70                        ),
71                    }
72                } else {
73                    (
74                        Err(anyhow::anyhow!("missing code")),
75                        callback_page(false, "请求无效", "缺少授权码。"),
76                    )
77                };
78                *result.lock().await = Some(outcome);
79
80                Html(page)
81            }
82        }),
83    );
84
85    let socket = tokio::net::TcpSocket::new_v4().context("create tcp socket")?;
86    socket.set_reuseaddr(true).context("set SO_REUSEADDR")?;
87    socket
88        .bind(SocketAddr::new(Ipv4Addr::LOCALHOST.into(), port))
89        .with_context(|| format!("bind callback port {port}"))?;
90    let listener = socket
91        .listen(128)
92        .with_context(|| format!("listen callback port {port}"))?;
93
94    tokio::select! {
95        _ = axum::serve(listener, app) => {}
96        _ = tokio::time::sleep(timeout) => {}
97    }
98
99    result
100        .lock()
101        .await
102        .take()
103        .unwrap_or_else(|| Err(anyhow::anyhow!("no callback received")))
104}
atman_runtime/oauth_server.rs

atman_runtime/
oauth_server.rs
