mobius_gateway/config/

store.rs

use super::*;
use mobius::backend::model::ModelCredentialLifetime;

/// File owner for gateway configuration and aggregate usage.
#[derive(Debug, Clone)]
pub struct ConfigStore {
    #[cfg(test)]
    pub(crate) runtime_operations: std::sync::Arc<RuntimeOperations>,
    state_dir: PathBuf,
    path: PathBuf,
}

/// Owner-only API-key storage kept outside frontend-readable configuration.
pub struct CredentialStore {
    path: PathBuf,
    values: Mutex<BTreeMap<String, StoredCredential>>,
}

/// A credential resolved atomically with the lifetime of the same stored secret.
#[derive(Clone)]
pub struct ResolvedCredential {
    /// Secret used only for provider assembly.
    pub api_key: String,
    /// Expiry and revocation shared by all routes using this secret.
    pub lifetime: ModelCredentialLifetime,
}

#[derive(Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct StoredCredential {
    provider: String,
    api_key: String,
    base_url: Option<String>,
    expires_at: Option<u64>,
    #[serde(skip)]
    revocation: Option<tokio::sync::watch::Sender<()>>,
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(super) struct UsageHistory {
    pub(super) days: BTreeMap<u64, BTreeMap<String, TokenUsage>>,
}

impl ConfigStore {
    /// Initializes an owner-only state directory and new config file.
    /// # Errors
    ///
    /// Returns an error if configuration is invalid or a required resource cannot be initialized.
    pub fn initialize(
        state_dir: PathBuf,
        listen: SocketAddr,
        tls: Option<TlsConfig>,
    ) -> Result<(Self, GatewayConfig)> {
        let config = GatewayConfig::new(listen, tls)?;
        let state_dir = prepare_state_dir(state_dir)?;
        let store = Self::at(state_dir);
        store.save_with_mode(&config, true)?;
        Ok((store, config))
    }

    /// Initializes state for an account-free Cloudflare Quick Tunnel.
    /// # Errors
    ///
    /// Returns an error if configuration is invalid or a required resource cannot be initialized.
    pub fn initialize_quick_cloudflare(
        state_dir: PathBuf,
        listen: SocketAddr,
    ) -> Result<(Self, GatewayConfig)> {
        Self::initialize_cloudflare(state_dir, listen, CloudflareConfig::Quick, None)
    }

    /// Initializes state for one user-owned Cloudflare Tunnel.
    /// # Errors
    ///
    /// Returns an error if configuration is invalid or a required resource cannot be initialized.
    pub fn initialize_named_cloudflare(
        state_dir: PathBuf,
        listen: SocketAddr,
        hostname: &str,
        token: &str,
    ) -> Result<(Self, GatewayConfig)> {
        Self::initialize_cloudflare(
            state_dir,
            listen,
            CloudflareConfig::named(hostname)?,
            Some(validate_cloudflare_token(token)?),
        )
    }

    fn initialize_cloudflare(
        state_dir: PathBuf,
        listen: SocketAddr,
        cloudflare: CloudflareConfig,
        token: Option<&str>,
    ) -> Result<(Self, GatewayConfig)> {
        let config = GatewayConfig::new_cloudflare(listen, cloudflare)?;
        let state_dir = prepare_state_dir(state_dir)?;
        let store = Self::at(state_dir);
        let result = token
            .map_or(Ok(()), |token| store.save_cloudflare_token(token))
            .and_then(|()| store.save_with_mode(&config, true));
        if let Err(error) = result {
            fs::remove_dir_all(&store.state_dir).map_err(|cleanup| {
                Error::Config(format!(
                    "{error}; failed to remove incomplete gateway state at {}: {cleanup}",
                    store.state_dir.display()
                ))
            })?;
            return Err(error);
        }
        Ok((store, config))
    }

