yggdrasil/api/
mcp_tokens.rs

1//! MCP 访问令牌管理：Dioxus server functions。
2//!
3//! 管理员在后台 `/admin/mcp` 签发/查看/编辑权限/撤销为 AI 客户端（Claude Code / Cursor /
4//! Cline）准备的 bearer 令牌。明文 token 仅在签发与「重新查看」时返回给管理员，
5//! 数据库只存 AES-GCM 密文（`token_enc`，可解密重查）+ SHA-256 哈希（`token_hash`，
6//! 每请求 O(1) 常量查找，见 `src/mcp/auth.rs`）。
7//!
8//! 鉴权走 cookie session（`get_current_admin_user`），与其它后台 server-fn 一致；
9//! MCP 工具路径（bearer）无法调用这些 server-fn——那是 `src/mcp/tools/*` 的职责。
10
11#![allow(clippy::unused_unit, deprecated)]
12
13use dioxus::prelude::*;
14
15#[cfg(feature = "server")]
16use crate::models::mcp_token::McpToken;
17use crate::models::mcp_token::{CreateTokenResponse, McpTokenSummary, NoteGrant, TokenScope};
18
19/// 令牌有效期预设：管理员在 UI 上从下拉菜单选择。
20///
21/// 序列化形式供前端选择回传（`days1` / `days7` / `days30` / `days90` / `never`）。
22/// `Never` 对应 `expires_at = NULL`（长期令牌）；其余按当前时间 + N 天计算。
23#[derive(Debug, Clone, Copy, serde::Deserialize, serde::Serialize, PartialEq, Eq)]
24#[serde(rename_all = "lowercase")]
25pub enum TokenLifetime {
26    /// 1 天（默认推荐：最小权限、轮换友好）。
27    Days1,
28    /// 7 天。
29    Days7,
30    /// 30 天。
31    Days30,
32    /// 90 天。
33    Days90,
34    /// 永不过期（`expires_at = NULL`）。仅用于可信长期客户端。
35    Never,
36}
37
38impl TokenLifetime {
39    /// 计算签发时刻对应的过期时间戳（UTC）。`Never` 返回 `None`。
40    #[cfg(feature = "server")]
41    fn expires_at(self) -> Option<chrono::DateTime<chrono::Utc>> {
42        let now = chrono::Utc::now();
43        match self {
44            TokenLifetime::Days1 => Some(now + chrono::Duration::days(1)),
45            TokenLifetime::Days7 => Some(now + chrono::Duration::days(7)),
46            TokenLifetime::Days30 => Some(now + chrono::Duration::days(30)),
47            TokenLifetime::Days90 => Some(now + chrono::Duration::days(90)),
48            TokenLifetime::Never => None,
49        }
50    }
51}
52
53/// 签发新的 MCP 令牌。
54///
55/// 生成明文 `ygg_<32 hex>`，AES-GCM 加密后存密文 + SHA-256 哈希；明文随响应一次性
56/// 返回给管理员（后续可经 [`reveal_mcp_token`] 重新查看）。仅 admin。
57#[server(CreateMcpToken, "/api")]
58pub async fn create_mcp_token(
59    name: String,
60    scope: TokenScope,
61    lifetime: TokenLifetime,
62    notes: Option<crate::models::mcp_token::NoteGrant>,
63) -> Result<CreateTokenResponse, ServerFnError> {
64    #[cfg(feature = "server")]
65    {
66        use crate::api::auth::get_current_admin_user;
67        use crate::api::error::AppError;
68        use crate::db::pool::get_conn;
69        use crate::mcp::auth::{hash_token, TOKEN_PREFIX};
70        use crate::mcp::crypto::encrypt_token;
71
72        let admin = get_current_admin_user().await?;
73
74        // 名称规范化与校验：去空白后非空，限制长度。
75        let name = name.trim().to_string();
76        crate::models::mcp_token::validate_token_name(&name)
77            .map_err(|message| AppError::BadRequest(message.to_string()))?;
78
79        // 加密主密钥必须已配置，否则无法安全存储明文。
80        if crate::mcp::crypto::mcp_enc_key().is_none() {
81            return Err(AppError::Internal("MCP_TOKEN_ENC_KEY 未设置").into());
82        }
83
84        // 明文 token：`ygg_` + 32 字节随机数 hex（64 hex 字符）。
85        let mut bytes = [0u8; 32];
86        rand::TryRng::try_fill_bytes(&mut rand::rngs::SysRng, &mut bytes)
87            .map_err(|_| AppError::Internal("令牌随机数生成失败"))?;
88        let plaintext = format!("{TOKEN_PREFIX}{}", hex::encode(bytes));
89        let hash = hash_token(&plaintext);
90        let enc = encrypt_token(&plaintext).ok_or(AppError::Internal("MCP 令牌加密失败"))?;
91        let id = uuid::Uuid::new_v4();
92        let expires_at = lifetime.expires_at();
93        let scope_str = scope.as_str();
94
95        let client = get_conn().await.map_err(AppError::db_conn)?;
96
97        let notes = validate_note_grant(&client, admin.id, notes.unwrap_or_default()).await?;
98
99        let row = client
100            .query_one(
101                "INSERT INTO mcp_tokens \
102                    (id, user_id, name, scope, token_enc, token_hash, expires_at, notes_read, notes_write, notebook_ids) \
103                 VALUES ($1::uuid, $2, $3, $4, $5, $6, $7, $8, $9, $10) \
104                 RETURNING id::text, user_id, name, scope, created_at, expires_at, \
105                           last_used_at, revoked_at, notes_read, notes_write, notebook_ids",
106                &[&id, &admin.id, &name, &scope_str, &enc, &hash, &expires_at, &notes.read, &notes.write, &notes.notebook_ids],
107            )
108            .await
109            .map_err(AppError::query)?;
110
111        let token = row_to_mcp_token_meta(&row);
112        Ok(CreateTokenResponse {
113            summary: token.into(),
114            plaintext,
115        })
116    }
117    #[cfg(not(feature = "server"))]
118    unreachable!()
119}
120
121/// 修改本人名下有效令牌的权限，不轮换密钥，也不修改有效期。
122#[server(UpdateMcpTokenPermissions, "/api")]
123pub async fn update_mcp_token_permissions(
124    id: String,
125    scope: TokenScope,
126    notes: NoteGrant,
127) -> Result<(), ServerFnError> {
128    let admin = crate::api::auth::get_current_admin_user().await?;
129    let client = crate::db::pool::get_conn()
130        .await
131        .map_err(crate::api::error::AppError::db_conn)?;
132    save_token_permissions(&client, admin.id, &id, scope, notes).await?;
133    Ok(())
134}
135
136#[cfg(feature = "server")]
137async fn validate_note_grant(
138    client: &tokio_postgres::Client,
139    owner_id: i32,
140    mut notes: NoteGrant,
141) -> Result<NoteGrant, crate::api::error::AppError> {
142    use crate::api::error::AppError;
143    notes.read |= notes.write;
144    if !notes.read {
145        notes.notebook_ids = None;
146    }
147    if let Some(ids) = notes.notebook_ids.as_mut() {
148        ids.sort_unstable();
149        ids.dedup();
150        let owned: i64 = client
151            .query_one(
152                "SELECT COUNT(*) FROM notebooks WHERE owner_id=$1 AND id=ANY($2)",
153                &[&owner_id, &ids.as_slice()],
154            )
155            .await
156            .map_err(AppError::query)?
157            .get(0);
158        if ids.is_empty() || owned != ids.len() as i64 {
159            return Err(AppError::BadRequest("请选择自己拥有的笔记本".into()));
160        }
161    }
162    Ok(notes)
163}
164
165#[cfg(feature = "server")]
166async fn save_token_permissions(
167    client: &tokio_postgres::Client,
168    owner_id: i32,
169    id: &str,
170    scope: TokenScope,
171    notes: NoteGrant,
172) -> Result<(), crate::api::error::AppError> {
173    use crate::api::error::AppError;
174    let invalid_token = || AppError::BadRequest("令牌不存在或已失效，请刷新列表".into());
175    let id = uuid::Uuid::parse_str(id).map_err(|_| invalid_token())?;
176    let notes = validate_note_grant(client, owner_id, notes).await?;
177    let changed = client
178        .execute(
179            "UPDATE mcp_tokens SET scope=$3, notes_read=$4, notes_write=$5, notebook_ids=$6 \
180             WHERE id=$1 AND user_id=$2 AND revoked_at IS NULL \
181               AND (expires_at IS NULL OR expires_at>NOW())",
182            &[
183                &id,
184                &owner_id,
185                &scope.as_str(),
186                &notes.read,
187                &notes.write,
188                &notes.notebook_ids,
189            ],
190        )
191        .await
192        .map_err(AppError::query)?;
193    if changed == 0 {
194        return Err(invalid_token());
195    }
196    Ok(())
197}
198
199#[cfg(all(test, feature = "server"))]
200mod permission_tests {
201    use super::*;
202
203    #[tokio::test]
204    #[ignore = "requires YGGDRASIL_TEST_DATABASE_URL; uses connection-local temporary tables"]
205    async fn edit_permissions_preserves_token_and_enforces_ownership_and_validity() {
206        let url = std::env::var("YGGDRASIL_TEST_DATABASE_URL").unwrap();
207        let (client, connection) = tokio_postgres::connect(&url, tokio_postgres::NoTls)
208            .await
209            .unwrap();
210        let connection = tokio::spawn(connection);
211        client.batch_execute(
212            "CREATE TEMP TABLE notebooks (id INT PRIMARY KEY, owner_id INT NOT NULL);
213             INSERT INTO notebooks VALUES (10,1),(20,2);
214             CREATE TEMP TABLE mcp_tokens (
215                 id UUID PRIMARY KEY, user_id INT NOT NULL, scope TEXT NOT NULL,
216                 notes_read BOOLEAN DEFAULT false, notes_write BOOLEAN DEFAULT false,
217                 notebook_ids INT[], revoked_at TIMESTAMPTZ, expires_at TIMESTAMPTZ,
218                 token_hash TEXT DEFAULT 'original-hash', token_enc TEXT DEFAULT 'original-ciphertext'
219             )"
220        ).await.unwrap();
221        let id = uuid::Uuid::new_v4();
222        client
223            .execute(
224                "INSERT INTO mcp_tokens(id,user_id,scope) VALUES($1,1,'admin')",
225                &[&id],
226            )
227            .await
228            .unwrap();
229        let grant = NoteGrant {
230            read: false,
231            write: true,
232            notebook_ids: Some(vec![10, 10]),
233        };
234        save_token_permissions(&client, 1, &id.to_string(), TokenScope::Read, grant)
235            .await
236            .unwrap();
237        let row = client
238            .query_one("SELECT * FROM mcp_tokens WHERE id=$1", &[&id])
239            .await
240            .unwrap();
241        assert!(row.get::<_, bool>("notes_read"));
242        assert!(row.get::<_, bool>("notes_write"));
243        assert_eq!(row.get::<_, Vec<i32>>("notebook_ids"), vec![10]);
244        assert_eq!(row.get::<_, String>("scope"), "read");
245        assert_eq!(row.get::<_, String>("token_hash"), "original-hash");
246        assert_eq!(row.get::<_, String>("token_enc"), "original-ciphertext");
247        assert!(row
248            .get::<_, Option<chrono::DateTime<chrono::Utc>>>("expires_at")
249            .is_none());
250
251        for ids in [vec![], vec![20], vec![10, 20], vec![999]] {
252            assert!(save_token_permissions(
253                &client,
254                1,
255                &id.to_string(),
256                TokenScope::Admin,
257                NoteGrant {
258                    read: true,
259                    write: false,
260                    notebook_ids: Some(ids)
261                }
262            )
263            .await
264            .is_err());
265        }
266        for (owner, target) in [
267            (2, id.to_string()),
268            (1, uuid::Uuid::new_v4().to_string()),
269            (1, "invalid".into()),
270        ] {
271            assert!(save_token_permissions(
272                &client,
273                owner,
274                &target,
275                TokenScope::Admin,
276                NoteGrant::default()
277            )
278            .await
279            .is_err());
280        }
281        let row = client
282            .query_one(
283                "SELECT scope,notes_write,notebook_ids FROM mcp_tokens WHERE id=$1",
284                &[&id],
285            )
286            .await
287            .unwrap();
288        assert_eq!(row.get::<_, String>(0), "read");
289        assert!(row.get::<_, bool>(1));
290        assert_eq!(row.get::<_, Vec<i32>>(2), vec![10]);
291
292        // 改为全部笔记本，然后收回笔记访问；关闭授权时清除无效的旧范围。
293        save_token_permissions(
294            &client,
295            1,
296            &id.to_string(),
297            TokenScope::Write,
298            NoteGrant {
299                read: true,
300                write: false,
301                notebook_ids: None,
302            },
303        )
304        .await
305        .unwrap();
306        save_token_permissions(
307            &client,
308            1,
309            &id.to_string(),
310            TokenScope::Write,
311            NoteGrant {
312                read: false,
313                write: false,
314                notebook_ids: Some(vec![999]),
315            },
316        )
317        .await
318        .unwrap();
319        let row = client
320            .query_one(
321                "SELECT notes_read,notes_write,notebook_ids FROM mcp_tokens WHERE id=$1",
322                &[&id],
323            )
324            .await
325            .unwrap();
326        assert!(!row.get::<_, bool>(0));
327        assert!(!row.get::<_, bool>(1));
328        assert!(row.get::<_, Option<Vec<i32>>>(2).is_none());
329        for sql in [
330            "UPDATE mcp_tokens SET expires_at=NOW()-INTERVAL '1 second' WHERE id=$1",
331            "UPDATE mcp_tokens SET expires_at=NULL,revoked_at=NOW() WHERE id=$1",
332        ] {
333            client.execute(sql, &[&id]).await.unwrap();
334            assert!(save_token_permissions(
335                &client,
336                1,
337                &id.to_string(),
338                TokenScope::Admin,
339                NoteGrant::default()
340            )
341            .await
342            .is_err());
343            let scope: String = client
344                .query_one("SELECT scope FROM mcp_tokens WHERE id=$1", &[&id])
345                .await
346                .unwrap()
347                .get(0);
348            assert_eq!(scope, "write");
349        }
350        drop(client);
351        connection.await.unwrap().unwrap();
352    }
353}
354
355/// 列出当前管理员名下的全部令牌（不含任何密钥材料，仅展示用元数据）。
356///
357/// 按 `created_at DESC` 排序，最近签发的在前。仅 admin。
358#[server(ListMcpTokens, "/api")]
359pub async fn list_mcp_tokens() -> Result<Vec<McpTokenSummary>, ServerFnError> {
360    #[cfg(feature = "server")]
361    {
362        use crate::api::auth::get_current_admin_user;
363        use crate::api::error::AppError;
364        use crate::db::pool::get_conn;
365
366        let admin = get_current_admin_user().await?;
367        let client = get_conn().await.map_err(AppError::db_conn)?;
368
369        let rows = client
370            .query(
371                "SELECT id::text, user_id, name, scope, created_at, expires_at, \
372                        last_used_at, revoked_at, notes_read, notes_write, notebook_ids \
373                 FROM mcp_tokens \
374                 WHERE user_id = $1 \
375                 ORDER BY created_at DESC",
376                &[&admin.id],
377            )
378            .await
379            .map_err(AppError::query)?;
380
381        Ok(rows
382            .iter()
383            .map(row_to_mcp_token_meta)
384            .map(McpTokenSummary::from)
385            .collect())
386    }
387    #[cfg(not(feature = "server"))]
388    unreachable!()
389}
390
391/// 重新查看令牌明文（可多次调用：明文以密文形式落库，可解密还原）。
392///
393/// 找不到令牌、或令牌不属于当前管理员 → 返回 `None`（不区分原因，避免探测）。
394/// 仅 admin。
395#[server(RevealMcpToken, "/api")]
396pub async fn reveal_mcp_token(id: String) -> Result<Option<String>, ServerFnError> {
397    #[cfg(feature = "server")]
398    {
399        use crate::api::auth::get_current_admin_user;
400        use crate::api::error::AppError;
401        use crate::db::pool::get_conn;
402        use crate::mcp::crypto::decrypt_token;
403
404        let admin = get_current_admin_user().await?;
405        let client = get_conn().await.map_err(AppError::db_conn)?;
406
407        // id 由前端以字符串传入（表 id 列是 uuid）：解析失败视作令牌不存在。
408        let id = match uuid::Uuid::parse_str(&id) {
409            Ok(u) => u,
410            Err(_) => return Ok(None),
411        };
412
413        // 仅取属于当前管理员的令牌的密文，避免越权解密他人令牌。
414        let row = client
415            .query_opt(
416                "SELECT token_enc FROM mcp_tokens WHERE id = $1::uuid AND user_id = $2",
417                &[&id, &admin.id],
418            )
419            .await
420            .map_err(AppError::query)?;
421
422        // 解密失败（密钥缺失/密文被篡改）也归一到 None：调用方无法区分，
423        // 按「该令牌不可解密」处理（等同于失效）。
424        Ok(row
425            .map(|r| r.get::<_, String>("token_enc"))
426            .and_then(|enc| decrypt_token(&enc)))
427    }
428    #[cfg(not(feature = "server"))]
429    unreachable!()
430}
431
432/// 撤销令牌（软删除：置 `revoked_at = now()`，行保留以备审计）。
433///
434/// 找不到或非本人令牌 → 静默无操作（不报错，避免探测）。仅 admin。
435#[server(RevokeMcpToken, "/api")]
436pub async fn revoke_mcp_token(id: String) -> Result<(), ServerFnError> {
437    #[cfg(feature = "server")]
438    {
439        use crate::api::auth::get_current_admin_user;
440        use crate::api::error::AppError;
441        use crate::db::pool::get_conn;
442
443        let admin = get_current_admin_user().await?;
444        let client = get_conn().await.map_err(AppError::db_conn)?;
445
446        // id 解析失败视作令牌不存在（静默无操作，避免探测）。
447        let id = match uuid::Uuid::parse_str(&id) {
448            Ok(u) => u,
449            Err(_) => return Ok(()),
450        };
451
452        client
453            .execute(
454                "UPDATE mcp_tokens SET revoked_at = NOW() \
455                 WHERE id = $1::uuid AND user_id = $2 AND revoked_at IS NULL",
456                &[&id, &admin.id],
457            )
458            .await
459            .map_err(AppError::query)?;
460
461        Ok(())
462    }
463    #[cfg(not(feature = "server"))]
464    unreachable!()
465}
466
467/// 把数据库行解析为令牌元数据（不含明文；密文/哈希已 `#[serde(skip)]`，这里置空）。
468///
469/// `scope` 列存的是字符串；非法值（理论不可能，除非手工改库）按 read 兜底并记日志，
470/// 不 panic。
471#[cfg(feature = "server")]
472fn row_to_mcp_token_meta(row: &tokio_postgres::Row) -> McpToken {
473    let scope_str: String = row.get("scope");
474    let scope = TokenScope::from_db(&scope_str).unwrap_or_else(|| {
475        tracing::warn!(scope = %scope_str, "mcp_tokens.scope 非法值，兜底为 read");
476        TokenScope::Read
477    });
478    McpToken {
479        id: row.get("id"),
480        user_id: row.get("user_id"),
481        name: row.get("name"),
482        scope,
483        notes: crate::models::mcp_token::NoteGrant {
484            read: row.get("notes_read"),
485            write: row.get("notes_write"),
486            notebook_ids: row.get("notebook_ids"),
487        },
488        token_enc: String::new(),
489        token_hash: String::new(),
490        created_at: row.get("created_at"),
491        expires_at: row.get("expires_at"),
492        last_used_at: row.get("last_used_at"),
493        revoked_at: row.get("revoked_at"),
494    }
495}
496
497/// 单个客户端配置片段：标题 + 原始文本（供复制）+ syntect 高亮 HTML（供展示）。
498///
499/// 高亮 HTML 由 `crate::highlight::server::highlight_code` 生成（spaced CSS class 风格，
500/// 配合 `public/highlight.css`）；前端需将其置于 `.md-content pre code` 作用域下，
501/// 否则高亮 CSS 选择器不匹配（见 `src/bin/generate_highlight_css.rs` 的 base 重写）。
502#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
503pub struct McpConfigSnippet {
504    /// 显示标题（含客户端名与目标文件路径）。
505    pub title: String,
506    /// 原始配置文本（供「复制」按钮复制，未高亮）。
507    pub content: String,
508    /// syntect 高亮后的 HTML（`<span>` 序列，无 `<pre>/<code>` 外壳）。
509    pub content_html: String,
510}
511
512/// 各客户端 MCP 配置片段集合。
513///
514/// 由 `get_mcp_client_configs` server fn 返回。`ClientConfigs`（在 `src/mcp/config.rs`）
515/// 是 server-only（`mcp` 模块整体 `#[cfg(feature = "server")]` 门控）；这里把每个片段的
516/// 原始文本与高亮 HTML 打包为可两端共享的 DTO，让 WASM 前端单次请求即可渲染带高亮的配置块。
517#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
518pub struct McpClientConfigs {
519    /// 有序配置片段（标题、原始文本、高亮 HTML）。
520    pub snippets: Vec<McpConfigSnippet>,
521}
522
523/// 根据明文令牌生成各客户端配置片段（含 syntect 高亮 HTML）。
524///
525/// 配置生成与高亮均在服务端完成（`crate::mcp::config` / `crate::highlight` 均为
526/// server-only），返回给前端展示。`APP_BASE_URL` 环境变量也只在服务端读取。仅 admin。
527#[server(GetMcpClientConfigs, "/api")]
528pub async fn get_mcp_client_configs(token: String) -> Result<McpClientConfigs, ServerFnError> {
529    #[cfg(feature = "server")]
530    {
531        use crate::api::auth::get_current_admin_user;
532        use crate::highlight::server::highlight_code;
533
534        let _admin = get_current_admin_user().await?;
535        let c = crate::mcp::config::generate_client_configs(
536            &crate::mcp::config::base_url_from_env(),
537            &token,
538        );
539        // (标题, 内容, 语言)：按客户端配置格式选择语法高亮。
540        let entries: [(&str, String, &str); 8] = [
541            (
542                "Oh-My-Pi（项目根 .mcp.json / ~/.omp/agent/mcp.json 或 ~/.mcp.json）",
543                c.omp_json,
544                "json",
545            ),
546            (
547                "Codex（追加到 ~/.codex/config.toml 或项目根 .codex/config.toml）",
548                c.codex_toml,
549                "toml",
550            ),
551            (
552                "OpenCode（~/.config/opencode/opencode.json 或项目根 opencode.json）",
553                c.opencode_json,
554                "json",
555            ),
556            (
557                "Claude Code（.mcp.json / ~/.claude.json）",
558                c.claude_code_json,
559                "json",
560            ),
561            ("Cursor（~/.cursor/mcp.json）", c.cursor_json, "json"),
562            ("Cline（cline_mcp_settings.json）", c.cline_json, "json"),
563            ("通用（单 server entry）", c.generic_json, "json"),
564            ("Claude Code CLI", c.claude_cli, "bash"),
565        ];
566        let snippets = entries
567            .into_iter()
568            .map(|(title, content, lang)| McpConfigSnippet {
569                title: title.to_string(),
570                content_html: highlight_code(&content, Some(lang)),
571                content,
572            })
573            .collect();
574        Ok(McpClientConfigs { snippets })
575    }
576    #[cfg(not(feature = "server"))]
577    unreachable!()
578}
579
580#[cfg(all(test, feature = "server"))]
581mod tests {
582    use super::*;
583
584    #[test]
585    fn lifetime_expires_at_days() {
586        let now = chrono::Utc::now();
587        let d1 = TokenLifetime::Days1.expires_at().unwrap();
588        let d7 = TokenLifetime::Days7.expires_at().unwrap();
589        assert!(d1 > now);
590        assert!(d7 > d1);
591        // 7 天与 1 天的差应≈6 天（容忍微量时钟漂移）。
592        let delta = (d7 - d1).num_seconds() as f64 / 86400.0;
593        assert!((5.9..6.1).contains(&delta));
594    }
595
596    #[test]
597    fn lifetime_never_is_none() {
598        assert!(TokenLifetime::Never.expires_at().is_none());
599    }
600
601    #[test]
602    fn lifetime_serde_roundtrip() {
603        let json = serde_json::to_string(&TokenLifetime::Days30).unwrap();
604        assert_eq!(json, "\"days30\"");
605        let back: TokenLifetime = serde_json::from_str(&json).unwrap();
606        assert_eq!(back, TokenLifetime::Days30);
607    }
608}
yggdrasil/api/mcp_tokens.rs

yggdrasil/api/
mcp_tokens.rs
