drep/

auth.rs

//! The credential store: API keys drep holds on the user's behalf.
//!
//! `drep.toml` is a repository file. It names an endpoint, a model and a
//! protocol, all of which are shareable, and it is meant to survive being
//! committed - which is why `api_key = "${VAR}"` names an environment variable
//! rather than holding a secret. That indirection is the right answer for CI,
//! where the key arrives as a secret and nobody is at a keyboard, and the wrong
//! answer for a person setting drep up on their laptop: it makes the first-run
//! experience "now go and export something into the right shell profile".
//!
//! So keys live here instead, once per machine, outside any repository:
//!
//! ```text
//! ~/.config/drep/auth.toml     (macOS: ~/Library/Application Support/dev.slb350.drep)
//! ```
//!
//! ## Keyed by endpoint, not by provider name
//!
//! A key authenticates a *host*, so the endpoint is what it belongs to. Keying
//! by a preset name instead would mean a config that named no preset - a custom
//! endpoint, or one edited by hand after `drep init` - could not find its own
//! credential, and two presets pointed at the same host would each need their
//! own copy of one key.
//!
//! The endpoint is normalised before use (see [`normalise`]) so a trailing
//! slash or a difference in case does not hide a key from the config that
//! stored it.
//!
//! There is deliberately no `load_default`/`save_default` pair. Both were thin
//! wrappers over `default_path()`, which reads the environment - so nothing
//! could test them without `std::env::set_var`, and the mutation gate found
//! them undetectable. Every caller resolves the path once, at its own entry
//! point, and passes it down; that is also what keeps tests off the real store.
//!
//! ## Resolution order
//!
//! [`resolve`] fills in what `drep.toml` left unset, in this order:
//!
//! 1. an explicit `api_key` - a literal, or a `${VAR}` the loader has already
//!    substituted;
//! 2. `api_key_command`, an argv drep runs to mint one;
//! 3. a key held here for the same endpoint;
//! 4. nothing, which `LlmClient::new` turns into an empty key so the SDK sends
//!    no protocol authentication header.
//!
//! An explicit value in the file always wins. A user who writes
//! `api_key = "${OPENROUTER_API_KEY}"` has said where the key comes from, and
//! silently preferring a stored one would make the file lie about what the run
//! used. The command sits above the store for the same reason: a file naming a
//! command has not left the question unanswered.

use std::borrow::Cow;
use std::collections::{BTreeMap, BTreeSet};
use std::path::PathBuf;

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::config::Config;

mod command;
mod store_file;
pub use command::KeyCommandError;
pub(crate) use store_file::ensure_dir_private;
#[cfg(test)]
pub(crate) use store_file::{restrict, temporary_parent};

/// Where a provider's key came from, for [`doctor`](crate::cli::doctor) to report.
///
/// The point of naming the source is that "it works on my machine" and "it works
/// in CI" are different configurations, and the difference is invisible in
/// `drep.toml` once a stored key exists.
///
/// The variant order is the resolution order.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum KeySource {
    /// The config named an environment variable or a literal, and it resolved.
    Config,
    /// The config named an `api_key_command`, which drep runs to mint a key.
    Command,
    /// The config named nothing and the store had one for this endpoint.
    Store,
    /// Neither. The provider sends no protocol authentication header; a custom
    /// header may still provide the endpoint's complete authentication scheme.
    Missing,
}

impl KeySource {
    /// Every source, in resolution order.
    ///
    /// Exists so the label-collision test cannot pass on a stale subset, the way
    /// `Severity::ALL` drives its parser test: a variant added without wording of
    /// its own has to fail something.
    pub const ALL: [Self; 4] = [Self::Config, Self::Command, Self::Store, Self::Missing];

    /// What `doctor` prints for this source.
    ///
    /// The single definition of the wording. `doctor` hand-wrote its own copy
    /// of these three strings while this method sat unused, which is two places
    /// to keep in step for no gain.
    pub fn label(&self) -> &'static str {
        match self {
            Self::Config => "from drep.toml",
            Self::Command => "from api_key_command",
            Self::Store => "from the drep auth store",
            Self::Missing => "not set - run `drep auth login` or add `api_key` to drep.toml",
        }
    }
}