    /// Opens and validates persisted gateway configuration.
    /// # Errors
    ///
    /// Returns an error if the resource cannot be read, decoded, or validated.
    pub fn open(state_dir: PathBuf) -> Result<(Self, GatewayConfig)> {
        let state_dir = fs::canonicalize(state_dir)?;
        validate_private_state_dir(&state_dir)?;
        let store = Self::at(state_dir);
        let mut file = fs::File::open(&store.path)?;
        let mut contents = Vec::new();
        std::io::Read::by_ref(&mut file)
            .take(MAX_CONFIG_BYTES + 1)
            .read_to_end(&mut contents)?;
        if u64::try_from(contents.len()).unwrap_or(u64::MAX) > MAX_CONFIG_BYTES {
            return Err(Error::Config("gateway configuration is too large".into()));
        }
        let config = toml::from_slice(&contents).map_err(|error| {
            Error::Config(format!(
                "gateway state at {} is incompatible with this release; remove that directory and run `mobius` again: {error}",
                store.state_dir.display()
            ))
        })?;
        store.validate_config(&config)?;
        Ok((store, config))
    }

    /// Atomically replaces validated persistent configuration.
    /// # Errors
    ///
    /// Returns an error if the value cannot be encoded or persisted.
    pub fn save(&self, config: &GatewayConfig) -> Result<()> {
        self.save_with_mode(config, false)
    }

    /// Returns the protected state directory.
    #[must_use]
    pub fn state_dir(&self) -> &Path {
        &self.state_dir
    }

    /// Returns the owner-managed extension store outside sandbox-masked gateway state.
    #[must_use]
    pub(crate) fn extensions_path(&self) -> PathBuf {
        crate::extensions::extensions_path(&self.state_dir)
    }

    /// Returns the provider credential file path.
    #[must_use]
    pub fn credentials_path(&self) -> PathBuf {
        self.state_dir.join("credentials.json")
    }

    /// Returns the provider browser-auth file path.
    #[must_use]
    pub fn provider_auth_path(&self) -> PathBuf {
        self.state_dir.join("provider-auth.json")
    }

    /// Returns the checkpoint database path.
    #[must_use]
    pub fn checkpoints_path(&self) -> PathBuf {
        self.state_dir.join("checkpoints.sqlite3")
    }

    /// Returns the authentication state path.
    #[must_use]
    pub fn auth_path(&self) -> PathBuf {
        self.state_dir.join("auth.json")
    }

    /// Returns the owner-only Cloudflare connector-token path.
    #[must_use]
    pub fn cloudflare_token_path(&self) -> PathBuf {
        self.state_dir.join(CLOUDFLARE_TOKEN_FILE)
    }

    fn at(state_dir: PathBuf) -> Self {
        let path = state_dir.join(CONFIG_FILE);
        Self {
            state_dir,
            path,
            #[cfg(test)]
            runtime_operations: Default::default(),
        }
    }

    fn save_with_mode(&self, config: &GatewayConfig, create_new: bool) -> Result<()> {
        self.validate_config(config)?;
        let config = toml::to_string_pretty(config).map_err(|error| {
            Error::Config(format!("cannot encode gateway configuration: {error}"))
        })?;
        let contents = config;
        if u64::try_from(contents.len()).unwrap_or(u64::MAX) > MAX_CONFIG_BYTES {
            return Err(Error::Config("gateway configuration is too large".into()));
        }
        crate::publication::publish(&self.path, contents.as_bytes(), create_new)
    }

    fn validate_config(&self, config: &GatewayConfig) -> Result<()> {
        config.validate()?;
        if matches!(
            config.cloudflare.as_ref(),
            Some(CloudflareConfig::Named { .. })
        ) {
            load_secret_file(&self.cloudflare_token_path())?;
        }
        Ok(())
    }

    fn save_cloudflare_token(&self, token: &str) -> Result<()> {
        let token = validate_cloudflare_token(token)?;
        crate::publication::publish(&self.cloudflare_token_path(), token.as_bytes(), true)
    }
}

