agnt5-sdk-core 0.3.9

//! Certificate-bound identity lifecycle for customer-hosted workers.
//!
//! This module is deliberately the only SDK-core component allowed to read the
//! bootstrap credential or persist private session material. Callers receive a
//! TLS configuration and a rotating bearer-token interceptor, never raw keys.

use crate::error::{ErrorCode, Result, SdkError};
use base64::Engine;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use rcgen::{CertificateParams, KeyPair, PKCS_ECDSA_P256_SHA256};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock, RwLock, Weak};
use std::time::Duration;
use tokio::sync::{Mutex, OnceCell};
use tonic::metadata::MetadataValue;
use tonic::transport::{Certificate, ClientTlsConfig, Identity};
use tracing::{error, warn};

const SESSION_OPEN_PATH: &str = "/api/v1/external-worker-sessions";
const TOKEN_REFRESH_PATH: &str = "/api/v1/external-worker-sessions/token";
const SESSION_RENEW_PATH: &str = "/api/v1/external-worker-sessions/renew";
const SESSION_FILE_NAME: &str = "worker-session.json";
const REFRESH_MARGIN: ChronoDuration = ChronoDuration::seconds(60);

static MANAGER: OnceCell<Arc<IdentityManager>> = OnceCell::const_new();
static SESSION_FILE_LOCK: OnceLock<Mutex<()>> = OnceLock::new();

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub(crate) struct Authority {
    pub(crate) project_id: String,
    pub(crate) environment_id: String,
    pub(crate) deployment_id: String,
    pub(crate) worker_pool_id: String,
    #[serde(default)]
    pub(crate) runtime_endpoint: String,
    #[serde(default)]
    pub(crate) protocol: String,
}

#[derive(Debug, Serialize)]
struct OpenRequest<'a> {
    project_id: &'a str,
    environment_id: &'a str,
    deployment_id: &'a str,
    worker_pool_id: &'a str,
    csr_der_base64: &'a str,
}

#[derive(Debug, Serialize)]
struct RenewRequest<'a> {
    request_id: &'a str,
    csr_der_base64: &'a str,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct Session {
    session_id: String,
    project_id: String,
    environment_id: String,
    deployment_id: String,
    worker_pool_id: String,
    worker_id: String,
    spiffe_id: String,
    #[serde(default)]
    runtime_endpoint: String,
    certificate_der_base64: String,
    certificate_chain_der_base64: Vec<String>,
    trust_bundle_der_base64: Vec<String>,
    trust_bundle_version: String,
    certificate_expires_at: DateTime<Utc>,
    renew_after: DateTime<Utc>,
    workload_token: String,
    token_type: String,
    token_expires_at: DateTime<Utc>,
    #[serde(default)]
    private_key_pem: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pending_renewal: Option<PendingRenewal>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct PendingRenewal {
    request_id: String,
    private_key_pem: String,
    csr_der_base64: String,
}

