dex_sdk/

credential.rs

// Copyright (c) 2026 Super Durable, Inc.
//
// Licensed under the Sustainable Use License 1.0.
// You may not use this file except in compliance with the License.
// See the LICENSE file in the repository root.
//
// SPDX-License-Identifier: LicenseRef-Sustainable-Use-1.0

use std::error::Error;
use std::fmt;
use std::fs;
use std::future::Future;
use std::io;
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant, SystemTime};

use bytes::Bytes;
use http_body::{Body as HttpBody, Frame};
use http_body_util::BodyExt;
use tonic::body::Body;
use tonic::codegen::{BoxFuture, Service, http};
use tonic::metadata::MetadataValue;
use tonic::service::Interceptor;
use tonic::service::interceptor::{InterceptedService, ResponseBody};
use tonic::transport::Channel;
use tonic::{Request, Status};

type RejectedTokenRefreshFuture<'a> =
    Pin<Box<dyn Future<Output = Result<String, CredentialError>> + Send + 'a>>;

#[cfg(feature = "api-key-credential")]
#[path = "credential/api_key_credential.rs"]
mod api_key_credential;
#[cfg(feature = "api-key-credential")]
pub use api_key_credential::{ApiKeyCredential, ApiKeyCredentialBuilder, ApiKeyCredentialError};

/// Supplies the current bearer token before each Dex request.
///
/// Return the bare token without a `Bearer ` prefix. Implementations may be called concurrently.
/// The synchronous call can run on a Tokio worker thread; it must return promptly without network
/// I/O. A call deadline cannot interrupt a blocked implementation. Sources requiring slow I/O
/// should refresh in the background and return a cached token. An error fails the request locally
/// with gRPC `Unauthenticated`. Pull Worker streams evaluate the token when each stream opens; a
/// failed open normally backs off for about one minute. An API-key credential can reconnect once
/// immediately after a token rejection. The token travels over Dex's plaintext gRPC connection,
/// so protect the network path.
///
/// ```no_run
/// use std::sync::{Arc, RwLock};
/// use dex_sdk::{ClientOptions, Credential, CredentialError};
/// struct RotatingCredential { current: RwLock<Option<String>> }
/// impl Credential for RotatingCredential {
///     fn token(&self) -> Result<String, CredentialError> {
///         self.current.read()
///             .map_err(|_| CredentialError::new("credential lock poisoned"))?
///             .clone()
///             .ok_or_else(|| CredentialError::new("token unavailable"))
///     }
/// }
/// let credential = Arc::new(RotatingCredential {
///     current: RwLock::new(Some("example-token".to_string())),
/// });
/// let client_options = ClientOptions::new().credential(credential);
/// ```
pub trait Credential: Send + Sync {
    /// Returns the current bare token, or an error that fails the request locally.
    ///
    /// Implementations must be safe for concurrent synchronous calls and return promptly. Do not
    /// include the `Bearer ` prefix. A credential failure maps to gRPC `Unauthenticated` without
    /// contacting Dex.
    fn token(&self) -> Result<String, CredentialError>;

    /// Discards a token rejected by Dex if it is still the credential's cached token.
    ///
    /// The default does nothing, keeping existing credential implementations compatible.
    /// Sources that cache tokens can compare `rejected_token` with their current token before
    /// invalidating it. Worker streams call this on rejection and reconnect once immediately when
    /// refresh is supported. The SDK never logs this value.
    fn discard_rejected_token(&self, _rejected_token: &str) {}

    /// Supplies an asynchronous fresh token after rejection when supported.
    ///
    /// The default returns `None`, so existing credentials continue handling rejection as before.
    /// An API-key credential returns a future bounded by its exchange timeout. The SDK uses
    /// this for one unary retry or one immediate Worker stream reconnect after
    /// [`Self::discard_rejected_token`]. Implementations must avoid returning the rejected token.
    /// Errors fail the attempt without exposing token contents.
    fn token_after_rejection(&self) -> Option<RejectedTokenRefreshFuture<'_>> {
        None
    }
}

impl fmt::Debug for dyn Credential {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("Credential")
    }
}

