spiffe 0.19.0

use crate::constants::DEFAULT_SVID;
use crate::workload_api::pb::workload::{
    JwtBundlesRequest, JwtBundlesResponse, Jwtsvid, JwtsvidRequest, JwtsvidResponse,
    ValidateJwtsvidRequest, ValidateJwtsvidResponse,
};
use crate::{
    JwtBundle, JwtBundleSet, JwtSvid, SpiffeId, TrustDomain, WorkloadApiClient, WorkloadApiError,
};
use futures::{Stream, StreamExt as _};
use std::str::FromStr as _;
use std::sync::Arc;

impl WorkloadApiClient {
    /// Fetches the current set of JWT bundles from the SPIFFE Workload API.
    ///
    /// This method establishes a streaming gRPC request to the Workload API
    /// and returns the latest JWT bundle set received from the server.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the gRPC request fails, the stream
    /// terminates unexpectedly, or an invalid response is received.
    pub async fn fetch_jwt_bundles(&self) -> Result<JwtBundleSet, WorkloadApiError> {
        let request = JwtBundlesRequest::default();

        let mut client = self.client.clone();

        let grpc_stream_response: tonic::Response<tonic::Streaming<JwtBundlesResponse>> = client
            .fetch_jwt_bundles(request)
            .await
            .map_err(WorkloadApiError::from_status)?;

        let response = Self::first_message(grpc_stream_response.into_inner()).await?;
        Self::parse_jwt_bundle_set_from_grpc_response(response)
    }

    /// Fetches a `JwtSvid` for the given audience and optional SPIFFE ID.
    ///
    /// If `spiffe_id` is `None`, the Workload API returns the default identity.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the JWT-SVID request fails or the Workload API
    /// returns an invalid or empty response.
    pub async fn fetch_jwt_svid<I>(
        &self,
        audience: I,
        spiffe_id: Option<&SpiffeId>,
    ) -> Result<JwtSvid, WorkloadApiError>
    where
        I: IntoIterator,
        I::Item: AsRef<str>,
    {
        #[cfg(test)]
        if let Some(fetch) = &self.jwt_fetch_hook {
            let audience = audience
                .into_iter()
                .map(|a| a.as_ref().to_string())
                .collect();
            return fetch(audience, spiffe_id.cloned()).await;
        }

        let response = self.fetch_jwt(audience, spiffe_id).await?;
        let mut selected = None;
        for (index, entry) in response.svids.iter().enumerate() {
            let svid = Self::parse_jwt_svid_entry(entry)?;
            if index == DEFAULT_SVID {
                selected = Some(svid);
            }
        }
        selected.ok_or(WorkloadApiError::EmptyResponse)
    }

    /// Fetches all JWT-SVIDs for the given audience and optional SPIFFE ID.
    ///
    /// The Workload API can return more than one JWT-SVID. Each returned [`JwtSvid`] may include an
    /// optional **hint** (via [`JwtSvid::hint`]) that can be used to disambiguate which SVID to use.
    ///
    /// If `spiffe_id` is `None`, the Workload API returns JWT-SVIDs for the default identity.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the JWT-SVID request fails, the Workload API response is
    /// invalid or empty, or any returned token cannot be parsed.
    pub async fn fetch_all_jwt_svids<I>(
        &self,
        audience: I,
        spiffe_id: Option<&SpiffeId>,
    ) -> Result<Vec<JwtSvid>, WorkloadApiError>
    where
        I: IntoIterator,
        I::Item: AsRef<str>,
    {
        let response = self.fetch_jwt(audience, spiffe_id).await?;

        response
            .svids
            .iter()
            .map(Self::parse_jwt_svid_entry)
            .collect()
    }

