A Google Drive sync renews its access token, a disconnect spares a grant still in use, and only the connector reconnects

crates/utopia-server/src/api/google_drive_routes.rs

@@ -76,6 +76,36 @@ fn require_connector(source: &Source, user_id: Uuid) -> Result<(), AppError> {
     })
 }
 
+/// 别人连着的来源不能直接重连：一点就会换掉他的 Drive、丢掉他选的文件夹。
+/// 要换人，先断开（断开会问一句，说清是谁的 Drive），再由新的人来连
+fn refuse_someone_elses_connection(source: &Source, user_id: Uuid) -> Result<(), AppError> {
+    match google_drive::connected_by(&source.config) {
+        Some(connector) if connector != user_id => Err(AppError::CodedForbidden {
+            code: "google_drive_connected_by_someone_else",
+            message: "This source reads someone else's Google Drive; disconnect it first".into(),
+        }),
+        _ => Ok(()),
+    }
+}
+
+/// 同一个 Google 账号是不是还连着别的来源。Google 撤销一枚刷新令牌，撤的是这个人对
+/// 我们这个客户端的整份授权，他别的 Drive 来源会一起失效；还有别的在用就不撤
+pub(super) async fn grant_still_used(
+    pool: &sqlx::PgPool,
+    source_id: Uuid,
+    account_email: &str,
+) -> Result<bool, sqlx::Error> {
+    sqlx::query_scalar(
+        "SELECT EXISTS (SELECT 1 FROM sources
+          WHERE kind = 'google_drive' AND id <> $1
+            AND config->>'account_email' = $2 AND config ? 'refresh_token')",
+    )
+    .bind(source_id)
+    .bind(account_email)
+    .fetch_one(pool)
+    .await
+}
+
 /// 落一次连接流程。总数与这个人的在途数都在同一条插入里判，满了就不插，回 `None`
 pub(super) async fn record_flow(
     pool: &sqlx::PgPool,
@@ -117,7 +147,8 @@ pub async fn connect(
     Path((kb_id, source_id)): Path<(Uuid, Uuid)>,
 ) -> ApiResult<Json<Value>> {
     require_kb(&state, &user, kb_id, Role::Editor).await?;
-    drive_source(&state, kb_id, source_id).await?;
+    let source = drive_source(&state, kb_id, source_id).await?;
+    refuse_someone_elses_connection(&source, user.id)?;
     let cfg = configured()?;
     let (flow_state, verifier) = record_flow(&state.pool, user.id, source_id)
         .await?
@@ -176,6 +207,7 @@ pub async fn callback(
                     "session" => "This window is not signed in to Utopia as the person who started the connection. Close it and start again from the source settings.",
                     "expired" => "This connection attempt has expired. Close this window and start again.",
                     "scope" => "Google did not grant read access to Drive. Close this window, connect again and allow it.",
+                    "taken" => "Someone else connected this source while you were on Google's page. Close this window.",
                     _ => "Google could not be reached or refused the request. Close this window and try again.",
                 },
             )
@@ -225,6 +257,13 @@ pub(super) async fn finish(
                 "user is no longer an editor of the knowledge base".to_string(),
             )
         })?;
+    // 同意页上的这几分钟里，别人可能先连上了
+    refuse_someone_elses_connection(&source, user.id).map_err(|_| {
+        (
+            "taken",
+            "someone else connected this source in the meantime".to_string(),
+        )
+    })?;
 
     let cfg = google_drive::oauth_config().ok_or(("unavailable", "not configured".to_string()))?;
     let grant = google_drive::exchange_code(&cfg, &code, &verifier)
@@ -241,20 +280,14 @@ pub(super) async fn finish(
             )
         })?;
 
-    // 换了一个人连：上一个人选的文件夹是照他的 Drive 选的，不留给新的人
-    let drop_folders: &[&str] = if google_drive::connected_by(&source.config) == Some(user.id) {
-        &[]
-    } else {
-        &["folders"]
-    };
     let grant_keys = json!({
         "refresh_token": grant.refresh_token,
         "account_email": grant.account_email,
         // 列表里看不见令牌，界面凭这一项知道「连上了」：没验证过的邮箱不显示，不能拿它当标志
         "connected_at": chrono::Utc::now().to_rfc3339(),
         "connected_by": user.id,
     });
-    utopia_store::sources::patch_config(&state.pool, source_id, &grant_keys, drop_folders)
+    utopia_store::sources::patch_config(&state.pool, source_id, &grant_keys, &[])
         .await
         .map_err(|e| ("unavailable", format!("saving the grant: {e}")))?;
     state.emit_source(source.kb_id);
@@ -272,7 +305,8 @@ pub(super) async fn finish(
     Ok(())
 }
 