/// Keys held for this machine, keyed by normalised endpoint.
///
/// `Serialize`/`Deserialize` drive the on-disk TOML directly; there is no
/// separate wire type because the file is drep's own and has one shape.
#[derive(Default, Serialize, Deserialize)]
pub struct AuthStore {
    /// Endpoint to key. A `BTreeMap` so the file is written in a stable order
    /// and a re-save produces no spurious diff.
    #[serde(default)]
    keys: BTreeMap<String, String>,

    /// Mutations made since this snapshot was loaded.
    ///
    /// Saving replays only these operations over the latest on-disk store, so
    /// two drep processes changing different endpoints cannot overwrite one
    /// another. The journal is process state, never part of the TOML file.
    #[serde(skip)]
    pending: std::sync::Mutex<BTreeSet<String>>,
}

/// Hand-written so a key cannot reach a log.
///
/// The same reasoning as `LlmConfig` and `LlmClient`: a derived `Debug` prints
/// every value, so one `{:?}` anywhere would emit every credential the user has.
/// The endpoints are printed because they are not secret and they are the useful
/// half when debugging.
impl std::fmt::Debug for AuthStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AuthStore")
            .field("endpoints", &self.keys.keys().collect::<Vec<_>>())
            .finish()
    }
}

/// What can go wrong reading or writing the store.
#[derive(Debug, Error)]
pub enum AuthError {
    #[error("no config directory for this platform; set the key in drep.toml instead")]
    NoConfigDir,

    #[error("could not read {0}: {1}")]
    Read(PathBuf, std::io::Error),

    #[error("could not write {0}: {1}")]
    Write(PathBuf, std::io::Error),

    #[error("could not lock auth store {0}: {1}")]
    Lock(PathBuf, std::io::Error),

    #[error("could not parse {0}: {1}")]
    Parse(PathBuf, String),

    /// Serializing the store failed. Distinct from [`Self::Parse`], which is
    /// about reading: a write failure reported as "could not parse" sends the
    /// reader looking at the file rather than at what drep tried to write.
    #[error("could not serialize the auth store: {0}")]
    Serialize(String),

    #[error("refusing to store an empty key for {0}")]
    EmptyKey(String),

    /// A key carrying a control character cannot be sent as an HTTP header, so
    /// storing it would defer a guaranteed failure to the first request.
    #[error("the key for {0} contains a character that cannot be sent in a header")]
    UnusableKey(String),

    /// A configured `api_key_command` did not produce a credential.
    ///
    /// Fatal, and fatal *here*: the provider chain does not exist yet, so there
    /// is nothing to fail over to - and failing over would be the wrong answer
    /// anyway, for the reason a 401 does not fail over. Routing around a broken
    /// credential path is what hides it.
    ///
    /// The entry is numbered in file order, matching `ConfigError`, because that
    /// is the numbering a user can count blocks in `drep.toml` to check.
    #[error("[[llm]] #{} in file order: {cause}", index + 1)]
    KeyCommand {
        index: usize,
        cause: KeyCommandError,
    },
}

/// The environment variable that relocates the store.
pub const PATH_VAR: &str = "DREP_AUTH_PATH";

/// The store location: [`PATH_VAR`] if set, else the platform's config dir.
///
/// The platform path comes from `directories::ProjectDirs` with the same triple
/// `Cache::default_root` uses, so drep's two user-level directories are siblings
/// under one application identity rather than two unrelated paths.
///
/// The override exists because there is otherwise **no way to run drep against a
/// scratch store**. `directories` follows each platform's own convention rather
/// than the XDG variables, so on macOS `XDG_CONFIG_HOME` is ignored entirely and
/// a command run to try something out writes into the real store - which is how
/// a test key ended up in one. It also serves the ordinary case of keeping
/// credentials somewhere deliberate, such as a mounted volume.
pub fn default_path() -> Result<PathBuf, AuthError> {
    path_from(std::env::var_os(PATH_VAR))
}

/// [`default_path`] with the override supplied rather than read.
///
/// Split out so the override can be tested without writing to the process
/// environment. `std::env::set_var` is `unsafe` in edition 2024 because another
/// thread reading the environment concurrently is a data race, and `cargo test`
/// runs tests on several threads - a "single-threaded test process" safety
/// comment would simply have been untrue.
pub fn path_from(overridden: Option<std::ffi::OsString>) -> Result<PathBuf, AuthError> {
    if let Some(path) = overridden {
        return Ok(PathBuf::from(path));
    }
    directories::ProjectDirs::from("dev", "slb350", "drep")
        .map(|dirs| dirs.config_dir().join("auth.toml"))
        .ok_or(AuthError::NoConfigDir)
}

