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#[derive(Debug, Clone, PartialEq, Eq)]
156pub struct DeviceCode {
157    /// URL the user opens to approve the login.
158    pub verification_url: String,
159    /// Short confirmation code the user checks against the browser page.
160    pub code: String,
161}
162
163/// An in-flight `muse login` OAuth device-code flow.
164///
165/// Plain stdout, no pseudo-terminal: the CLI prints the verification URL
166/// and code, then blocks polling for browser approval. Dropping the flow
167/// cancels the login (the child is killed).
168pub struct DeviceLoginFlow {
169    child: Child,
170    lines: Lines<BufReader<ChildStdout>>,
171}
172
173impl DeviceLoginFlow {
174    /// Spawn `muse login` from `PATH`.
175    pub async fn start() -> Result<Self> {
176        Self::start_with_binary("muse").await
177    }
178
179    /// [`start`](Self::start) with a specific CLI binary.
180    pub async fn start_with_binary(binary: &str) -> Result<Self> {
181        let mut cmd = tokio::process::Command::new(resolve(binary)?);
182        cmd.arg("login")
183            .stdin(Stdio::null())
184            .stdout(Stdio::piped())
185            .stderr(Stdio::piped())
186            .kill_on_drop(true);
187        let mut child = cmd.spawn()?;
188        let stdout = child
189            .stdout
190            .take()
191            .ok_or_else(|| Error::Protocol("failed to get stdout".to_string()))?;
192        Ok(Self {
193            child,
194            lines: BufReader::new(stdout).lines(),
195        })
196    }
197
198    /// Read output until the verification URL and code are both seen.
199    ///
200    /// Show them to your user; the flow keeps polling in the background
201    /// until [`wait_approved`](Self::wait_approved) resolves.
202    pub async fn device_code(&mut self, timeout: Duration) -> Result<DeviceCode> {
203        let read = async {
204            let mut url: Option<String> = None;
205            let mut code: Option<String> = None;
206            while let Some(line) = self.lines.next_line().await? {
207                if let Some(u) = extract_url(&line) {
208                    url = Some(u);
209                }
210                if let Some(c) = extract_code_from_url_or_line(&line) {
211                    code = Some(c);
212                }
213                if let (Some(u), Some(c)) = (&url, &code) {
214                    return Ok(DeviceCode {
215                        verification_url: u.clone(),
216                        code: c.clone(),
217                    });
218                }
219            }
220            Err(Error::Protocol(
221                "muse login ended before printing a device code".to_string(),
222            ))
223        };
224        tokio::time::timeout(timeout, read)
225            .await
226            .map_err(|_| Error::Protocol("timed out waiting for device code".to_string()))?
227    }
228
229    /// Wait for the user to approve in the browser: resolves when the CLI
230    /// exits successfully (credential saved) or errors on failure/timeout.
231    pub async fn wait_approved(mut self, timeout: Duration) -> Result<()> {
232        let status = tokio::time::timeout(timeout, self.child.wait())
233            .await
234            .map_err(|_| Error::Protocol("timed out waiting for login approval".to_string()))??;
235        if status.success() {
236            Ok(())
237        } else {
238            Err(Error::Protocol(format!(
239                "muse login exited with {:?} before approval",
240                status.code()
241            )))
242        }
243    }
244
245    /// Cancel the login.
246    pub async fn cancel(mut self) -> Result<()> {
247        self.child.kill().await?;
248        Ok(())
249    }
250}
251
252/// First `https://` run in a line (device URLs carry no trailing prose on
253/// the observed wire, but trim conservatively anyway).
254fn extract_url(line: &str) -> Option<String> {
255    let start = line.find("https://")?;
256    let url: String = line[start..]
257        .chars()
258        .take_while(|c| !c.is_whitespace())
259        .collect();
260    Some(url)
261}
262
263/// The device code: from the URL's `code=` parameter, or a bare
264/// `XXXX-XXXX`-shaped token on its own line (the CLI prints both forms).
265fn extract_code_from_url_or_line(line: &str) -> Option<String> {
266    if let Some(pos) = line.find("code=") {
267        let code: String = line[pos + 5..]
268            .chars()
269            .take_while(|c| c.is_ascii_alphanumeric() || *c == '-')
270            .collect();
271        if !code.is_empty() {
272            return Some(code);
273        }
274    }
275    let t = line.trim();
276    let is_code_shaped = t.len() >= 7
277        && t.len() <= 12
278        && t.chars()
279            .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == '-')
280        && t.contains('-');
281    is_code_shaped.then(|| t.to_string())
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    /// Verbatim `muse login` output captured from Muse Code 0.1.0.
289    const CAPTURE: &str = "Open this page to sign in:\n  \
290        https://auth.meta.com/oauth/device/?code=TBSS-QJWM\n\
291        confirm this code matches:\n  TBSS-QJWM\n\nWaiting for approval…\n";
292
293    #[test]
294    fn device_code_extracted_from_captured_output() {
295        let mut url = None;
296        let mut code = None;
297        for line in CAPTURE.lines() {
298            if let Some(u) = extract_url(line) {
299                url = Some(u);
300            }
301            if let Some(c) = extract_code_from_url_or_line(line) {
302                code = Some(c);
303            }
304        }
305        assert_eq!(
306            url.as_deref(),
307            Some("https://auth.meta.com/oauth/device/?code=TBSS-QJWM")
308        );
309        assert_eq!(code.as_deref(), Some("TBSS-QJWM"));
310    }
311
312    #[test]
313    fn bare_code_line_matches_and_prose_does_not() {
314        assert_eq!(
315            extract_code_from_url_or_line("  TBSS-QJWM"),
316            Some("TBSS-QJWM".to_string())
317        );
318        assert_eq!(extract_code_from_url_or_line("Waiting for approval…"), None);
319        assert_eq!(
320            extract_code_from_url_or_line("Open this page to sign in:"),
321            None
322        );
323    }
324}