Skip to main content

muse_codes/
auth.rs

1//! Login support tooling for Muse Code (feature `async-client`).
2//!
3//! Muse's auth surface is automation-friendly — no TUI to drive:
4//!
5//! - [`auth_set`] wraps `muse auth set --api-key-stdin` (the key travels
6//!   over stdin; Muse refuses to take secrets as arguments).
7//! - [`DeviceLoginFlow`] wraps `muse login`, a plain-stdout OAuth
8//!   device-code flow: the CLI prints a verification URL and a short code,
9//!   then polls until the user approves in a browser.
10//! - [`logout`] wraps `muse logout` (removes the saved credential;
11//!   `META_API_KEY` in the environment is never touched).
12//! - [`credentials_present`] reports whether a run could authenticate
13//!   right now (env key or saved credential file).
14//!
15//! Credential resolution order (per the CLI): `META_API_KEY` env always
16//! wins, then the saved credential at `~/.config/muse/auth.json`.
17
18use crate::error::{Error, Result};
19use std::path::PathBuf;
20use std::process::Stdio;
21use std::time::Duration;
22use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, Lines};
23use tokio::process::{Child, ChildStdout};
24
25/// Environment variable that overrides any saved credential.
26pub const META_API_KEY_VAR: &str = "META_API_KEY";
27
28/// Path of the saved credential file (`~/.config/muse/auth.json`),
29/// honoring `XDG_CONFIG_HOME`. `None` when no home directory resolves.
30pub fn credentials_path() -> Option<PathBuf> {
31    if let Some(xdg) = std::env::var_os("XDG_CONFIG_HOME") {
32        return Some(PathBuf::from(xdg).join("muse/auth.json"));
33    }
34    let home = std::env::var_os("HOME").or_else(|| std::env::var_os("USERPROFILE"))?;
35    Some(PathBuf::from(home).join(".config/muse/auth.json"))
36}
37
38/// True when a headless run could authenticate right now: `META_API_KEY`
39/// is set (non-empty) or the saved credential file carries at least one
40/// provider. (`muse logout` empties the providers map but keeps the file,
41/// so bare existence is not enough.)
42pub fn credentials_present() -> bool {
43    if std::env::var(META_API_KEY_VAR).map(|v| !v.trim().is_empty()) == Ok(true) {
44        return true;
45    }
46    let Some(path) = credentials_path() else {
47        return false;
48    };
49    match std::fs::read_to_string(&path) {
50        Ok(raw) => serde_json::from_str::<AuthFile>(&raw)
51            .map(|f| !f.providers.is_empty())
52            // Unparseable file: assume it authenticates (newer schema).
53            .unwrap_or(true),
54        Err(_) => false,
55    }
56}
57
58/// Shape of `~/.config/muse/auth.json` (observed schema_version 1).
59///
60/// `muse logout` rewrites this with an empty `providers` map rather than
61/// deleting the file.
62#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
63pub struct AuthFile {
64    pub schema_version: u32,
65    pub providers: std::collections::BTreeMap<String, ProviderCredential>,
66    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
67    pub extra: serde_json::Map<String, serde_json::Value>,
68}
69
70/// One saved provider credential.
71#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
72pub struct ProviderCredential {
73    #[serde(default, skip_serializing_if = "Option::is_none")]
74    pub api_key: Option<String>,
75    #[serde(flatten, default, skip_serializing_if = "serde_json::Map::is_empty")]
76    pub extra: serde_json::Map<String, serde_json::Value>,
77}
78
79fn resolve(binary: &str) -> Result<PathBuf> {
80    which::which(binary).map_err(|_| Error::BinaryNotFound {
81        name: binary.to_string(),
82    })
83}
84
85/// Save a provider API key: `muse auth set --provider <p> --api-key-stdin`.
86///
87/// The key is written to the child's stdin and never appears on a command
88/// line. `provider` defaults to `"meta"` when `None`.
89pub async fn auth_set(api_key: &str, provider: Option<&str>) -> Result<()> {
90    auth_set_with_binary("muse", api_key, provider).await
91}
92
93/// [`auth_set`] against a specific CLI binary.
94pub async fn auth_set_with_binary(
95    binary: &str,
96    api_key: &str,
97    provider: Option<&str>,
98) -> Result<()> {
99    let mut cmd = tokio::process::Command::new(resolve(binary)?);
100    cmd.args([
101        "auth",
102        "set",
103        "--provider",
104        provider.unwrap_or("meta"),
105        "--api-key-stdin",
106    ])
107    .stdin(Stdio::piped())
108    .stdout(Stdio::piped())
109    .stderr(Stdio::piped())
110    .kill_on_drop(true);
111    let mut child = cmd.spawn()?;
112    let mut stdin = child
113        .stdin
114        .take()
115        .ok_or_else(|| Error::Protocol("failed to get stdin".to_string()))?;
116    stdin.write_all(api_key.trim().as_bytes()).await?;
117    stdin.write_all(b"\n").await?;
118    drop(stdin); // EOF tells the CLI the key is complete.
119    let out = child.wait_with_output().await?;
120    if out.status.success() {
121        Ok(())
122    } else {
123        Err(Error::Protocol(format!(
124            "muse auth set failed (exit {:?}): {}",
125            out.status.code(),
126            String::from_utf8_lossy(&out.stderr).trim()
127        )))
128    }
129}
130
131/// Remove the saved credential: `muse logout`.
132pub async fn logout() -> Result<()> {
133    logout_with_binary("muse").await
134}
135
136/// [`logout`] against a specific CLI binary.
137pub async fn logout_with_binary(binary: &str) -> Result<()> {
138    let out = tokio::process::Command::new(resolve(binary)?)
139        .arg("logout")
140        .stdin(Stdio::null())
141        .output()
142        .await?;
143    if out.status.success() {
144        Ok(())
145    } else {
146        Err(Error::Protocol(format!(
147            "muse logout failed (exit {:?}): {}",
148            out.status.code(),
149            String::from_utf8_lossy(&out.stderr).trim()
150        )))
151    }
152}
153
154/// The verification details a [`DeviceLoginFlow`] presents to the user.
155/// Serde-serializable for relay to remote UIs.
156#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
157pub struct DeviceCode {
158    /// URL the user opens to approve the login.
159    pub verification_url: String,
160    /// Short confirmation code the user checks against the browser page.
161    pub code: String,
162}
163
164/// An in-flight `muse login` OAuth device-code flow.
165///
166/// Plain stdout, no pseudo-terminal: the CLI prints the verification URL
167/// and code, then blocks polling for browser approval. Dropping the flow
168/// cancels the login (the child is killed).
169pub struct DeviceLoginFlow {
170    child: Child,
171    lines: Lines<BufReader<ChildStdout>>,
172}
173
174impl DeviceLoginFlow {
175    /// Spawn `muse login` from `PATH`.
176    pub async fn start() -> Result<Self> {
177        Self::start_with_binary("muse").await
178    }
179
180    /// [`start`](Self::start) with a specific CLI binary.
181    pub async fn start_with_binary(binary: &str) -> Result<Self> {
182        let mut cmd = tokio::process::Command::new(resolve(binary)?);
183        cmd.arg("login")
184            .stdin(Stdio::null())
185            .stdout(Stdio::piped())
186            .stderr(Stdio::piped())
187            .kill_on_drop(true);
188        let mut child = cmd.spawn()?;
189        let stdout = child
190            .stdout
191            .take()
192            .ok_or_else(|| Error::Protocol("failed to get stdout".to_string()))?;
193        Ok(Self {
194            child,
195            lines: BufReader::new(stdout).lines(),
196        })
197    }
198
199    /// Read output until the verification URL and code are both seen.
200    ///
201    /// Show them to your user; the flow keeps polling in the background
202    /// until [`wait_approved`](Self::wait_approved) resolves.
203    pub async fn device_code(&mut self, timeout: Duration) -> Result<DeviceCode> {
204        let read = async {
205            let mut url: Option<String> = None;
206            let mut code: Option<String> = None;
207            while let Some(line) = self.lines.next_line().await? {
208                if let Some(u) = extract_url(&line) {
209                    url = Some(u);
210                }
211                if let Some(c) = extract_code_from_url_or_line(&line) {
212                    code = Some(c);
213                }
214                if let (Some(u), Some(c)) = (&url, &code) {
215                    return Ok(DeviceCode {
216                        verification_url: u.clone(),
217                        code: c.clone(),
218                    });
219                }
220            }
221            Err(Error::Protocol(
222                "muse login ended before printing a device code".to_string(),
223            ))
224        };
225        tokio::time::timeout(timeout, read)
226            .await
227            .map_err(|_| Error::Protocol("timed out waiting for device code".to_string()))?
228    }
229
230    /// Wait for the user to approve in the browser: resolves when the CLI
231    /// exits successfully (credential saved) or errors on failure/timeout.
232    pub async fn wait_approved(mut self, timeout: Duration) -> Result<()> {
233        let status = tokio::time::timeout(timeout, self.child.wait())
234            .await
235            .map_err(|_| Error::Protocol("timed out waiting for login approval".to_string()))??;
236        if status.success() {
237            Ok(())
238        } else {
239            Err(Error::Protocol(format!(
240                "muse login exited with {:?} before approval",
241                status.code()
242            )))
243        }
244    }
245
246    /// Cancel the login.
247    pub async fn cancel(mut self) -> Result<()> {
248        self.child.kill().await?;
249        Ok(())
250    }
251}
252
253/// First `https://` run in a line (device URLs carry no trailing prose on
254/// the observed wire, but trim conservatively anyway).
255fn extract_url(line: &str) -> Option<String> {
256    let start = line.find("https://")?;
257    let url: String = line[start..]
258        .chars()
259        .take_while(|c| !c.is_whitespace())
260        .collect();
261    Some(url)
262}
263
264/// The device code: from the URL's `code=` parameter, or a bare
265/// `XXXX-XXXX`-shaped token on its own line (the CLI prints both forms).
266fn extract_code_from_url_or_line(line: &str) -> Option<String> {
267    if let Some(pos) = line.find("code=") {
268        let code: String = line[pos + 5..]
269            .chars()
270            .take_while(|c| c.is_ascii_alphanumeric() || *c == '-')
271            .collect();
272        if !code.is_empty() {
273            return Some(code);
274        }
275    }
276    let t = line.trim();
277    let is_code_shaped = t.len() >= 7
278        && t.len() <= 12
279        && t.chars()
280            .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == '-')
281        && t.contains('-');
282    is_code_shaped.then(|| t.to_string())
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    /// Verbatim `muse login` output captured from Muse Code 0.1.0.
290    const CAPTURE: &str = "Open this page to sign in:\n  \
291        https://auth.meta.com/oauth/device/?code=TBSS-QJWM\n\
292        confirm this code matches:\n  TBSS-QJWM\n\nWaiting for approval…\n";
293
294    #[test]
295    fn device_code_extracted_from_captured_output() {
296        let mut url = None;
297        let mut code = None;
298        for line in CAPTURE.lines() {
299            if let Some(u) = extract_url(line) {
300                url = Some(u);
301            }
302            if let Some(c) = extract_code_from_url_or_line(line) {
303                code = Some(c);
304            }
305        }
306        assert_eq!(
307            url.as_deref(),
308            Some("https://auth.meta.com/oauth/device/?code=TBSS-QJWM")
309        );
310        assert_eq!(code.as_deref(), Some("TBSS-QJWM"));
311    }
312
313    #[test]
314    fn bare_code_line_matches_and_prose_does_not() {
315        assert_eq!(
316            extract_code_from_url_or_line("  TBSS-QJWM"),
317            Some("TBSS-QJWM".to_string())
318        );
319        assert_eq!(extract_code_from_url_or_line("Waiting for approval…"), None);
320        assert_eq!(
321            extract_code_from_url_or_line("Open this page to sign in:"),
322            None
323        );
324    }
325}