impl AuthStore {
    /// An empty store, for a machine that has never stored a key.
    pub fn new() -> Self {
        Self::default()
    }

    /// The key held for `endpoint`, if any.
    pub fn get(&self, endpoint: &str) -> Option<&str> {
        self.keys.get(&normalise(endpoint)).map(String::as_str)
    }

    /// Store `key` for `endpoint`, replacing any previous one.
    ///
    /// The rule about what a credential may be is `vet`'s, not this method's;
    /// only the wording is here, because a paste at the prompt and a helper's
    /// stdout send the reader to different fixes.
    pub fn set(&mut self, endpoint: &str, key: &str) -> Result<(), AuthError> {
        let key = vet(key).map_err(|defect| match defect {
            CredentialDefect::Empty => AuthError::EmptyKey(endpoint.to_string()),
            CredentialDefect::Unusable => AuthError::UnusableKey(endpoint.to_string()),
        })?;
        let endpoint = normalise(endpoint);
        self.keys.insert(endpoint.clone(), key);
        self.pending
            .get_mut()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .insert(endpoint);
        Ok(())
    }

    /// Forget the key for `endpoint`. Returns whether one was held.
    pub fn remove(&mut self, endpoint: &str) -> bool {
        let endpoint = normalise(endpoint);
        let removed = self.keys.remove(&endpoint).is_some();
        if removed {
            self.pending
                .get_mut()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .insert(endpoint);
        }
        removed
    }

    /// Every endpoint with a stored key, in sorted order. Never the keys.
    pub fn endpoints(&self) -> Vec<&str> {
        self.keys.keys().map(String::as_str).collect()
    }

    /// Whether the store holds nothing.
    pub fn is_empty(&self) -> bool {
        self.keys.is_empty()
    }
}

/// Why a credential cannot be used, whatever produced it.
///
/// Two variants because they send the reader to different fixes: nothing arrived
/// at all, or something arrived that cannot be sent.
pub(crate) enum CredentialDefect {
    Empty,
    Unusable,
}

/// Trim a candidate credential and refuse one that cannot become a header value.
///
/// One definition, because this is a safety property, and `crate::http`'s module
/// doc already records what happens to a safety property written twice: "written
/// once and forgotten once". It was written twice - here for a key pasted at the
/// prompt, and in `auth::command::run` for one a helper printed - and the two
/// copies had already drifted cosmetically. The minted path is the one that never
/// passes through [`AuthStore::set`], so a defect class added at the prompt would
/// silently not have applied to the credential a gateway helper produces, which is
/// the only path whose value nobody ever sees.
///
/// The wording stays with each caller: this returns [`CredentialDefect`] rather
/// than an `AuthError`, so the store can name the endpoint and the helper can name
/// the program. That is the mechanism-shared, classification-private split
/// `crate::http` documents for the same pair of callers.
///
/// **Both ends** are trimmed. Every helper that prints a token prints a newline
/// after it, and a leading space is one `printf ' %s'` or one `cut` field away;
/// trimming only the tail accepted `" sk-live-..."`, which passes the
/// control-character guard and is then stripped again by the transport, so the
/// value sent was not the value checked and the endpoint answered 401.
///
/// A control character is refused because every use of a credential is an HTTP
/// header value, which cannot carry one. Rejecting it here turns a paste or a
/// helper that picked up a stray escape sequence into a message the user can act
/// on, rather than a transport failure on the first file of the first push.
pub(crate) fn vet(candidate: &str) -> Result<String, CredentialDefect> {
    let candidate = candidate.trim();
    if candidate.is_empty() {
        return Err(CredentialDefect::Empty);
    }
    if candidate.chars().any(char::is_control) {
        return Err(CredentialDefect::Unusable);
    }
    Ok(candidate.to_owned())
}