#[derive(Debug, Deserialize)]
struct RefreshedToken {
    workload_token: String,
    token_type: String,
    token_expires_at: DateTime<Utc>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum IdentityReadiness {
    Enrolling,
    Ready {
        certificate_expires_at: DateTime<Utc>,
        token_expires_at: DateTime<Utc>,
    },
    Degraded,
}

#[derive(Clone)]
pub struct ConnectionIdentity {
    pub endpoint: String,
    pub worker_id: String,
    pub tls: ClientTlsConfig,
    pub authorization: Arc<RwLock<Option<MetadataValue<tonic::metadata::Ascii>>>>,
}

struct Subscriber {
    session_id: String,
    token: Weak<RwLock<Option<MetadataValue<tonic::metadata::Ascii>>>>,
}

struct IdentityManager {
    identity_url: String,
    authority: Authority,
    session_path: PathBuf,
    session: RwLock<Session>,
    readiness: RwLock<IdentityReadiness>,
    subscribers: Mutex<Vec<Subscriber>>,
}

pub(crate) async fn connection_identity(
    control_plane_url: String,
    identity_url: String,
    credential: String,
    authority: Authority,
) -> Result<ConnectionIdentity> {
    let expected_authority = authority.clone();
    let expected_identity_url = identity_url.clone();
    let manager = MANAGER
        .get_or_try_init(|| async {
            IdentityManager::initialize(control_plane_url, identity_url, credential, authority)
                .await
                .map(Arc::new)
        })
        .await?
        .clone();
    if manager.authority != expected_authority || manager.identity_url != expected_identity_url {
        return Err(connection_error(
            "external worker discovery changed the process identity authority; restart with its own session directory",
        ));
    }
    manager.connection_identity().await
}

pub async fn readiness() -> Option<IdentityReadiness> {
    Some(MANAGER.get()?.readiness.read().ok()?.clone())
}

impl IdentityManager {
    async fn initialize(
        control_plane_url: String,
        identity_url: String,
        credential: String,
        authority: Authority,
    ) -> Result<Self> {
        validate_identity_endpoint(&identity_url)?;
        let session_dir = PathBuf::from(required_env("AGNT5_WORKER_SESSION_DIR")?);
        ensure_private_directory(&session_dir)?;
        let session_path = session_dir.join(SESSION_FILE_NAME);
        let _guard = SESSION_FILE_LOCK
            .get_or_init(|| Mutex::new(()))
            .lock()
            .await;
        let session = if session_path
            .try_exists()
            .map_err(|_| connection_error("cannot inspect worker session file"))?
        {
            load_session(&session_path, &authority)?
        } else {
            let (key, csr) = new_key_and_csr()?;
            let mut session =
                open_session(&control_plane_url, &credential, &authority, &csr).await?;
            session.private_key_pem = key;
            validate_session(&session, &authority)?;
            write_session(&session_path, &session)?;
            session
        };
        let readiness = IdentityReadiness::Ready {
            certificate_expires_at: session.certificate_expires_at,
            token_expires_at: session.token_expires_at,
        };
        let manager = Self {
            identity_url,
            authority,
            session_path,
            session: RwLock::new(session),
            readiness: RwLock::new(readiness),
            subscribers: Mutex::new(Vec::new()),
        };
        let current = manager
            .session
            .read()
            .map_err(|_| connection_error("worker identity state unavailable"))?
            .clone();
        if current.pending_renewal.is_some() || Utc::now() >= current.renew_after {
            manager.renew(&current).await?;
        } else if current.token_expires_at <= Utc::now() + REFRESH_MARGIN {
            manager.refresh(&current).await?;
        }
        Ok(manager)
    }

    async fn connection_identity(self: &Arc<Self>) -> Result<ConnectionIdentity> {
        let session = self
            .session
            .read()
            .map_err(|_| connection_error("worker identity state is unavailable"))?
            .clone();
        validate_session(&session, &self.authority)?;
        let authorization = Arc::new(RwLock::new(Some(bearer_metadata(&session.workload_token)?)));
        self.subscribers.lock().await.push(Subscriber {
            session_id: session.session_id.clone(),
            token: Arc::downgrade(&authorization),
        });
        let manager = self.clone();
        static MAINTENANCE_STARTED: OnceLock<()> = OnceLock::new();
        if MAINTENANCE_STARTED.set(()).is_ok() {
            tokio::spawn(async move { manager.maintain().await });
        }
        Ok(ConnectionIdentity {
            endpoint: session.runtime_endpoint.clone(),
            worker_id: session.worker_id.clone(),
            tls: tls_config(&session)?,
            authorization,
        })
    }

    async fn maintain(self: Arc<Self>) {
        loop {
            tokio::time::sleep(Duration::from_secs(5)).await;
            let session = match self.session.read() {
                Ok(value) => value.clone(),
                Err(_) => continue,
            };
            let now = Utc::now();
            let result = if session.pending_renewal.is_some() || now >= session.renew_after {
                self.renew(&session).await
            } else if now + REFRESH_MARGIN >= session.token_expires_at {
                self.refresh(&session).await
            } else {
                continue;
            };
            if let Err(e) = result {
                if let Ok(mut readiness) = self.readiness.write() {
                    *readiness = IdentityReadiness::Degraded;
                }
                error!("external worker identity maintenance failed: {e}");
            }
        }
    }