/// An error obtaining a bearer token before a Dex request.
///
/// Return this from a custom [`Credential`] when no usable token is available. Avoid including
/// token contents in the error message.
///
/// ```
/// use dex_sdk::CredentialError;
/// let error = CredentialError::new("token unavailable");
/// assert_eq!(error.to_string(), "token unavailable");
/// ```
#[derive(Debug, Clone)]
pub struct CredentialError {
    message: String,
}

impl CredentialError {
    /// Creates an error for a credential source failure.
    ///
    /// `message` is shown in the local service error and must not contain a token. See the
    /// [`CredentialError`] example.
    pub fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for CredentialError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl Error for CredentialError {}

/// Reads a token file at construction and after its modification time or size changes.
///
/// File metadata is checked at most once per second. Surrounding whitespace is removed. If the
/// file becomes empty or unreadable, requests fail until a later check succeeds; no stale token
/// is sent. The constructor fails when the initial file is unreadable or empty. Checks and reads
/// use a mutex, briefly serializing concurrent callers. Use a local filesystem; slow network
/// filesystems can delay calls. Replace the token atomically by writing a temporary file and
/// renaming it. An empty or partial token can fail pull stream opens, which back off for about one
/// minute.
///
/// ```no_run
/// use std::sync::Arc;
/// use dex_sdk::{ClientOptions, FileCredential, WorkerOptions};
/// let credential = Arc::new(FileCredential::new("/var/run/secrets/dex/token")?);
/// let client_options = ClientOptions::new().credential(credential.clone());
/// let worker_options = WorkerOptions::new().credential(credential);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub struct FileCredential {
    path: PathBuf,
    state: Mutex<FileCredentialState>,
}

struct FileCredentialState {
    token: Result<String, CredentialError>,
    modified: SystemTime,
    size: u64,
    checked_at: Instant,
}

impl fmt::Debug for FileCredential {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("FileCredential")
            .field("path", &self.path)
            .field("token", &"<redacted>")
            .finish()
    }
}

impl FileCredential {
    /// Opens `path` and reads its initial nonempty token.
    ///
    /// # Errors
    ///
    /// Returns an I/O error if the file cannot be inspected or read, or contains only whitespace.
    /// See the [`FileCredential`] example for sharing the credential with a Client and Worker.
    pub fn new(path: impl AsRef<Path>) -> io::Result<Self> {
        let path = path.as_ref().to_path_buf();
        let metadata = fs::metadata(&path)?;
        let token = read_token(&path)?;
        Ok(Self {
            path,
            state: Mutex::new(FileCredentialState {
                token: Ok(token),
                modified: metadata.modified()?,
                size: metadata.len(),
                checked_at: Instant::now(),
            }),
        })
    }
}

impl Credential for FileCredential {
    fn token(&self) -> Result<String, CredentialError> {
        let mut state = self.state.lock().expect("credential file lock poisoned");
        if state.checked_at.elapsed() >= Duration::from_secs(1) {
            state.checked_at = Instant::now();
            state.token = match fs::metadata(&self.path) {
                Ok(metadata) => match metadata.modified() {
                    Ok(modified)
                        if state.token.is_ok()
                            && modified == state.modified
                            && metadata.len() == state.size =>
                    {
                        state.token.clone()
                    }
                    Ok(modified) => {
                        let token = read_token(&self.path).map_err(credential_file_error);
                        if token.is_ok() {
                            state.modified = modified;
                            state.size = metadata.len();
                        }
                        token
                    }
                    Err(error) => Err(credential_file_error(error)),
                },
                Err(error) => Err(credential_file_error(error)),
            };
        }
        state.token.clone()
    }
}

fn read_token(path: &Path) -> io::Result<String> {
    let token = fs::read_to_string(path)?.trim().to_string();
    if token.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("credential file {} is empty", path.display()),
        ));
    }
    Ok(token)
}

fn credential_file_error(error: io::Error) -> CredentialError {
    CredentialError::new(format!("read credential file: {error}"))
}

#[derive(Clone)]
pub(crate) struct CredentialChannel {
    service: InterceptedService<RetryChannel, CredentialInterceptor>,
    credential: Option<Arc<dyn Credential>>,
}

impl CredentialChannel {
    pub(crate) fn new(channel: Channel, interceptor: CredentialInterceptor) -> Self {
        let credential = interceptor.credential.clone();
        let retry_channel = RetryChannel {
            channel,
            credential: credential.clone(),
        };
        Self {
            service: InterceptedService::new(retry_channel, interceptor),
            credential,
        }
    }