/// Canonical form of an endpoint for use as a store key.
///
/// Lowercased and stripped of trailing slashes. `https://API.Z.AI/v1/` and
/// `https://api.z.ai/v1` are the same host and the same credential, and a store
/// that disagreed would report "no key" for a config that had just stored one -
/// a failure with no visible cause, since both spellings look right.
///
/// Deliberately no more than that: the path is significant (`/v1` and
/// `/anthropic/v1` are different APIs on the same host, and can carry different
/// keys), so nothing beyond the trailing slash is trimmed.
pub fn normalise(endpoint: &str) -> String {
    let trimmed = endpoint.trim().trim_end_matches('/');

    // Only the scheme and authority are case-insensitive; a URL *path* is not.
    // Lowercasing the whole thing collapsed `/API/v1` and `/api/v1` onto one
    // entry, which for a host serving both would hand one endpoint's key to the
    // other - the same class of mistake as keying on the model alone.
    match trimmed.find("://") {
        Some(scheme_end) => format!(
            "{}://{}",
            trimmed[..scheme_end].to_ascii_lowercase(),
            lower_authority(&trimmed[scheme_end + 3..])
        ),
        // No scheme, which is what `localhost:11434/v1` looks like. It still has
        // an authority and a path, and the same rule applies to both halves -
        // lowercasing the whole string collapsed `/V1` onto `/v1` for exactly
        // the endpoints a user is most likely to type by hand.
        None => lower_authority(trimmed),
    }
}

/// Lowercase everything before the first `/` and leave the rest alone.
fn lower_authority(rest: &str) -> String {
    let host_len = rest.find('/').unwrap_or(rest.len());
    format!(
        "{}{}",
        rest[..host_len].to_ascii_lowercase(),
        &rest[host_len..]
    )
}

/// What an endpoint prints as wherever drep shows one to a person or a log.
const REDACTED: &str = "<redacted>";

/// An endpoint as drep prints it: without the credential a URL can carry.
///
/// An endpoint's `user:password@` is a credential, not an address. reqwest
/// removes it from the URL and sends it as HTTP Basic auth, so a gateway behind
/// a password is configured as `https://drep:${GATEWAY_PASSWORD}@host/v1`, and
/// by the time a report prints the loaded config the variable holds the secret.
/// Every place an endpoint is shown - `check`'s failover block and JSON,
/// `doctor`, `init`, `auth`, and `Debug` - goes through here, and prints
/// `https://<redacted>@host/v1` instead: the credential's presence stays
/// visible, its value does not.
///
/// The URL is read with the parser reqwest itself uses, rather than by finding
/// an `@`, so this strips what would be sent: `http://evil.example\@host/v1`
/// names the host `evil.example` and carries no credential, and is printed as
/// written. An endpoint without one is returned untouched, spelling and all.
/// One that does not parse, or parses with no host (a scheme-less
/// `user:pass@localhost:1234/v1` reads as an opaque path), cannot be split
/// reliably, so it prints as `<redacted>` whole when it contains an `@`.
///
/// Display only. The cache identity keeps the raw endpoint, so rotating a
/// password still starts the cache cold.
pub fn display_endpoint(endpoint: &str) -> Cow<'_, str> {
    let unsplittable = || {
        Cow::Borrowed(if endpoint.contains('@') {
            REDACTED
        } else {
            endpoint
        })
    };
    let Ok(mut url) = reqwest::Url::parse(endpoint) else {
        return unsplittable();
    };
    if url.username().is_empty() && url.password().is_none() {
        return if url.has_host() {
            Cow::Borrowed(endpoint)
        } else {
            unsplittable()
        };
    }
    // Cannot fail: only a URL with no host, or a `file:` one, refuses a
    // username, and neither parses with a credential in the first place.
    if url.set_username("").and(url.set_password(None)).is_err() {
        return Cow::Borrowed(REDACTED);
    }
    match url.as_str().split_once("://") {
        Some((scheme, rest)) => Cow::Owned(format!("{scheme}://{REDACTED}@{rest}")),
        None => Cow::Borrowed(REDACTED),
    }
}