    async fn refresh(&self, current: &Session) -> Result<()> {
        let client = identity_http_client(current)?;
        let response = client
            .post(format!("{}{}", self.identity_url, TOKEN_REFRESH_PATH))
            .send()
            .await
            .map_err(|e| connection_error(format!("worker token refresh failed: {e}")))?;
        if !response.status().is_success() {
            return Err(connection_error(format!(
                "worker token refresh returned {}",
                response.status()
            )));
        }
        let refreshed: RefreshedToken = response
            .json()
            .await
            .map_err(|e| connection_error(format!("decode worker token refresh: {e}")))?;
        let value = bearer_metadata(&refreshed.workload_token)?;
        let mut next = current.clone();
        next.workload_token = refreshed.workload_token;
        next.token_type = refreshed.token_type;
        next.token_expires_at = refreshed.token_expires_at;
        write_session(&self.session_path, &next)?;
        *self
            .session
            .write()
            .map_err(|_| connection_error("worker identity state is unavailable"))? = next.clone();
        let mut subscribers = self.subscribers.lock().await;
        subscribers.retain(|subscriber| {
            if let Some(token) = subscriber.token.upgrade() {
                if subscriber.session_id == current.session_id {
                    if let Ok(mut stored) = token.write() {
                        *stored = Some(value.clone());
                    }
                }
                true
            } else {
                false
            }
        });
        self.mark_ready(&next);
        Ok(())
    }

    async fn renew(&self, current: &Session) -> Result<()> {
        let mut pending_session = current.clone();
        if pending_session.pending_renewal.is_none() {
            let (private_key_pem, csr_der_base64) = new_key_and_csr()?;
            pending_session.pending_renewal = Some(PendingRenewal {
                request_id: uuid::Uuid::new_v4().to_string(),
                private_key_pem,
                csr_der_base64,
            });
            // Commit intent before the network call. A lost response or restart
            // must resend the same CSR and retain its private key.
            write_session(&self.session_path, &pending_session)?;
            *self
                .session
                .write()
                .map_err(|_| connection_error("worker identity state unavailable"))? =
                pending_session.clone();
        }
        let pending = pending_session.pending_renewal.as_ref().unwrap();
        let client = identity_http_client(current)?;
        let response = client
            .post(format!("{}{}", self.identity_url, SESSION_RENEW_PATH))
            .bearer_auth(&current.workload_token)
            .json(&RenewRequest {
                request_id: &pending.request_id,
                csr_der_base64: &pending.csr_der_base64,
            })
            .send()
            .await
            .map_err(|e| connection_error(format!("worker identity renewal failed: {e}")))?;
        if !response.status().is_success() {
            return Err(connection_error(format!(
                "worker identity renewal returned {}",
                response.status()
            )));
        }
        let mut next: Session = response
            .json()
            .await
            .map_err(|e| connection_error(format!("decode worker identity renewal: {e}")))?;
        next.private_key_pem = pending.private_key_pem.clone();
        next.pending_renewal = None;
        if next.runtime_endpoint.trim().is_empty() {
            next.runtime_endpoint = current.runtime_endpoint.clone();
        }
        validate_session(&next, &self.authority)?;
        write_session(&self.session_path, &next)?;
        *self
            .session
            .write()
            .map_err(|_| connection_error("worker identity state is unavailable"))? = next.clone();
        // A token bound to the replacement certificate must never be sent on
        // the old TLS channel. Removing authorization from current clients
        // makes the next pull request fail closed and enter the worker's
        // normal reconnect path, which obtains this replacement identity.
        let mut subscribers = self.subscribers.lock().await;
        subscribers.retain(|subscriber| {
            if let Some(token) = subscriber.token.upgrade() {
                if subscriber.session_id == current.session_id {
                    if let Ok(mut stored) = token.write() {
                        *stored = None;
                    }
                }
                true
            } else {
                false
            }
        });
        drop(subscribers);
        self.mark_ready(&next);
        warn!("external worker identity rotated; the runtime will reconnect the worker channel");
        Ok(())
    }