impl CredentialStore {
    /// Opens credential state, treating a missing file as an empty store.
    /// # Errors
    ///
    /// Returns an error if the resource cannot be read, decoded, or validated.
    pub fn open(path: PathBuf) -> Result<Self> {
        let values = match fs::read(&path) {
            Ok(contents) => {
                if contents.len() > MAX_CREDENTIAL_STATE_BYTES {
                    return Err(Error::Config(
                        "provider credential state is too large".into(),
                    ));
                }
                serde_json::from_slice(&contents)?
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => BTreeMap::new(),
            Err(error) => return Err(error.into()),
        };
        validate_credential_state(&values)?;
        Ok(Self {
            path,
            values: Mutex::new(values),
        })
    }

    /// Atomically replaces one instance's API key after provider and size validation.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn set(
        &self,
        instance: &str,
        provider_id: &str,
        api_key: &str,
        base_url: Option<&str>,
        expires_at: Option<u64>,
    ) -> Result<()> {
        let api_key = api_key.trim();
        validate_new_api_key(api_key)?;
        let credential = StoredCredential {
            provider: provider_id.into(),
            api_key: api_key.into(),
            base_url: base_url.map(str::to_owned),
            expires_at,
            revocation: None,
        };
        validate_stored_credential(instance, &credential)?;
        let mut values = self
            .values
            .lock()
            .map_err(|_| Error::Config("provider credential lock is poisoned".into()))?;
        if let Some(credential) = values.get(instance)
            && credential.provider != provider_id
        {
            return Err(Error::Config(format!(
                "provider instance `{instance}` already belongs to `{}`",
                credential.provider
            )));
        }
        if values.get(instance).is_some_and(|current| {
            current.provider == credential.provider
                && current.api_key == credential.api_key
                && current.base_url == credential.base_url
                && current.expires_at == credential.expires_at
        }) {
            return Ok(());
        }
        let mut next = values.clone();
        next.insert(instance.into(), credential);
        save_private_map(&self.path, &next)?;
        *values = next;
        Ok(())
    }

    /// Resolves a secret and its lifetime under the same lock.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn get(
        &self,
        instance: &str,
        provider_id: &str,
        base_url: Option<&str>,
    ) -> Result<Option<ResolvedCredential>> {
        let mut values = self
            .values
            .lock()
            .map_err(|_| Error::Config("provider credential lock is poisoned".into()))?;
        let Some(credential) = values.get_mut(instance).filter(|credential| {
            credential.provider == provider_id && credential.base_url.as_deref() == base_url
        }) else {
            return Ok(None);
        };
        let revoked = credential
            .revocation
            .get_or_insert_with(|| tokio::sync::watch::channel(()).0)
            .subscribe();
        Ok(Some(ResolvedCredential {
            api_key: credential.api_key.clone(),
            lifetime: ModelCredentialLifetime {
                expires_at: credential
                    .expires_at
                    .map(|seconds| UNIX_EPOCH + std::time::Duration::from_secs(seconds)),
                revoked: Some(revoked),
            },
        }))
    }

    /// Returns a non-secret suffix for identifying one stored credential.
    pub(crate) fn hint(
        &self,
        instance: &str,
        provider_id: &str,
        base_url: Option<&str>,
    ) -> Result<Option<String>> {
        let values = self
            .values
            .lock()
            .map_err(|_| Error::Config("provider credential lock is poisoned".into()))?;
        Ok(values
            .get(instance)
            .filter(|credential| {
                credential.provider == provider_id && credential.base_url.as_deref() == base_url
            })
            .and_then(|credential| {
                let suffix = credential.api_key.chars().rev().take(4).collect::<String>();
                (suffix.chars().count() == 4).then(|| suffix.chars().rev().collect())
            }))
    }

    /// Atomically removes one instance-scoped API-key credential.
    /// # Errors
    ///
    /// Returns an error if validation or an operation required by this function fails.
    pub fn remove(&self, instance: &str) -> Result<bool> {
        super::validation::validate_instance_id(instance)?;
        let mut values = self
            .values
            .lock()
            .map_err(|_| Error::Config("provider credential lock is poisoned".into()))?;
        if !values.contains_key(instance) {
            return Ok(false);
        }
        let mut next = values.clone();
        next.remove(instance);
        save_private_map(&self.path, &next)?;
        *values = next;
        Ok(true)
    }
}