    /// Fetches the JWT-SVID whose Workload API hint matches `hint`.
    ///
    /// Wrapper around [`WorkloadApiClient::fetch_all_jwt_svids`] that selects
    /// a single [`JwtSvid`] by its hint.
    ///
    /// The hint is **not** part of the JWT token; it is transport metadata provided by the SPIFFE
    /// Workload API to help identify a specific SVID when multiple are available.
    ///
    /// If `spiffe_id` is `None`, the Workload API returns JWT-SVIDs for the default identity.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the JWT-SVID request fails, the Workload API response is
    /// invalid, or no JWT-SVID with the requested hint is found.
    pub async fn fetch_jwt_svid_by_hint<I>(
        &self,
        audience: I,
        spiffe_id: Option<&SpiffeId>,
        hint: impl AsRef<str>,
    ) -> Result<JwtSvid, WorkloadApiError>
    where
        I: IntoIterator,
        I::Item: AsRef<str>,
    {
        let hint = hint.as_ref();
        let all = self.fetch_all_jwt_svids(audience, spiffe_id).await?;
        all.into_iter()
            .find(|s| s.hint() == Some(hint))
            .ok_or_else(|| WorkloadApiError::HintNotFound(hint.to_owned()))
    }

    /// Fetches a JWT-SVID token string for the given audience and optional SPIFFE ID.
    ///
    /// If `spiffe_id` is `None`, the Workload API returns the default identity.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the token request fails or the Workload API
    /// returns an invalid or empty response.
    #[cfg(feature = "jwt")]
    pub async fn fetch_jwt_token<I>(
        &self,
        audience: I,
        spiffe_id: Option<&SpiffeId>,
    ) -> Result<String, WorkloadApiError>
    where
        I: IntoIterator,
        I::Item: AsRef<str>,
    {
        let response = self.fetch_jwt(audience, spiffe_id).await?;
        response
            .svids
            .get(DEFAULT_SVID)
            .map(|r| r.svid.clone())
            .ok_or(WorkloadApiError::EmptyResponse)
    }

    /// Validates a JWT-SVID token for the given audience and returns the parsed [`JwtSvid`].
    ///
    /// Validation is performed by the SPIRE agent via the Workload API. After successful
    /// validation, the token is parsed locally for structured access. The use of
    /// `parse_insecure` is safe here because the security property comes from the agent's
    /// validation, not from local signature verification.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if validation fails, the successful response
    /// omits mandatory identity/claims or has malformed required claims, or the
    /// token cannot be parsed. Response structure checks do not perform local
    /// signature, audience-match, or temporal verification.
    pub async fn validate_jwt_token(
        &self,
        audience: &str,
        jwt_token: &str,
    ) -> Result<JwtSvid, WorkloadApiError> {
        // Validate via the SPIRE agent (security property comes from agent validation)
        let response = self.validate_jwt(audience, jwt_token).await?;
        Self::validate_jwt_response(&response)?;
        // Parse locally for structured access (safe because agent already validated)
        let jwt_svid = JwtSvid::parse_insecure(jwt_token)?;
        Ok(jwt_svid)
    }

    /// Streams JWT bundle set updates from the Workload API.
    ///
    /// The stream ends when the server closes the connection. This stream does not
    /// automatically reconnect; if you need resilience and automatic reconnection,
    /// use [`crate::X509Source`] for X.509 material or handle reconnection manually.
    ///
    /// # Errors
    ///
    /// Returns a [`WorkloadApiError`] if the Workload API stream cannot be
    /// established or the initial request fails.
    pub async fn stream_jwt_bundles(
        &self,
    ) -> Result<
        impl Stream<Item = Result<JwtBundleSet, WorkloadApiError>> + Send + 'static + use<>,
        WorkloadApiError,
    > {
        let request = JwtBundlesRequest::default();

        let mut client = self.client.clone();

        let response = client
            .fetch_jwt_bundles(request)
            .await
            .map_err(WorkloadApiError::from_status)?;
        let stream = response.into_inner().map(|message| {
            message
                .map_err(WorkloadApiError::from_status)
                .and_then(Self::parse_jwt_bundle_set_from_grpc_response)
        });
        Ok(Box::pin(stream))
    }
}

impl WorkloadApiClient {
    fn parse_jwt_svid_entry(entry: &Jwtsvid) -> Result<JwtSvid, WorkloadApiError> {
        let mut svid = JwtSvid::from_str(&entry.svid)?;
        if !entry.hint.is_empty() {
            svid = svid.with_hint(Arc::<str>::from(entry.hint.as_str()));
        }
        Ok(svid)
    }