    fn mark_ready(&self, session: &Session) {
        if let Ok(mut readiness) = self.readiness.write() {
            *readiness = IdentityReadiness::Ready {
                certificate_expires_at: session.certificate_expires_at,
                token_expires_at: session.token_expires_at,
            };
        }
    }
}

async fn open_session(
    control_plane_url: &str,
    credential: &str,
    authority: &Authority,
    csr: &str,
) -> Result<Session> {
    let response = bootstrap_http_client()?
        .post(format!("{control_plane_url}{SESSION_OPEN_PATH}"))
        .bearer_auth(credential)
        .json(&OpenRequest {
            project_id: &authority.project_id,
            environment_id: &authority.environment_id,
            deployment_id: &authority.deployment_id,
            worker_pool_id: &authority.worker_pool_id,
            csr_der_base64: csr,
        })
        .send()
        .await
        .map_err(|e| connection_error(format!("worker identity enrollment failed: {e}")))?;
    if !response.status().is_success() {
        return Err(connection_error(format!(
            "worker identity enrollment returned {}",
            response.status()
        )));
    }
    response
        .json()
        .await
        .map_err(|e| connection_error(format!("decode worker identity enrollment: {e}")))
}

pub(crate) fn bootstrap_http_client() -> Result<reqwest::Client> {
    server_http_builder()?
        .build()
        .map_err(|e| connection_error(format!("create worker enrollment client: {e}")))
}

pub(crate) fn server_ca_pem() -> Result<Option<Vec<u8>>> {
    match std::env::var("AGNT5_WORKER_SERVER_CA_FILE").or_else(|e| {
        if matches!(e, std::env::VarError::NotPresent) {
            std::env::var("SSL_CERT_FILE")
        } else {
            Err(e)
        }
    }) {
        Ok(path) if !path.trim().is_empty() => {
            let bytes = std::fs::read(path.trim())
                .map_err(|_| connection_error("cannot read AGNT5_WORKER_SERVER_CA_FILE"))?;
            if bytes.is_empty() {
                return Err(connection_error("worker server CA file is empty"));
            }
            Ok(Some(bytes))
        }
        Ok(_) | Err(std::env::VarError::NotPresent) => Ok(None),
        Err(_) => Err(connection_error("worker server CA file path must be UTF-8")),
    }
}

fn server_http_builder() -> Result<reqwest::ClientBuilder> {
    let mut builder = reqwest::Client::builder()
        .timeout(Duration::from_secs(15))
        .redirect(reqwest::redirect::Policy::none());
    if let Some(pem) = server_ca_pem()? {
        let roots = reqwest::Certificate::from_pem_bundle(&pem)
            .map_err(|_| connection_error("invalid worker server CA PEM"))?;
        if roots.is_empty() {
            return Err(connection_error(
                "worker server CA file contains no certificates",
            ));
        }
        for root in roots {
            builder = builder.add_root_certificate(root);
        }
    }
    Ok(builder)
}

fn identity_http_client(session: &Session) -> Result<reqwest::Client> {
    let identity = reqwest::Identity::from_pem(&identity_pem(session)?)
        .map_err(|e| connection_error(format!("load worker client identity: {e}")))?;
    server_http_builder()?
        .identity(identity)
        .build()
        .map_err(|e| connection_error(format!("create worker mTLS client: {e}")))
}

fn tls_config(session: &Session) -> Result<ClientTlsConfig> {
    let mut tls = ClientTlsConfig::new()
        .with_webpki_roots()
        .identity(Identity::from_pem(
            certificate_chain_pem(session)?,
            session.private_key_pem.clone(),
        ));
    if let Some(pem) = server_ca_pem()? {
        tls = tls.ca_certificate(Certificate::from_pem(pem));
    }
    Ok(tls)
}

fn identity_pem(session: &Session) -> Result<Vec<u8>> {
    let mut bytes = certificate_chain_pem(session)?.into_bytes();
    bytes.extend_from_slice(session.private_key_pem.as_bytes());
    Ok(bytes)
}

fn certificate_chain_pem(session: &Session) -> Result<String> {
    let mut values = Vec::with_capacity(1 + session.certificate_chain_der_base64.len());
    values.push(session.certificate_der_base64.as_str());
    values.extend(
        session
            .certificate_chain_der_base64
            .iter()
            .map(String::as_str),
    );
    der_values_to_pem("CERTIFICATE", values)
}

fn der_values_to_pem<'a>(label: &str, values: impl IntoIterator<Item = &'a str>) -> Result<String> {
    let mut result = String::new();
    for value in values {
        let der = base64::engine::general_purpose::STANDARD
            .decode(value.trim())
            .map_err(|_| connection_error("worker session contained invalid certificate data"))?;
        result.push_str(&pem::encode(&pem::Pem::new(label, der)));
    }
    Ok(result)
}

fn new_key_and_csr() -> Result<(String, String)> {
    let key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)
        .map_err(|e| connection_error(format!("generate worker identity key: {e}")))?;
    let csr = CertificateParams::default()
        .serialize_request(&key)
        .map_err(|e| connection_error(format!("generate worker identity CSR: {e}")))?;
    Ok((
        key.serialize_pem(),
        base64::engine::general_purpose::STANDARD.encode(csr.der().as_ref()),
    ))
}

fn validate_session(session: &Session, authority: &Authority) -> Result<()> {
    validate_session_material(session, authority)?;
    if session.token_expires_at <= Utc::now() + ChronoDuration::seconds(30) {
        return Err(connection_error("worker token requires refresh"));
    }
    Ok(())
}