    pub(crate) fn credential(&self) -> Option<Arc<dyn Credential>> {
        self.credential.clone()
    }
}

impl Service<http::Request<Body>> for CredentialChannel {
    type Response = http::Response<ResponseBody<Body>>;
    type Error = tonic::transport::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(context)
    }

    fn call(&mut self, request: http::Request<Body>) -> Self::Future {
        Box::pin(self.service.call(request))
    }
}

#[derive(Clone)]
struct RetryChannel {
    channel: Channel,
    credential: Option<Arc<dyn Credential>>,
}

impl Service<http::Request<Body>> for RetryChannel {
    type Response = http::Response<Body>;
    type Error = tonic::transport::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn poll_ready(&mut self, _context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, request: http::Request<Body>) -> Self::Future {
        let mut channel = self.channel.clone();
        let credential = self.credential.clone();
        let deadline = request_deadline(&request);
        Box::pin(async move {
            let attempt = async {
                let mut request = request;
                let mut rejected_token = match request.headers().get("authorization") {
                    Some(value) => match value.to_str() {
                        Ok(value) => value.strip_prefix("Bearer ").map(str::to_owned),
                        Err(_) => None,
                    },
                    None => None,
                };
                if rejected_token.is_none()
                    && let Some(credential) = &credential
                    && let Some(fresh_token) = credential.token_after_rejection()
                {
                    let fresh_token = match fresh_token.await {
                        Ok(token) => token,
                        Err(error) => {
                            return Ok(Status::unauthenticated(error.to_string()).into_http());
                        }
                    };
                    let authorization = match bearer_header(&fresh_token) {
                        Ok(value) => value,
                        Err(status) => return Ok(status.into_http()),
                    };
                    request.headers_mut().insert("authorization", authorization);
                    rejected_token = Some(fresh_token);
                }
                let is_poll_stream =
                    request.uri().path() == "/dex.WorkerTaskService/PollWorkerTask";
                if is_poll_stream || credential.is_none() {
                    let response = send_channel_request(&mut channel, request).await?;
                    let Some(credential) = credential else {
                        return Ok(response);
                    };
                    return Ok(observe_stream_rejection(
                        response,
                        credential,
                        rejected_token,
                    ));
                }
                let (parts, body) = request.into_parts();
                let request_body = match body.collect().await {
                    Ok(body) => body,
                    Err(status) => return Ok(status.into_http()),
                };
                let request_trailers = request_body.trailers().cloned();
                let request_bytes = request_body.to_bytes();
                let retry_method = parts.method.clone();
                let retry_uri = parts.uri.clone();
                let retry_version = parts.version;
                let retry_headers = parts.headers.clone();
                let request = http::Request::from_parts(
                    parts,
                    buffered_body(request_bytes.clone(), request_trailers.clone()),
                );
                let response = send_channel_request(&mut channel, request).await?;
                let (response_parts, response_body) = response.into_parts();
                let response_body = match response_body.collect().await {
                    Ok(body) => body,
                    Err(status) => return Ok(status.into_http()),
                };
                let response_trailers = response_body.trailers().cloned();
                let is_rejected = is_unauthenticated(&response_parts.headers)
                    || response_trailers.as_ref().is_some_and(is_unauthenticated);
                if !is_rejected {
                    return Ok(http::Response::from_parts(
                        response_parts,
                        buffered_body(response_body.to_bytes(), response_trailers),
                    ));
                }
                let Some(credential) = credential else {
                    return Ok(http::Response::from_parts(
                        response_parts,
                        buffered_body(response_body.to_bytes(), response_trailers),
                    ));
                };
                let Some(rejected_token) = rejected_token else {
                    return Ok(http::Response::from_parts(
                        response_parts,
                        buffered_body(response_body.to_bytes(), response_trailers),
                    ));
                };
                credential.discard_rejected_token(&rejected_token);
                let Some(fresh_token) = credential.token_after_rejection() else {
                    return Ok(http::Response::from_parts(
                        response_parts,
                        buffered_body(response_body.to_bytes(), response_trailers),
                    ));
                };
                let fresh_token = match fresh_token.await {
                    Ok(token) => token,
                    Err(error) => return Ok(Status::unauthenticated(error.to_string()).into_http()),
                };
                let authorization = match bearer_header(&fresh_token) {
                    Ok(value) => value,
                    Err(status) => return Ok(status.into_http()),
                };
                let mut retry = http::Request::new(buffered_body(request_bytes, request_trailers));
                *retry.method_mut() = retry_method;
                *retry.uri_mut() = retry_uri;
                *retry.version_mut() = retry_version;
                *retry.headers_mut() = retry_headers;
                retry.headers_mut().insert("authorization", authorization);
                if let Some(deadline) = deadline
                    && !set_remaining_timeout(&mut retry, deadline)
                {
                    return Ok(Status::deadline_exceeded("request deadline exceeded").into_http());
                }
                let response = send_channel_request(&mut channel, retry).await?;
                let (parts, body) = response.into_parts();
                let response = observe_stream_rejection(
                    http::Response::from_parts(parts, body),
                    credential,
                    Some(fresh_token),
                );
                Ok(response)
            };
            if let Some(deadline) = deadline {
                match tokio::time::timeout_at(deadline.into(), attempt).await {
                    Ok(result) => result,
                    Err(_) => {
                        Ok(Status::deadline_exceeded("request deadline exceeded").into_http())
                    }
                }
            } else {
                attempt.await
            }
        })
    }
}