/// Resolves the gateway state directory from the environment or home directory.
/// # Errors
///
/// Returns an error if validation or an operation required by this function fails.
pub fn state_dir() -> Result<PathBuf> {
    if let Some(path) = env::var_os("MOBIUS_GATEWAY_STATE_DIR") {
        if path.is_empty() {
            return Err(Error::Config("MOBIUS_GATEWAY_STATE_DIR is empty".into()));
        }
        return Ok(path.into());
    }
    env::home_dir()
        .filter(|path| !path.as_os_str().is_empty())
        .map(|path| path.join(".mobius").join("gateway"))
        .ok_or_else(|| {
            Error::Config(
                "cannot determine the home directory; set MOBIUS_GATEWAY_STATE_DIR".into(),
            )
        })
}

/// Loads a connector token from an owner-only regular file without exposing its contents.
/// # Errors
///
/// Returns an error if the resource cannot be read, decoded, or validated.
pub fn load_secret_file(path: &Path) -> Result<String> {
    #[cfg(unix)]
    let file = fs::OpenOptions::new()
        .read(true)
        .custom_flags(nix::libc::O_NOFOLLOW | nix::libc::O_NONBLOCK)
        .open(path)
        .map_err(|error| -> Error {
            if error.raw_os_error() == Some(nix::libc::ELOOP) {
                Error::Config("secret must be stored in a regular file".into())
            } else {
                error.into()
            }
        })?;
    #[cfg(not(unix))]
    let file = {
        let metadata = fs::symlink_metadata(path)?;
        if !metadata.file_type().is_file() {
            return Err(Error::Config(
                "secret must be stored in a regular file".into(),
            ));
        }
        fs::File::open(path)?
    };
    let metadata = file.metadata()?;
    if !metadata.file_type().is_file() {
        return Err(Error::Config(
            "secret must be stored in a regular file".into(),
        ));
    }
    #[cfg(unix)]
    if metadata.permissions().mode() & 0o077 != 0 {
        return Err(Error::Config(
            "secret file must not be accessible by group or others (use mode 0600)".into(),
        ));
    }
    if metadata.len() > MAX_CLOUDFLARE_TOKEN_BYTES as u64 {
        return Err(Error::Config("secret file token is too large".into()));
    }
    let mut contents = String::new();
    file.take(MAX_CLOUDFLARE_TOKEN_BYTES as u64 + 1)
        .read_to_string(&mut contents)?;
    let token = validate_cloudflare_token(&contents)
        .map_err(|_| Error::Config("secret file contains an invalid token".into()))?;
    Ok(token.to_owned())
}

fn prepare_state_dir(path: PathBuf) -> Result<PathBuf> {
    let name = path
        .file_name()
        .ok_or_else(|| Error::Config("gateway state directory must have a name".into()))?
        .to_owned();
    let parent = path
        .parent()
        .filter(|parent| !parent.as_os_str().is_empty())
        .unwrap_or(Path::new("."));
    fs::create_dir_all(parent)?;
    let path = fs::canonicalize(parent)?.join(name);
    match fs::create_dir(&path) {
        Ok(()) => {}
        Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
            return Err(Error::Config(
                "gateway state directory already exists".into(),
            ));
        }
        Err(error) => return Err(error.into()),
    }
    #[cfg(unix)]
    fs::set_permissions(&path, mobius::owner_only::dir())?;
    Ok(path)
}

fn validate_private_state_dir(path: &Path) -> Result<()> {
    let metadata = fs::metadata(path)?;
    if !metadata.is_dir() {
        return Err(Error::Config(
            "gateway state path must be a directory".into(),
        ));
    }
    #[cfg(unix)]
    if metadata.permissions().mode() & 0o077 != 0 {
        return Err(Error::Config(
            "gateway state directory must not be accessible by group or others (use mode 0700)"
                .into(),
        ));
    }
    Ok(())
}

fn validate_stored_credential(instance: &str, credential: &StoredCredential) -> Result<()> {
    if credential
        .expires_at
        .is_some_and(|seconds| seconds == 0 || seconds > 253_402_300_799)
    {
        return Err(Error::Config(
            "credential expiry must be a valid Unix timestamp".into(),
        ));
    }
    super::validation::validate_instance_id(instance)?;
    let definition = provider(&credential.provider)?;
    if !matches!(definition.auth(), ProviderAuth::ApiKey(_)) {
        return Err(Error::Config(format!(
            "provider `{}` does not accept an API key",
            credential.provider
        )));
    }
    if credential.api_key.trim().is_empty() || credential.api_key.len() > MAX_PROVIDER_API_KEY_BYTES
    {
        return Err(Error::Config(format!(
            "API key must be 1–{MAX_PROVIDER_API_KEY_BYTES} bytes"
        )));
    }
    definition.validate_base_url(credential.base_url.as_deref())?;
    Ok(())
}