fn validate_session_material(session: &Session, authority: &Authority) -> Result<()> {
    if session.project_id != authority.project_id
        || session.environment_id != authority.environment_id
        || session.deployment_id != authority.deployment_id
        || session.worker_pool_id != authority.worker_pool_id
        || session.session_id.trim().is_empty()
        || session.worker_id.trim().is_empty()
        || session.private_key_pem.trim().is_empty()
        || session.workload_token.trim().is_empty()
        || session.certificate_der_base64.trim().is_empty()
        || session.trust_bundle_der_base64.is_empty()
        || session.certificate_expires_at <= Utc::now() + ChronoDuration::seconds(30)
    {
        return Err(connection_error(
            "stored worker identity is expired, incomplete, or outside discovery authority",
        ));
    }
    Ok(())
}

const AUTH_PROFILE_FILE: &str = "worker-auth-profile";

pub(crate) struct RolloutReadiness {
    pub(crate) ready: bool,
    pub(crate) pinned: bool,
    directory: PathBuf,
}

impl RolloutReadiness {
    pub(crate) fn from_env() -> Result<Self> {
        Self::prepare(
            &std::env::var("AGNT5_WORKER_MTLS_ENABLED").unwrap_or_default(),
            &std::env::var("AGNT5_WORKER_SESSION_DIR").unwrap_or_default(),
        )
    }

    fn prepare(enabled: &str, directory: &str) -> Result<Self> {
        let opted_in = match enabled.trim() {
            "" | "false" => false,
            "true" => true,
            _ => {
                return Err(connection_error(
                    "AGNT5_WORKER_MTLS_ENABLED must be true or false",
                ))
            }
        };
        let directory = PathBuf::from(directory.trim());
        let mut pinned = false;
        if !directory.as_os_str().is_empty() {
            for name in [AUTH_PROFILE_FILE, SESSION_FILE_NAME] {
                let path = directory.join(name);
                match std::fs::symlink_metadata(&path) {
                    Ok(_) => {
                        ensure_private_session_file(&path)?;
                        if name == AUTH_PROFILE_FILE {
                            let value = std::fs::read_to_string(&path).map_err(|e| {
                                connection_error(format!("read worker auth profile: {e}"))
                            })?;
                            if value != "bootstrap-mtls" {
                                return Err(connection_error(
                                    "invalid persisted worker auth profile",
                                ));
                            }
                        }
                        pinned = true;
                    }
                    Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
                    Err(e) => {
                        return Err(connection_error(format!(
                            "inspect worker identity state: {e}"
                        )))
                    }
                }
            }
        }
        let ready = opted_in || pinned;
        if ready {
            if directory.as_os_str().is_empty() {
                return Err(connection_error(
                    "AGNT5_WORKER_SESSION_DIR is required for mTLS readiness",
                ));
            }
            ensure_private_directory(&directory)?;
            let probe = directory.join(format!(".worker-readiness-{}", uuid::Uuid::new_v4()));
            write_private_file(&probe, b"")?;
            std::fs::remove_file(probe)
                .map_err(|e| connection_error(format!("remove worker readiness probe: {e}")))?;
        }
        Ok(Self {
            ready,
            pinned,
            directory,
        })
    }

    pub(crate) fn profiles(&self) -> Vec<&'static str> {
        if self.pinned {
            vec!["bootstrap-mtls"]
        } else if self.ready {
            vec!["bootstrap-mtls", "token-auth"]
        } else {
            vec!["token-auth"]
        }
    }

    pub(crate) fn current_profile(&self) -> &'static str {
        if self.pinned {
            "bootstrap-mtls"
        } else {
            ""
        }
    }

    pub(crate) fn accept(&self, profile: &str) -> Result<()> {
        if profile == "token-auth" {
            return if self.pinned {
                Err(connection_error(
                    "refusing to downgrade an enrolled mTLS worker to token-auth",
                ))
            } else {
                Ok(())
            };
        }
        if profile != "bootstrap-mtls" || !self.ready {
            return Err(connection_error(
                "discovery selected an authentication profile the worker is not ready to use",
            ));
        }
        // Persist selection before enrollment so retries and process restarts
        // cannot turn a failed enrollment into an implicit bearer fallback.
        let temporary = self
            .directory
            .join(format!(".worker-auth-profile-{}", uuid::Uuid::new_v4()));
        write_private_file(&temporary, b"bootstrap-mtls")?;
        if let Err(error) = std::fs::rename(&temporary, self.directory.join(AUTH_PROFILE_FILE)) {
            let _ = std::fs::remove_file(&temporary);
            return Err(connection_error(format!(
                "persist worker auth profile: {error}"
            )));
        }
        sync_directory(&self.directory)
    }
}