    async fn fetch_jwt<I>(
        &self,
        audience: I,
        spiffe_id: Option<&SpiffeId>,
    ) -> Result<JwtsvidResponse, WorkloadApiError>
    where
        I: IntoIterator,
        I::Item: AsRef<str>,
    {
        let request = JwtsvidRequest {
            spiffe_id: spiffe_id.map(ToString::to_string).unwrap_or_default(),
            audience: audience
                .into_iter()
                .map(|a| a.as_ref().to_string())
                .collect(),
        };

        let mut client = self.client.clone();

        let response = client
            .fetch_jwtsvid(request)
            .await
            .map_err(WorkloadApiError::from_status)?
            .into_inner();
        if response.svids.is_empty() {
            return Err(WorkloadApiError::MissingRequiredField {
                field: "JWTSVIDResponse.svids",
            });
        }
        for svid in &response.svids {
            for (empty, field) in [
                (svid.spiffe_id.is_empty(), "JWTSVID.spiffe_id"),
                (svid.svid.is_empty(), "JWTSVID.svid"),
            ] {
                if empty {
                    return Err(WorkloadApiError::MissingRequiredField { field });
                }
            }
        }
        Ok(response)
    }

    async fn validate_jwt(
        &self,
        audience: &str,
        jwt_svid: &str,
    ) -> Result<ValidateJwtsvidResponse, WorkloadApiError> {
        let request = ValidateJwtsvidRequest {
            audience: audience.to_string(),
            svid: jwt_svid.into(),
        };
        let mut client = self.client.clone();
        Ok(client
            .validate_jwtsvid(request)
            .await
            .map_err(WorkloadApiError::from_status)?
            .into_inner())
    }

    fn parse_jwt_bundle_set_from_grpc_response(
        response: JwtBundlesResponse,
    ) -> Result<JwtBundleSet, WorkloadApiError> {
        if response.bundles.is_empty() {
            return Err(WorkloadApiError::MissingRequiredField {
                field: "JWTBundlesResponse.bundles",
            });
        }
        let mut bundle_set = JwtBundleSet::new();

        for (td, bundle_data) in response.bundles {
            let trust_domain = TrustDomain::try_from(td)?;
            let bundle = JwtBundle::from_jwt_authorities(trust_domain, &bundle_data)
                .map_err(WorkloadApiError::from)?;

            bundle_set.add_bundle(bundle);
        }

        Ok(bundle_set)
    }

    fn validate_jwt_response(response: &ValidateJwtsvidResponse) -> Result<(), WorkloadApiError> {
        use prost_types::value::Kind;

        if response.spiffe_id.is_empty() {
            return Err(WorkloadApiError::MissingRequiredField {
                field: "ValidateJWTSVIDResponse.spiffe_id",
            });
        }
        let claims = response
            .claims
            .as_ref()
            .ok_or(WorkloadApiError::MissingRequiredField {
                field: "ValidateJWTSVIDResponse.claims",
            })?;
        // The pinned Workload API contract requires the SPIFFE claims even though
        // the vendored protobuf comment calls this message field optional.
        for (name, field) in [
            ("sub", "ValidateJWTSVIDResponse.claims.sub"),
            ("aud", "ValidateJWTSVIDResponse.claims.aud"),
            ("exp", "ValidateJWTSVIDResponse.claims.exp"),
        ] {
            let kind = claims
                .fields
                .get(name)
                .and_then(|value| value.kind.as_ref())
                .ok_or(WorkloadApiError::MissingRequiredField { field })?;
            let valid = match (name, kind) {
                ("sub", Kind::StringValue(subject)) => SpiffeId::from_str(subject).is_ok(),
                ("aud", Kind::StringValue(_)) => true,
                ("aud", Kind::ListValue(audiences)) => {
                    !audiences.values.is_empty()
                        && audiences
                            .values
                            .iter()
                            .all(|value| matches!(value.kind, Some(Kind::StringValue(_))))
                }
                ("exp", Kind::NumberValue(expiration)) => expiration.is_finite(),
                _ => false,
            };
            if !valid {
                return Err(WorkloadApiError::InvalidResponseField { field });
            }
        }
        Ok(())
    }
}