fn bearer_header(token: &str) -> Result<http::HeaderValue, Status> {
    http::HeaderValue::from_str(&format!("Bearer {token}"))
        .map_err(|_| Status::unauthenticated("invalid bearer token metadata"))
}

fn request_deadline(request: &http::Request<Body>) -> Option<Instant> {
    let timeout = request.headers().get("grpc-timeout")?.to_str().ok()?;
    let (amount, unit) = timeout.split_at(timeout.len().checked_sub(1)?);
    if amount.is_empty() || amount.len() > 8 {
        return None;
    }
    let amount = amount.parse::<u64>().ok()?;
    let duration = match unit {
        "H" => Duration::from_secs(amount.checked_mul(3_600)?),
        "M" => Duration::from_secs(amount.checked_mul(60)?),
        "S" => Duration::from_secs(amount),
        "m" => Duration::from_millis(amount),
        "u" => Duration::from_micros(amount),
        "n" => Duration::from_nanos(amount),
        _ => return None,
    };
    Instant::now().checked_add(duration)
}

fn set_remaining_timeout(request: &mut http::Request<Body>, deadline: Instant) -> bool {
    let remaining = deadline.saturating_duration_since(Instant::now());
    if remaining.is_zero() {
        return false;
    }
    let nanoseconds = remaining.as_nanos();
    let timeout = [
        (1, "n"),
        (1_000, "u"),
        (1_000_000, "m"),
        (1_000_000_000, "S"),
        (60_000_000_000, "M"),
        (3_600_000_000_000, "H"),
    ]
    .into_iter()
    .find_map(|(unit, suffix)| {
        let amount = nanoseconds.div_ceil(unit);
        (amount <= 99_999_999).then(|| format!("{amount}{suffix}"))
    })
    .expect("gRPC timeout fits an eight-digit hour value");
    request.headers_mut().insert(
        "grpc-timeout",
        http::HeaderValue::from_str(&timeout).expect("numeric gRPC timeout header"),
    );
    true
}

async fn send_channel_request(
    channel: &mut Channel,
    request: http::Request<Body>,
) -> Result<http::Response<Body>, tonic::transport::Error> {
    std::future::poll_fn(|context| channel.poll_ready(context)).await?;
    channel.call(request).await
}

fn is_unauthenticated(headers: &http::HeaderMap) -> bool {
    headers
        .get("grpc-status")
        .is_some_and(|status| status == "16")
}

fn observe_stream_rejection(
    response: http::Response<Body>,
    credential: Arc<dyn Credential>,
    rejected_token: Option<String>,
) -> http::Response<Body> {
    let (parts, body) = response.into_parts();
    if is_unauthenticated(&parts.headers)
        && let Some(token) = &rejected_token
    {
        credential.discard_rejected_token(token);
    }
    http::Response::from_parts(
        parts,
        Body::new(ObservedBody {
            body,
            credential,
            rejected_token,
        }),
    )
}