fn load_session(path: &Path, authority: &Authority) -> Result<Session> {
    ensure_private_session_file(path)?;
    let bytes = std::fs::read(path)
        .map_err(|e| connection_error(format!("load worker session {}: {e}", path.display())))?;
    let session: Session = serde_json::from_slice(&bytes)
        .map_err(|e| connection_error(format!("decode worker session: {e}")))?;
    validate_session_material(&session, authority)?;
    Ok(session)
}

fn write_session(path: &Path, session: &Session) -> Result<()> {
    let parent = path
        .parent()
        .ok_or_else(|| connection_error("worker session path has no parent"))?;
    let temporary = parent.join(format!(".{SESSION_FILE_NAME}.{}.tmp", uuid::Uuid::new_v4()));
    let bytes = serde_json::to_vec(session)
        .map_err(|e| connection_error(format!("encode worker session: {e}")))?;
    write_private_file(&temporary, &bytes)?;
    if let Err(error) = std::fs::rename(&temporary, path) {
        let _ = std::fs::remove_file(&temporary);
        return Err(connection_error(format!(
            "replace worker session atomically: {error}"
        )));
    }
    sync_directory(parent)?;
    Ok(())
}

fn ensure_private_directory(path: &Path) -> Result<()> {
    std::fs::create_dir_all(path)
        .map_err(|e| connection_error(format!("create worker session directory: {e}")))?;
    let metadata = std::fs::symlink_metadata(path)
        .map_err(|e| connection_error(format!("inspect worker session directory: {e}")))?;
    if metadata.file_type().is_symlink() || !metadata.is_dir() {
        return Err(connection_error(
            "worker session directory must be a real directory, not a symlink",
        ));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))
            .map_err(|e| connection_error(format!("secure worker session directory: {e}")))?;
    }
    Ok(())
}

fn write_private_file(path: &Path, bytes: &[u8]) -> Result<()> {
    use std::io::Write;
    let mut options = std::fs::OpenOptions::new();
    options.create_new(true).write(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.mode(0o600);
    }
    let mut file = options
        .open(path)
        .map_err(|e| connection_error(format!("write worker session: {e}")))?;
    file.write_all(bytes)
        .and_then(|_| file.sync_all())
        .map_err(|e| connection_error(format!("flush worker session: {e}")))
}

fn ensure_private_session_file(path: &Path) -> Result<()> {
    let metadata = std::fs::symlink_metadata(path)
        .map_err(|e| connection_error(format!("inspect worker session {}: {e}", path.display())))?;
    if metadata.file_type().is_symlink() || !metadata.is_file() {
        return Err(connection_error(
            "worker session must be a regular file, not a symlink",
        ));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        if metadata.permissions().mode() & 0o077 != 0 {
            return Err(connection_error(
                "worker session permissions must not allow group or other access",
            ));
        }
    }
    Ok(())
}

fn sync_directory(path: &Path) -> Result<()> {
    #[cfg(unix)]
    std::fs::File::open(path)
        .and_then(|directory| directory.sync_all())
        .map_err(|e| connection_error(format!("flush worker session directory: {e}")))?;
    Ok(())
}

fn bearer_metadata(token: &str) -> Result<MetadataValue<tonic::metadata::Ascii>> {
    format!("Bearer {token}")
        .parse()
        .map_err(|_| connection_error("worker token cannot be encoded as authorization metadata"))
}

pub(crate) fn validate_identity_endpoint(value: &str) -> Result<()> {
    let parsed = reqwest::Url::parse(value)
        .map_err(|e| connection_error(format!("invalid discovered identity endpoint: {e}")))?;
    let loopback = parsed.host_str().is_some_and(|host| {
        host == "localhost"
            || host
                .trim_matches(['[', ']'])
                .parse::<std::net::IpAddr>()
                .is_ok_and(|ip| ip.is_loopback())
    });
    if !(parsed.scheme() == "https" || parsed.scheme() == "http" && loopback)
        || parsed.host_str().is_none()
        || !parsed.username().is_empty()
        || parsed.password().is_some()
        || parsed.query().is_some()
        || parsed.fragment().is_some()
    {
        return Err(connection_error("worker endpoint requires HTTPS without URL credentials, query or fragment; HTTP is loopback-only"));
    }
    Ok(())
}

