hyperhive/swarm-claude-login/src/main.rs

//! The claude login minter.
//!
//! Connects to the swarm queue as its own OIDC client, takes an operator's
//! login steps from swarm-controller on
//! [`swarm_wire::claude_login::SUBJECT`], runs `claude auth login`
//! for the agent named ([`login`]), and writes the credentials it produces to
//! that agent's store leaf, a new version over any stored before. The agent's
//! harness takes a version newer than the one it last installed.
//!
//! The store identity comes from the `BAO_*` environment. The queue client
//! secret is read at start and again when authelia refuses it; the store is
//! logged in to afresh for each write.

use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result, bail};
use futures_util::StreamExt as _;
use swarm_queue_client::{ClientSecret, QueueConfig};
use swarm_secret_client::claude_login::{Credentials, client_secret_path, credentials_path};
use swarm_secret_client::stored_secret::StoredSecret;
use swarm_wire::claude_login::{Reply, Request, SESSION_DEADLINE, STEP_DEADLINE, SUBJECT};

mod login;

/// The cert-auth role this service logs in to the store as. `swarm-bao.nix`
/// writes it, matching the leaf's subject.
const CERT_ROLE: &str = "swarm-claude-login";

/// The environment prefix `swarm-claude-login.nix` sets.
const PREFIX: &str = "SWARM_CLAUDE_LOGIN";

/// The agent's store leaf, written under this service's certificate.
struct AgentLeaf;