-/// 断开：删掉刷新令牌，并请 Google 把它作废。作废失败不拦着——库里已经没有它了
+/// 断开：删掉刷新令牌，并请 Google 把它作废——除非同一个账号还连着别的来源
+///（[`grant_still_used`]），或者不知道是哪个账号。作废失败不拦着：库里已经没有它了
 pub async fn disconnect(
     State(state): State<AppState>,
     AuthUser(user): AuthUser,
@@ -281,17 +315,30 @@ pub async fn disconnect(
     require_kb(&state, &user, kb_id, Role::Editor).await?;
     let source = drive_source(&state, kb_id, source_id).await?;
     let token = source.config["refresh_token"].as_str().map(String::from);
+    let account = source.config["account_email"].as_str().map(String::from);
     utopia_store::sources::patch_config(&state.pool, source_id, &json!({}), CONNECTION_KEYS)
         .await?;
-    if let Some(token) = token {
+    let revoke = match (token, account) {
+        (Some(token), Some(account))
+            if !grant_still_used(&state.pool, source_id, &account).await? =>
+        {
+            Some(token)
+        }
+        _ => None,
+    };
+    if let Some(token) = revoke {
         let revoked = reqwest::Client::new()
             .post("https://oauth2.googleapis.com/revoke")
             .form(&[("token", token.as_str())])
             .timeout(std::time::Duration::from_secs(10))
             .send()
             .await;
-        if let Err(e) = revoked {
-            tracing::warn!(error = %e, "Google 刷新令牌作废失败，库里已删除");
+        match revoked {
+            Ok(resp) if resp.status().is_success() => {}
+            Ok(resp) => {
+                tracing::warn!(status = %resp.status(), "Google 拒绝作废刷新令牌，库里已删除")
+            }
+            Err(e) => tracing::warn!(error = %e, "Google 刷新令牌作废失败，库里已删除"),
         }
     }
     state.emit_source(kb_id);

crates/utopia-server/src/api/google_drive_routes_tests.rs

@@ -118,7 +118,9 @@ impl Fixture {
 
 fn denied_code<T>(result: ApiResult<T>) -> Option<&'static str> {
     match result {
-        Err(ApiErr(AppError::CodedForbidden { code, .. })) => Some(code),
+        Err(ApiErr(AppError::CodedForbidden { code, .. } | AppError::Invalid { code, .. })) => {
+            Some(code)
+        }
         _ => None,
     }
 }
@@ -291,3 +293,91 @@ async fn a_sync_stops_once_the_connected_person_is_no_longer_an_editor() -> anyh
     );
     Ok(())
 }
+
+#[tokio::test]
+async fn a_colleague_cannot_reconnect_the_connected_persons_drive() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    let started = connect(
+        State(fx.state.clone()),
+        AuthUser(fx.colleague.clone()),
+        Path((fx.kb, fx.source)),
+    )
+    .await;
+    assert_eq!(
+        denied_code(started),
+        Some("google_drive_connected_by_someone_else")
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn a_callback_after_someone_else_connected_is_refused() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    let (flow_state, _) = record_flow(&fx.pool, fx.colleague.id, fx.source)
+        .await?
+        .expect("room for a flow");
+    assert_eq!(fx.callback(Some(&fx.colleague), &flow_state).await, "taken");
+    Ok(())
+}
+
+#[tokio::test]
+async fn a_grant_is_kept_while_another_source_uses_the_same_account() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    let account = format!("{}@drive.test", Uuid::now_v7());
+    let other = utopia_store::sources::create(
+        &fx.pool,
+        fx.kb,
+        "google_drive",
+        "Other Drive",
+        &json!({}),
+        None,
+        None,
+        None,
+    )
+    .await?
+    .id;
+    for id in [fx.source, other] {
+        let grant = json!({ "refresh_token": "1//grant", "account_email": account });
+        utopia_store::sources::patch_config(&fx.pool, id, &grant, &[]).await?;
+    }
+    let while_other_connected = grant_still_used(&fx.pool, fx.source, &account).await?;
+    utopia_store::sources::patch_config(&fx.pool, other, &json!({}), CONNECTION_KEYS).await?;
+    let after_other_disconnected = grant_still_used(&fx.pool, fx.source, &account).await?;
+    assert_eq!(
+        (while_other_connected, after_other_disconnected),
+        (true, false)
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn a_drive_source_cannot_be_created_without_an_oauth_client() -> anyhow::Result<()> {
+    let Some(fx) = Fixture::new().await? else {
+        return Ok(());
+    };
+    let created = super::super::sources_routes::create(
+        State(fx.state.clone()),
+        crate::auth::SessionOrToken {
+            user: fx.connector.clone(),
+            pat: None,
+        },
+        Path(fx.kb),
+        Json(super::super::sources_routes::CreateBody {
+            kind: "google_drive".into(),
+            name: "Drive".into(),
+            config: json!({}),
+            icon: None,
+            sync_interval_minutes: None,
+            sync_cron: None,
+        }),
+    )
+    .await;
+    assert_eq!(denied_code(created), Some("google_drive_not_configured"));
+    Ok(())
+}

crates/utopia-server/src/api/sources_routes.rs

@@ -109,11 +109,16 @@ pub async fn create(
     let user = caller.user;
     require_kb(&state, &user, kb_id, Role::Editor).await?;
     validate_rss_config(&body.kind, &body.config)?;
-    let config = if config_is_drive_owned(&body.kind) {
-        json!({})
-    } else {
-        body.config
-    };
+    let drive = config_is_drive_owned(&body.kind);
+    // 没配 OAuth 客户端的部署连不上 Google，建出来也只会一直失败（0074 决定 2）
+    if drive && crate::google_drive::oauth_config().is_none() {
+        return Err(utopia_core::AppError::invalid(
+            "google_drive_not_configured",
+            "Google Drive is not configured on this server",
+        )
+        .into());
+    }
+    let config = if drive { json!({}) } else { body.config };
     let source = utopia_store::sources::create(
         &state.pool,
         kb_id,

crates/utopia-server/src/google_drive.rs

@@ -75,6 +75,7 @@ const MAX_RETRIES: u32 = 3;
 const JSON_LIMIT: usize = 4 * 1024 * 1024;
 
 /// 部署级的 OAuth 客户端。三项都得有，缺一项就当没配
+#[derive(Clone)]
 pub struct OAuthConfig {
     pub client_id: String,
     pub client_secret: String,
@@ -350,24 +351,38 @@ fn export_as(mime: &str) -> Option<(&'static str, &'static str)> {
 /// 拿着一个访问令牌的 Drive 客户端
 pub struct Drive {
     http: reqwest::Client,
-    token: String,
+    token: std::sync::Mutex<String>,
+    /// 访问令牌一小时就过期，一轮同步可以更长：遇到 401 凭刷新令牌换一枚新的，再试一次
+    renewal: Option<Renewal>,
     // 私有地址入口让 HTTP 测试走真实的遍历逻辑，不增加用户配置
     api_root: String,
 }
 
+struct Renewal {
+    cfg: OAuthConfig,
+    refresh_token: String,
+    token_url: String,
+}
+
 impl Drive {
     /// 用刷新令牌换访问令牌，得到一个能用的客户端
     pub async fn connect(cfg: &OAuthConfig, refresh_token: &str) -> anyhow::Result<Self> {
         let http = http()?;
         let token = access_token_at(&http, TOKEN_URL, cfg, refresh_token).await?;
         Ok(Self {
             http,
-            token,
+            token: std::sync::Mutex::new(token),
+            renewal: Some(Renewal {
+                cfg: cfg.clone(),
+                refresh_token: refresh_token.into(),
+                token_url: TOKEN_URL.into(),
+            }),
             api_root: API_ROOT.into(),
         })
     }
 
-    /// 发一次 GET。429 与 5xx 退避重试；其余非 2xx 带上 Google 自己的 message 报错，
+    /// 发一次 GET。429 与 5xx 退避重试；401 换一枚访问令牌再试一次（只一次：换来的也被拒，
+    /// 就是授权本身出了问题）；其余非 2xx 带上 Google 自己的 message 报错，
     /// 状态码与 `reason` 留在错误里给调用方分辨（404 = 打不开）
     async fn get(
         &self,
@@ -378,15 +393,24 @@ impl Drive {
         let url = reqwest::Url::parse_with_params(url, query)
             .with_context(|| format!("google drive {what}: building the URL"))?;
         let mut attempt = 0u32;
+        let mut renewed = false;
         loop {
+            let token = self.token.lock().expect("token lock").clone();
             let resp = self
                 .http
                 .get(url.clone())
-                .bearer_auth(&self.token)
+                .bearer_auth(&token)
                 .send()
                 .await
                 .with_context(|| format!("google drive {what}"))?;
             let status = resp.status();
+            if let (StatusCode::UNAUTHORIZED, false, Some(r)) = (status, renewed, &self.renewal) {
+                renewed = true;
+                let fresh =
+                    access_token_at(&self.http, &r.token_url, &r.cfg, &r.refresh_token).await?;
+                *self.token.lock().expect("token lock") = fresh;
+                continue;
+            }
             if (status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error())
                 && attempt < MAX_RETRIES
             {