/// Fill in keys the config left unset, and report where each one came from.
///
/// Returns one [`KeySource`] per entry in `config.llm`, positionally - including
/// the disabled ones, so the position in the returned slice is the position in the
/// file. No production caller reads it: `check` discards it, and `doctor` calls
/// [`source_of`] against the raw TOML tree instead, because `config::load` fails
/// on an unset `${VAR}` and that is exactly the config `doctor` is most useful on.
/// It is retained for the tests, which assert the precedence rule per entry
/// against the one function that decides it, and the positional shape is what lets
/// them name an entry by the line the user wrote.
///
/// **Disabled entries are skipped**, matching every other pass over the provider
/// list: `${VAR}` expansion and field validation already leave a parked entry
/// alone, and looking a key up for one would report a missing credential for a
/// provider that is never contacted. For an `api_key_command` that also means no
/// subprocess: some helpers are rate-limited and some prompt for a fingerprint,
/// so spending a real credential call on a provider drep will not contact is
/// worse than the missing-key report it avoids.
///
/// This is the one pass where a credential command runs, so each entry's command
/// runs exactly once per process however many files the run reviews. A
/// short-lived credential re-minted per file would fail per file, which the
/// chain's demotion logic reads as an endpoint problem. There is deliberately no
/// disk cache and no TTL behind that: drep is a short-lived process, so a
/// credential written to disk buys nothing and adds a file worth stealing.
///
/// Once per process for every *enabled* entry, whether or not the chain reaches
/// it. Deferring to first use would put credential resolution inside the request
/// path, and "a broken credential is fatal rather than a provider drep quietly
/// asks instead" holds precisely because resolution happens before the chain
/// exists. An unset `${VAR}` in the same position is equally fatal, one layer up,
/// in `ConfigError::EnvVarUnset`. `enabled = false` is the control for a fallback
/// whose helper should not be spent.
pub async fn resolve(config: &mut Config, store: &AuthStore) -> Result<Vec<KeySource>, AuthError> {
    let mut sources = Vec::with_capacity(config.llm.len());
    for (index, llm) in config.llm.iter_mut().enumerate() {
        let source = source_of(
            Declared {
                api_key: llm.api_key.as_deref(),
                has_api_key_command: llm.api_key_command.is_some(),
                endpoint: llm.endpoint.as_deref(),
                enabled: llm.enabled,
            },
            store,
        );
        match source {
            // `source_of` returns `Command` only when the argv is present, so
            // the default here is unreachable through it; `command::run` reports
            // an empty argv rather than panicking if it ever becomes reachable.
            KeySource::Command => {
                let argv = llm.api_key_command.as_deref().unwrap_or_default();
                let key = command::run_bounded(argv)
                    .await
                    .map_err(|cause| AuthError::KeyCommand { index, cause })?;
                llm.api_key = Some(key);
            }
            KeySource::Store => {
                if let Some(endpoint) = llm.endpoint.as_deref()
                    && let Some(key) = store.get(endpoint)
                {
                    llm.api_key = Some(key.to_string());
                }
            }
            KeySource::Config | KeySource::Missing => {}
        }
        sources.push(source);
    }
    Ok(sources)
}

/// Run one entry's `api_key_command` and discard the credential.
///
/// For `doctor`, whose contract is "what will actually run here": it has to
/// really invoke the helper, because a helper that no longer authenticates is
/// precisely what `api_key_command` exists to make visible. Returning `()`
/// rather than the key is what makes printing it impossible at the call site,
/// instead of a rule the reporting code has to remember.
pub async fn probe_key_command(argv: &[String]) -> Result<(), KeyCommandError> {
    command::run_bounded(argv).await.map(drop)
}

/// What one `[[llm]]` entry declares about its own credential.
///
/// Named fields rather than positional arguments, following `ExplicitFields`:
/// `check` fills these from a loaded `Config` and `doctor` from the raw TOML
/// tree, and two `Option<&str>` plus two `bool` in a call is a swap waiting to
/// happen that the compiler would not catch.
pub struct Declared<'a> {
    pub api_key: Option<&'a str>,
    pub has_api_key_command: bool,
    pub endpoint: Option<&'a str>,
    pub enabled: bool,
}

/// Where a provider's key will come from, given what its config names.
///
/// The precedence rule itself, in one place. [`resolve`] applies it to a loaded
/// `Config`; `doctor` applies it to the *raw* TOML tree, which it has to read
/// separately so a `${VAR}` prints as itself rather than being swallowed by the
/// variable-not-set error. Two readers of one rule is the shape
/// `config::env_var_refs_in` already exists to prevent: doctor once carried a
/// narrower copy of that scanner and reported a config as fine that `check`
/// refused to load.
pub fn source_of(declared: Declared<'_>, store: &AuthStore) -> KeySource {
    if !declared.enabled {
        return KeySource::Missing;
    }
    if declared.api_key.is_some() {
        return KeySource::Config;
    }
    if declared.has_api_key_command {
        return KeySource::Command;
    }
    match declared.endpoint {
        Some(endpoint) if store.get(endpoint).is_some() => KeySource::Store,
        _ => KeySource::Missing,
    }
}

#[cfg(test)]
mod tests;