impl login::CredentialSink for AgentLeaf {
    async fn put(&self, agent: &str, credentials: &Credentials) -> Result<()> {
        let path = credentials_path(agent)?;
        let store = swarm_secret_client::SecretStore::from_env(CERT_ROLE)
            .await
            .context("logging in to the secret store")?;
        // `write`, not `create`: a relogin replaces what an earlier one
        // stored.
        store
            .write(&path, credentials)
            .await
            .with_context(|| format!("writing {path}"))
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    hive_log::init();
    swarm_secret_client::install_crypto_provider();
    let var = |name: &str| {
        let full = format!("{PREFIX}_{name}");
        std::env::var(&full).with_context(|| format!("{full} is unset"))
    };
    let client_id = var("OIDC_CLIENT_ID")?;
    let logins = Arc::new(login::Logins::new(
        PathBuf::from(var("CLAUDE")?),
        std::env::temp_dir(),
        STEP_DEADLINE,
    ));
    let secret =
        swarm_queue_client::read_with_startup_retry("reading the queue client secret", || {
            client_secret(&client_id)
        })
        .await?;
    let reread_id = client_id.clone();
    let secret = StoredSecret::new(client_id.clone(), secret, move || {
        let client_id = reread_id.clone();
        async move { client_secret(&client_id).await.map_err(Into::into) }
    });
    let queue = swarm_queue_client::connect(QueueConfig {
        url: var("NATS_URL")?,
        token_endpoint: var("OIDC_TOKEN_ENDPOINT")?,
        client_id,
        client_secret: ClientSecret::Stored(secret),
        ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE"))
            .ok()
            .map(PathBuf::from),
    })
    .await?;
    // No queue group: a running login lives in this process's memory, so
    // every step for it has to reach this process.
    let mut requests = queue
        .subscribe(SUBJECT)
        .await
        .with_context(|| format!("subscribing to {SUBJECT}"))?;
    tracing::info!(subject = SUBJECT, "claude login minter: awaiting logins");

    while let Some(msg) = requests.next().await {
        let Some(reply_to) = msg.reply.clone() else {
            tracing::warn!("a login step arrived with no reply subject; dropping it");
            continue;
        };
        let queue = queue.clone();
        let logins = Arc::clone(&logins);
        tokio::spawn(async move {
            let reply = handle(&msg.payload, &logins).await;
            match serde_json::to_vec(&reply) {
                Ok(body) => {
                    if let Err(e) = queue.publish(reply_to, body.into()).await {
                        tracing::warn!(error = %e, "publishing a login step's reply failed");
                    }
                }
                Err(e) => tracing::warn!(error = %e, "encoding a login step's reply failed"),
            }
        });
    }
    bail!("the subscription to {SUBJECT} ended")
}

/// Decode one step, run it and log it. Never logs the code or the URL.
async fn handle(payload: &[u8], logins: &Arc<login::Logins>) -> Reply {
    let request: Request = match serde_json::from_slice(payload) {
        Ok(r) => r,
        Err(e) => {
            tracing::warn!(error = %e, "an undecodable login step; answering with the error");
            return Reply::Failed {
                error: format!("the minter could not decode the request: {e}"),
            };
        }
    };
    // Before anything else: the name becomes a scratch directory's name and
    // a store path, and this service holds write on every agent's leaf.
    let (Request::Start { agent, .. } | Request::Code { agent, .. } | Request::Cancel { agent }) =
        &request;
    if let Err(reason) = hive_types::Ident::parse(agent) {
        tracing::warn!(agent = ?agent, reason, "a login step for an invalid agent name; refused");
        return Reply::Failed {
            error: format!("{agent:?} is not an agent name: {reason}"),
        };
    }
    let (agent, reply) = match request {
        Request::Start { agent, operator } => {
            tracing::info!(%agent, %operator, "starting a claude login");
            let (reply, id) = logins.start(&agent).await;
            if let Some(id) = id {
                let logins = Arc::clone(logins);
                let expiring = agent.clone();
                tokio::spawn(async move {
                    tokio::time::sleep(SESSION_DEADLINE).await;
                    logins.expire(&expiring, id).await;
                });
            }
            (agent, reply)
        }
        Request::Code { agent, code } => {
            let reply = logins.code(&agent, &code, &AgentLeaf).await;
            (agent, reply)
        }
        Request::Cancel { agent } => {
            let reply = logins.cancel(&agent).await;
            (agent, reply)
        }
    };
    match &reply {
        Reply::Started { .. } => tracing::info!(%agent, "claude login waiting for its code"),
        Reply::Stored => tracing::info!(%agent, "claude credentials stored"),
        Reply::Cancelled => tracing::info!(%agent, "claude login cancelled"),
        Reply::Failed { error } => tracing::warn!(%agent, %error, "claude login step failed"),
    }
    reply
}

async fn client_secret(client_id: &str) -> Result<String> {
    use swarm_secret_client::builder::ClientSecret;

    let path = client_secret_path(client_id)?;
    let store = swarm_secret_client::SecretStore::from_env(CERT_ROLE)
        .await
        .context("logging in to the secret store")?;
    let stored: Option<ClientSecret> = store
        .read_optional(&path)
        .await
        .with_context(|| format!("reading {path}"))?;
    match stored {
        Some(s) if !s.value.trim().is_empty() => Ok(s.value.trim().to_owned()),
        _ => bail!(
            "nothing at {path}: swarm-secret-publish, in authelia's container, has not \
             published this service's client secret"
        ),
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::Duration;

    use swarm_wire::claude_login::{Reply, Request};

    use super::{handle, login};

    #[tokio::test]
    async fn a_step_for_an_invalid_agent_name_is_refused_before_anything_is_made() {
        let root = std::env::temp_dir().join(format!(
            "swarm-claude-login-test-hostile-{}",
            std::process::id()
        ));
        let _ = std::fs::remove_dir_all(&root);
        let scratch = root.join("scratch");
        std::fs::create_dir_all(&scratch).expect("temp dir");
        let logins = Arc::new(login::Logins::new(
            login::tests::stubs().join("claude"),
            scratch.clone(),
            Duration::from_secs(10),
        ));
        for agent in ["../escaped", "..", "", "a/b", "Atlas"] {
            for request in [
                Request::Start {
                    agent: agent.to_owned(),
                    operator: "op".to_owned(),
                },
                Request::Code {
                    agent: agent.to_owned(),
                    code: "c".to_owned(),
                },
                Request::Cancel {
                    agent: agent.to_owned(),
                },
            ] {
                let payload = serde_json::to_vec(&request).expect("encodes");
                let reply = handle(&payload, &logins).await;
                assert!(
                    matches!(reply, Reply::Failed { .. }),
                    "{agent:?}: {reply:?}"
                );
            }
            let made: Vec<_> = std::fs::read_dir(&root)
                .expect("root")
                .map(|e| e.expect("entry").file_name())
                .collect();
            assert_eq!(made, vec!["scratch"], "{agent:?} made a directory");
            assert_eq!(
                std::fs::read_dir(&scratch).expect("scratch").count(),
                0,
                "{agent:?}"
            );
        }
        // The control: a valid name starts.
        let payload = serde_json::to_vec(&Request::Start {
            agent: "atlas".to_owned(),
            operator: "op".to_owned(),
        })
        .expect("encodes");
        assert!(matches!(
            handle(&payload, &logins).await,
            Reply::Started { .. }
        ));
        drop(logins);
        let _ = std::fs::remove_dir_all(&root);
    }
}