fn validate_new_api_key(api_key: &str) -> Result<()> {
    if api_key.is_empty() || api_key.len() > MAX_PROVIDER_API_KEY_BYTES {
        return Err(Error::Config(format!(
            "API key must be 1–{MAX_PROVIDER_API_KEY_BYTES} bytes"
        )));
    }
    if !api_key.bytes().all(|byte| byte.is_ascii_graphic()) {
        return Err(Error::Config(
            "API key must contain only visible ASCII characters without whitespace".into(),
        ));
    }
    Ok(())
}

fn validate_credential_state(values: &BTreeMap<String, StoredCredential>) -> Result<()> {
    for (instance, credential) in values {
        validate_stored_credential(instance, credential)?;
    }
    Ok(())
}

fn save_private_map(path: &Path, values: &BTreeMap<String, StoredCredential>) -> Result<()> {
    validate_credential_state(values)?;
    let contents = serde_json::to_vec(values)?;
    if contents.len() > MAX_CREDENTIAL_STATE_BYTES {
        return Err(Error::Config(
            "provider credential state is too large".into(),
        ));
    }
    crate::publication::publish(path, &contents, false)
}

impl UsageHistory {
    pub(super) fn observe(
        &mut self,
        provider: &str,
        usage: &TokenUsage,
        now: SystemTime,
    ) -> Result<bool> {
        validate_usage_provider(provider)?;
        validate_usage(usage)?;
        if usage == &TokenUsage::default() {
            return Ok(false);
        }
        let day = unix_day(now)?;
        let mut bucket = self
            .days
            .get(&day)
            .and_then(|providers| providers.get(provider))
            .cloned()
            .unwrap_or_default();
        bucket
            .checked_add(usage)
            .ok_or_else(|| Error::Config("daily token usage overflow".into()))?;
        self.days
            .entry(day)
            .or_default()
            .insert(provider.into(), bucket);
        let first_day = day.saturating_sub(USAGE_HISTORY_DAYS - 1);
        self.days.retain(|stored, _| *stored >= first_day);
        Ok(true)
    }
}

pub(super) fn unix_day(now: SystemTime) -> Result<u64> {
    Ok(now
        .duration_since(UNIX_EPOCH)
        .map_err(|_| Error::Config("system clock is before the Unix epoch".into()))?
        .as_secs()
        / SECONDS_PER_DAY)
}

pub(super) fn validate_usage(usage: &TokenUsage) -> Result<()> {
    if !usage_nonnegative(usage) {
        return Err(Error::Config("token usage cannot be negative".into()));
    }
    Ok(())
}

pub(super) fn validate_usage_provider(provider: &str) -> Result<()> {
    if provider.trim().is_empty()
        || provider != provider.trim()
        || provider.len() > 256
        || provider.chars().any(char::is_control)
    {
        return Err(Error::Config(
            "usage provider ID must be canonical and 1–256 bytes".into(),
        ));
    }
    Ok(())
}

fn usage_nonnegative(usage: &TokenUsage) -> bool {
    usage.input_tokens >= 0
        && usage.cached_input_tokens >= 0
        && usage.cache_write_input_tokens >= 0
        && usage.output_tokens >= 0
        && usage.reasoning_output_tokens >= 0
        && usage.total_tokens >= 0
}

#[cfg(test)]
#[derive(Debug, Default)]
pub(crate) struct RuntimeOperations {
    pub(crate) preparations: std::sync::atomic::AtomicUsize,
    pub(crate) assemblies: std::sync::atomic::AtomicUsize,
}

#[cfg(test)]
impl RuntimeOperations {
    pub(crate) fn counts(&self) -> (usize, usize) {
        use std::sync::atomic::Ordering::Relaxed;
        (
            self.preparations.load(Relaxed),
            self.assemblies.load(Relaxed),
        )
    }
}