struct ObservedBody {
    body: Body,
    credential: Arc<dyn Credential>,
    rejected_token: Option<String>,
}

impl HttpBody for ObservedBody {
    type Data = Bytes;
    type Error = Status;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        context: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
        let frame = std::task::ready!(Pin::new(&mut self.body).poll_frame(context));
        if let Some(Ok(frame)) = &frame
            && frame.trailers_ref().is_some_and(is_unauthenticated)
            && let Some(token) = &self.rejected_token
        {
            self.credential.discard_rejected_token(token);
        }
        Poll::Ready(frame)
    }
}

fn buffered_body(data: Bytes, trailers: Option<http::HeaderMap>) -> Body {
    Body::new(BufferedBody {
        data: Some(data),
        trailers,
    })
}

struct BufferedBody {
    data: Option<Bytes>,
    trailers: Option<http::HeaderMap>,
}

impl HttpBody for BufferedBody {
    type Data = Bytes;
    type Error = Status;

    fn poll_frame(
        mut self: Pin<&mut Self>,
        _context: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
        if let Some(data) = self.data.take() {
            return Poll::Ready(Some(Ok(Frame::data(data))));
        }
        Poll::Ready(
            self.trailers
                .take()
                .map(|trailers| Ok(Frame::trailers(trailers))),
        )
    }
}

#[derive(Clone)]
pub(crate) struct CredentialInterceptor {
    credential: Option<Arc<dyn Credential>>,
}

impl CredentialInterceptor {
    pub(crate) fn new(credential: Option<Arc<dyn Credential>>) -> Self {
        Self { credential }
    }
}

impl Interceptor for CredentialInterceptor {
    fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
        if let Some(credential) = &self.credential {
            match credential.token() {
                Ok(token) => {
                    let value = MetadataValue::try_from(format!("Bearer {token}"))
                        .map_err(|_| Status::unauthenticated("invalid bearer token metadata"))?;
                    request.metadata_mut().insert("authorization", value);
                }
                Err(_) if credential.token_after_rejection().is_some() => {}
                Err(error) => return Err(Status::unauthenticated(error.to_string())),
            }
        }
        Ok(request)
    }
}

#[cfg(test)]
mod tests {
    use std::convert::Infallible;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::task::{Context, Poll};