fn required_env(name: &str) -> Result<String> {
    let value = std::env::var(name)
        .map_err(|_| connection_error(format!("{name} is required for bootstrap identity")))?;
    if value.trim().is_empty() {
        return Err(connection_error(format!(
            "{name} is required for bootstrap identity"
        )));
    }
    Ok(value.trim().to_string())
}

fn connection_error(message: impl Into<String>) -> SdkError {
    SdkError::Connection {
        message: message.into(),
        code: ErrorCode::ConnectionFailed,
        source: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rolling_opt_in_survives_restart_and_prevents_downgrade() {
        let legacy = RolloutReadiness::prepare("", "").unwrap();
        assert_eq!(legacy.profiles(), vec!["token-auth"]);
        assert!(legacy.accept("bootstrap-mtls").is_err());
        assert!(legacy.accept("token-auth").is_ok());
        assert!(RolloutReadiness::prepare("true", "").is_err());
        assert!(RolloutReadiness::prepare("typo", "").is_err());
        let dir = tempfile::tempdir().unwrap();
        let directory = dir.path().to_str().unwrap();
        let ready = RolloutReadiness::prepare("true", directory).unwrap();
        assert!(ready.ready && !ready.pinned);
        ready.accept("bootstrap-mtls").unwrap();
        let restarted = RolloutReadiness::prepare("false", directory).unwrap();
        assert!(restarted.ready && restarted.pinned);
        assert_eq!(restarted.profiles(), vec!["bootstrap-mtls"]);
        assert_eq!(restarted.current_profile(), "bootstrap-mtls");
        assert!(restarted.accept("token-auth").is_err());
        std::fs::write(dir.path().join(AUTH_PROFILE_FILE), b"corrupt").unwrap();
        assert!(RolloutReadiness::prepare("", directory).is_err());
    }

    #[test]
    fn existing_session_pins_even_without_rollout_marker() {
        let dir = tempfile::tempdir().unwrap();
        write_private_file(&dir.path().join(SESSION_FILE_NAME), b"old SDK session").unwrap();
        let ready = RolloutReadiness::prepare("", dir.path().to_str().unwrap()).unwrap();
        assert!(ready.ready && ready.pinned);
        assert!(ready.accept("token-auth").is_err());
    }

    fn authority() -> Authority {
        Authority {
            project_id: "project".into(),
            environment_id: "environment".into(),
            deployment_id: "deployment".into(),
            worker_pool_id: "pool".into(),
            runtime_endpoint: "https://runtime.example.com".into(),
            protocol: "pull.v1".into(),
        }
    }

    #[test]
    fn generated_csr_uses_profile_key_algorithm() {
        let (key, csr) = new_key_and_csr().expect("key and csr");
        assert!(key.contains("PRIVATE KEY"));
        assert!(!csr.is_empty());
    }

    #[test]
    fn session_validation_rejects_cross_authority_and_expiry() {
        let mut session = Session {
            session_id: "session".into(),
            project_id: "project".into(),
            environment_id: "environment".into(),
            deployment_id: "deployment".into(),
            worker_pool_id: "pool".into(),
            worker_id: "worker".into(),
            spiffe_id: "spiffe://example/workload/project/environment/deployment/worker".into(),
            runtime_endpoint: "https://runtime.example.com".into(),
            certificate_der_base64: "Y2VydA==".into(),
            certificate_chain_der_base64: vec![],
            trust_bundle_der_base64: vec!["cm9vdA==".into()],
            trust_bundle_version: "v1".into(),
            certificate_expires_at: Utc::now() + ChronoDuration::minutes(10),
            renew_after: Utc::now() + ChronoDuration::minutes(5),
            workload_token: "token".into(),
            token_type: "Bearer".into(),
            token_expires_at: Utc::now() + ChronoDuration::minutes(5),
            private_key_pem: "key".into(),
            pending_renewal: None,
        };
        assert!(validate_session(&session, &authority()).is_ok());
        session.project_id = "other".into();
        assert!(validate_session(&session, &authority()).is_err());
    }
    #[tokio::test]
    async fn renewal_reuses_persisted_intent_after_response_loss_and_restart() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join(SESSION_FILE_NAME);
        let key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap();
        let cert = CertificateParams::default().self_signed(&key).unwrap();
        let current = Session {
            session_id: "predecessor".into(),
            project_id: "project".into(),
            environment_id: "environment".into(),
            deployment_id: "deployment".into(),
            worker_pool_id: "pool".into(),
            worker_id: "worker".into(),
            spiffe_id: "spiffe://example/workload/project/environment/deployment/worker".into(),
            runtime_endpoint: "https://runtime.example.com".into(),
            certificate_der_base64: base64::engine::general_purpose::STANDARD.encode(cert.der()),
            certificate_chain_der_base64: vec![],
            trust_bundle_der_base64: vec![
                base64::engine::general_purpose::STANDARD.encode(cert.der())
            ],
            trust_bundle_version: "v1".into(),
            certificate_expires_at: Utc::now() + ChronoDuration::hours(1),
            renew_after: Utc::now() - ChronoDuration::minutes(1),
            workload_token: "old-token".into(),
            token_type: "Bearer".into(),
            token_expires_at: Utc::now() - ChronoDuration::minutes(1),
            private_key_pem: key.serialize_pem(),
            pending_renewal: None,
        };
        write_session(&path, &current).unwrap();
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let endpoint = format!("http://{}", listener.local_addr().unwrap());
        let server_path = path.clone();
        let server = tokio::spawn(async move {
            let mut original = None;
            for attempt in 0..2 {
                let (mut stream, _) = listener.accept().await.unwrap();
                let mut bytes = Vec::new();
                let request: serde_json::Value = loop {
                    let mut chunk = [0; 4096];
                    let n = stream.read(&mut chunk).await.unwrap();
                    assert!(n > 0);
                    bytes.extend_from_slice(&chunk[..n]);
                    assert!(bytes.len() < 64 * 1024);
                    if let Some(header_end) = bytes.windows(4).position(|w| w == b"\r\n\r\n") {
                        let header = String::from_utf8_lossy(&bytes[..header_end]);
                        let length: usize = header
                            .lines()
                            .find_map(|line| {
                                let (k, v) = line.split_once(':')?;
                                k.eq_ignore_ascii_case("content-length")
                                    .then(|| v.trim().parse().unwrap())
                            })
                            .unwrap();
                        if bytes.len() >= header_end + 4 + length {
                            break serde_json::from_slice(
                                &bytes[header_end + 4..header_end + 4 + length],
                            )
                            .unwrap();
                        }
                    }
                };
                let stored = load_session(&server_path, &authority()).unwrap();
                let pending = stored
                    .pending_renewal
                    .as_ref()
                    .expect("intent must predate network request");
                assert_eq!(request["request_id"], pending.request_id);
                assert_eq!(request["csr_der_base64"], pending.csr_der_base64);
                let identity = (
                    pending.request_id.clone(),
                    pending.csr_der_base64.clone(),
                    pending.private_key_pem.clone(),
                );
                if attempt == 0 {
                    original = Some(identity);
                    drop(stream);
                    continue;
                }
                assert_eq!(original.as_ref().unwrap(), &identity);
                let key = KeyPair::from_pem(&pending.private_key_pem).unwrap();
                let cert = CertificateParams::default().self_signed(&key).unwrap();
                let mut next = stored.clone();
                next.session_id = "successor".into();
                next.pending_renewal = None;
                next.private_key_pem = String::new();
                next.certificate_der_base64 =
                    base64::engine::general_purpose::STANDARD.encode(cert.der());
                next.token_expires_at = Utc::now() + ChronoDuration::minutes(10);
                let body = serde_json::to_vec(&next).unwrap();
                let header = format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",body.len());
                stream.write_all(header.as_bytes()).await.unwrap();
                stream.write_all(&body).await.unwrap();
            }
            original.unwrap().2
        });
        let manager = |session: Session| IdentityManager {
            identity_url: endpoint.clone(),
            authority: authority(),
            session_path: path.clone(),
            session: RwLock::new(session),
            readiness: RwLock::new(IdentityReadiness::Degraded),
            subscribers: Mutex::new(Vec::new()),
        };
        let first = manager(current.clone());
        assert!(first.renew(&current).await.is_err());
        drop(first);
        let restarted_session = load_session(&path, &authority()).unwrap();
        let restarted = manager(restarted_session.clone());
        restarted.renew(&restarted_session).await.unwrap();
        let original_key = server.await.unwrap();
        let final_session = load_session(&path, &authority()).unwrap();
        assert_eq!(final_session.session_id, "successor");
        assert!(final_session.pending_renewal.is_none());
        assert_eq!(final_session.private_key_pem, original_key);
    }
}