    use dex_protocol::dex::worker_task_service_client::WorkerTaskServiceClient;
    use dex_protocol::dex::worker_task_service_server::{
        WorkerTaskService, WorkerTaskServiceServer,
    };
    use dex_protocol::dex::{
        HealthInfo, PollWorkerTaskRequest, PollWorkerTaskResponse, ServerInfo,
    };
    use tokio::net::TcpListener;
    use tokio::sync::{mpsc, oneshot};
    use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream};
    use tonic::codegen::{BoxFuture, Service, http};
    use tonic::server::{NamedService, UnaryService};
    use tonic::transport::{Endpoint, Server};
    use tonic::{Response, Streaming};

    use super::*;
    use crate::{
        BlobCache, BlobCacheConfig, Client, ClientOptions, PullWorkerOptions, Registry, Worker,
        WorkerOptions,
    };

    #[derive(Clone)]
    struct FakeFlowService {
        headers: Arc<Mutex<Vec<Option<String>>>>,
    }

    impl NamedService for FakeFlowService {
        const NAME: &'static str = "dex.FlowService";
    }

    impl<B> Service<http::Request<B>> for FakeFlowService
    where
        B: tonic::codegen::Body + Send + 'static,
        B::Error: Into<tonic::codegen::StdError> + Send + 'static,
    {
        type Response = http::Response<tonic::body::Body>;
        type Error = Infallible;
        type Future = BoxFuture<Self::Response, Self::Error>;

        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, request: http::Request<B>) -> Self::Future {
            let is_server_info = request.uri().path() == "/dex.FlowService/GetServerInfo";
            let header = request
                .headers()
                .get("authorization")
                .map(|value| value.to_str().expect("ASCII bearer header").to_string());
            self.headers.lock().expect("headers lock").push(header);
            Box::pin(async move {
                if is_server_info {
                    let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
                    Ok(grpc.unary(GetServerInfo, request).await)
                } else {
                    let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
                    Ok(grpc.unary(HealthCheck, request).await)
                }
            })
        }
    }

    struct HealthCheck;

    impl UnaryService<()> for HealthCheck {
        type Response = HealthInfo;
        type Future = BoxFuture<Response<Self::Response>, Status>;

        fn call(&mut self, _: Request<()>) -> Self::Future {
            Box::pin(async { Ok(Response::new(HealthInfo::default())) })
        }
    }

    struct GetServerInfo;

    impl UnaryService<()> for GetServerInfo {
        type Response = ServerInfo;
        type Future = BoxFuture<Response<Self::Response>, Status>;

        fn call(&mut self, _: Request<()>) -> Self::Future {
            Box::pin(async { Ok(Response::new(ServerInfo::default())) })
        }
    }

    #[derive(Clone)]
    struct FakeTaskService {
        headers: Arc<Mutex<Vec<Option<String>>>>,
    }

    #[tonic::async_trait]
    impl WorkerTaskService for FakeTaskService {
        type PollWorkerTaskStream = ReceiverStream<Result<PollWorkerTaskResponse, Status>>;

        async fn poll_worker_task(
            &self,
            request: Request<Streaming<PollWorkerTaskRequest>>,
        ) -> Result<Response<Self::PollWorkerTaskStream>, Status> {
            let header = request
                .metadata()
                .get("authorization")
                .map(|value| value.to_str().expect("ASCII bearer header").to_string());
            self.headers.lock().expect("headers lock").push(header);
            let (_sender, receiver) = mpsc::channel(1);
            Ok(Response::new(ReceiverStream::new(receiver)))
        }
    }

    struct MutableCredential {
        token: Mutex<Result<String, CredentialError>>,
        calls: AtomicUsize,
    }

    impl Credential for MutableCredential {
        fn token(&self) -> Result<String, CredentialError> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            self.token.lock().expect("token lock").clone()
        }
    }

    #[test]
    fn file_credential_rotates_and_fails_closed() {
        let directory = tempfile::tempdir().expect("token directory");
        let path = directory.path().join("token");
        assert!(FileCredential::new(&path).is_err());
        fs::write(&path, " \n").expect("write empty token");
        assert!(FileCredential::new(&path).is_err());
        fs::write(&path, "  first\n").expect("write initial token");
        let credential = FileCredential::new(&path).expect("initial token");
        assert_eq!(credential.token().expect("initial token"), "first");
        fs::write(&path, " second-token \n").expect("rotate token");
        assert_eq!(credential.token().expect("cached token"), "first");
        credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
        assert_eq!(credential.token().expect("rotated token"), "second-token");
        fs::write(&path, " \n").expect("empty token");
        credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
        assert!(credential.token().is_err());
        fs::remove_file(&path).expect("remove token");
        credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
        assert!(credential.token().is_err());
        fs::write(&path, "restored").expect("restore token");
        credential.state.lock().expect("file state lock").checked_at -= Duration::from_secs(1);
        assert_eq!(credential.token().expect("restored token"), "restored");
    }

    #[test]
    fn file_credential_debug_redacts_token() {
        let directory = tempfile::tempdir().expect("token directory");
        let path = directory.path().join("token");
        fs::write(&path, "private-token-value").expect("write token");
        let credential = Arc::new(FileCredential::new(&path).expect("file credential"));
        let debug = format!("{credential:?}");
        assert!(!debug.contains("private-token-value"));
        assert!(debug.contains("<redacted>"));
    }

    #[test]
    fn flow_service_unary_uses_current_credential_and_default_has_no_header() {
        let runtime = tokio::runtime::Runtime::new().expect("server runtime");
        let listener = runtime
            .block_on(TcpListener::bind("127.0.0.1:0"))
            .expect("bind fake FlowService");
        let address = listener.local_addr().expect("fake server address");
        let headers = Arc::new(Mutex::new(Vec::new()));
        let (shutdown, stopped) = oneshot::channel::<()>();
        runtime.spawn(
            Server::builder()
                .add_service(FakeFlowService {
                    headers: Arc::clone(&headers),
                })
                .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
                    stopped.await.expect("FlowService shutdown signal");
                }),
        );
        let credential = Arc::new(MutableCredential {
            token: Mutex::new(Ok("first".to_string())),
            calls: AtomicUsize::new(0),
        });
        let directory = tempfile::tempdir().expect("blob cache directory");
        let cache = Arc::new(
            BlobCache::open(
                BlobCacheConfig::new(directory.path(), 1024 * 1024, 0).expect("cache config"),
            )
            .expect("open cache"),
        );
        let without_credential = Client::try_new(
            Registry::new(),
            Arc::clone(&cache),
            ClientOptions::new().server_address(address.to_string()),
        )
        .expect("default Client");
        without_credential.health_check().expect("default call");
        let with_credential = Client::try_new(
            Registry::new(),
            Arc::clone(&cache),
            ClientOptions::new()
                .server_address(address.to_string())
                .credential(credential.clone()),
        )
        .expect("authenticated Client");
        with_credential.health_check().expect("first call");
        *credential.token.lock().expect("token lock") = Ok("second".to_string());
        with_credential.health_check().expect("rotated call");
        *credential.token.lock().expect("token lock") =
            Err(CredentialError::new("token unavailable"));
        let error = with_credential.health_check().expect_err("local error");
        assert_eq!(
            error.service_error().expect("service error").code(),
            tonic::Code::Unauthenticated
        );
        *credential.token.lock().expect("token lock") = Ok("worker".to_string());
        let worker = Worker::try_new(
            Registry::new(),
            cache,
            WorkerOptions::new()
                .server_address(address.to_string())
                .credential(credential.clone())
                .pull(PullWorkerOptions::new()),
        )
        .expect("authenticated Worker");
        worker
            .start()
            .expect_err("fake protocol interval is invalid");
        assert_eq!(credential.calls.load(Ordering::SeqCst), 4);
        assert_eq!(
            *headers.lock().expect("headers lock"),
            vec![
                None,
                Some("Bearer first".to_string()),
                Some("Bearer second".to_string()),
                Some("Bearer worker".to_string())
            ]
        );
        shutdown.send(()).expect("stop fake FlowService");
    }

    #[tokio::test(flavor = "multi_thread")]
    async fn pull_stream_open_reevaluates_credential() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("bind fake WorkerTaskService");
        let address = listener.local_addr().expect("fake server address");
        let headers = Arc::new(Mutex::new(Vec::new()));
        let (shutdown, stopped) = oneshot::channel::<()>();
        let server = tokio::spawn(
            Server::builder()
                .add_service(WorkerTaskServiceServer::new(FakeTaskService {
                    headers: Arc::clone(&headers),
                }))
                .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
                    stopped.await.expect("WorkerTaskService shutdown signal");
                }),
        );
        let credential = Arc::new(MutableCredential {
            token: Mutex::new(Ok("first".to_string())),
            calls: AtomicUsize::new(0),
        });
        let channel = Endpoint::from_shared(format!("http://{address}"))
            .expect("endpoint")
            .connect()
            .await
            .expect("connect fake WorkerTaskService");
        let mut default_client = WorkerTaskServiceClient::new(CredentialChannel::new(
            channel.clone(),
            CredentialInterceptor::new(None),
        ));
        default_client
            .poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
            .await
            .expect("open default task stream");
        let mut client = WorkerTaskServiceClient::new(CredentialChannel::new(
            channel,
            CredentialInterceptor::new(Some(credential.clone())),
        ));
        for token in ["first", "second"] {
            *credential.token.lock().expect("token lock") = Ok(token.to_string());
            client
                .poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
                .await
                .expect("open task stream");
        }
        *credential.token.lock().expect("token lock") =
            Err(CredentialError::new("token unavailable"));
        let error = client
            .poll_worker_task(tokio_stream::empty::<PollWorkerTaskRequest>())
            .await
            .expect_err("local task stream error");
        assert_eq!(error.code(), tonic::Code::Unauthenticated);
        assert_eq!(credential.calls.load(Ordering::SeqCst), 3);
        assert_eq!(
            *headers.lock().expect("headers lock"),
            vec![
                None,
                Some("Bearer first".to_string()),
                Some("Bearer second".to_string())
            ]
        );
        shutdown.send(()).expect("stop fake WorkerTaskService");
        server
            .await
            .expect("join fake WorkerTaskService")
            .expect("serve task service");
    }
}