gcloud-sdk 0.33.3

//! Fixtures shared by the tests: scripted credentials, an RSA key, a local HTTP stub
//! server and local gRPC endpoints.

// Without a JWT crypto provider the tests that sign or verify are not built, and most
// fixtures go unused.
#![cfg_attr(
    not(any(feature = "jwt-aws-lc-rs", feature = "jwt-rust-crypto")),
    allow(dead_code)
)]

use std::convert::Infallible;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;

use futures::future::BoxFuture;
use futures::StreamExt;
use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag};
use google_cloud_auth::errors::CredentialsError;
use hyper::header::{HeaderMap, HeaderValue};
use hyper::http::{Extensions, Uri};
use hyper_util::rt::TokioIo;
#[cfg(feature = "id-token-verify")]
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tonic::server::NamedService;
use tonic::transport::{Channel, Endpoint};
use tonic::Status;
use tonic_prost::ProstCodec;

/// Credentials that serve scripted headers the way google-cloud-auth's do: under an entity
/// tag that changes when the headers do, answering `NotModified` to a caller that holds
/// the current tag. Clones share their script and counters.
#[derive(Debug, Clone)]
pub(crate) struct StubCredentials {
    current: Arc<Mutex<TaggedHeaders>>,
    served: Arc<AtomicUsize>,
    not_modified: Arc<AtomicUsize>,
}

#[derive(Debug, Clone)]
struct TaggedHeaders {
    entity_tag: EntityTag,
    headers: HeaderMap,
}

impl StubCredentials {
    pub(crate) fn new(headers: HeaderMap) -> Self {
        Self {
            current: Arc::new(Mutex::new(TaggedHeaders {
                entity_tag: EntityTag::new(),
                headers,
            })),
            served: Arc::new(AtomicUsize::new(0)),
            not_modified: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Credentials that serve `authorization: Bearer {token}`, not marked sensitive.
    pub(crate) fn bearer(token: &str) -> Self {
        Self::new(bearer_headers(token))
    }

    /// Serves `authorization: Bearer {token}` from now on, under a new entity tag.
    pub(crate) fn rotate_bearer(&self, token: &str) {
        *self.current.lock().unwrap() = TaggedHeaders {
            entity_tag: EntityTag::new(),
            headers: bearer_headers(token),
        };
    }

    /// How many times the headers were served in full.
    pub(crate) fn served(&self) -> usize {
        self.served.load(Ordering::SeqCst)
    }

    /// How many times a caller holding the current tag was answered `NotModified`.
    pub(crate) fn not_modified(&self) -> usize {
        self.not_modified.load(Ordering::SeqCst)
    }
}

impl CredentialsProvider for StubCredentials {
    async fn headers(
        &self,
        extensions: Extensions,
    ) -> Result<CacheableResource<HeaderMap>, CredentialsError> {
        let current = self.current.lock().unwrap().clone();
        if extensions.get::<EntityTag>() == Some(&current.entity_tag) {
            self.not_modified.fetch_add(1, Ordering::SeqCst);
            return Ok(CacheableResource::NotModified);
        }
        self.served.fetch_add(1, Ordering::SeqCst);
        Ok(CacheableResource::New {
            entity_tag: current.entity_tag,
            data: current.headers,
        })
    }

    async fn universe_domain(&self) -> Option<String> {
        None
    }
}

fn bearer_headers(token: &str) -> HeaderMap {
    let mut headers = HeaderMap::new();
    headers.insert(
        hyper::header::AUTHORIZATION,
        HeaderValue::from_str(&format!("Bearer {token}")).unwrap(),
    );
    headers
}

/// A 2048-bit RSA key generated for these tests; it signs nothing outside them.
pub(crate) const TEST_RSA_PRIVATE_KEY: &str = r"-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDfuSBdVqRqLKTk
uFNL6aieALaeza3kcVXFzIMZw+SIXWMW372Hes+ettM5501atS2ND8FXSmP3NR9/
8ac4kINZLlt94Y2l2uu0NCagKYWRxGd96oYpywYNLJQWvPxv9ETO1tVeDlT/EQ3f
sLlRmhSvZBjS3mWf2lX7AkvLz6zR/h4rvREwoNbz6q8Go9hdsxTNFsQcbIE6ZhI2
75hZxR0rKZA0n8UUA552HoAa3OJx03iGABK7h/k3hm7ZmnPaD558RVHUc7imJwsz
5uMpoIiTzBJlewvI68OkVdfuvM8UVeItVGj1t3cgg0ClAOfuVxn83c1iNTBMobQ3
IUc9u8mBAgMBAAECggEAAZB3kbeCoExuHbxNNs1sNKigHkWlZuDE/wRSUVqNjUeL
4xPO7TTWYU95dCDyKUV1i1Q2H6BhSQ/5x6j+qJZYGMZKdPugBC4e8kxgDcQkjzOe
nqKGbCHjibGLWopZQIYmegTGCqmSfhMWM/92GEQ5y00poEXWU5MRAVePhJ4P4QsA
vVKVB408BdYl7/ll+N9jNNICKdlX8ETrBp+TPTi2VmGuCPSRGH9EUehmzxuRwUoK
U4uY10En8pyEMBjiNGiOjyFn6NvV4KtsMyUIh2DrUBEcyAYSRt33Lzcfflf3dS4p
6IN6uNFzPhtFBY4kKO01w7ssabBHO3bWhZ61PMvjAQKBgQD2yIVDvFegCW/t9/4i
VBUJP08XH7fKU3L1Y6hZnJodzsYZl4wYDP8j5uqcbsc5aocv5bAWDQugxK1OrBp3
NgozQiOjkIsPjS1Od6GGMIyDXvprki9iWyYnkfi1aiKS8+ZI2rN+xUeyBkBHGzcz
Nxqju1E1pa6ziSfCjXQfGu40IQKBgQDoFCEezWHknZHss/V853qG5hDeFooYvfgd
gB+ErvMZkJzFcmxVJPtc/j/rNXCoIT/tzFlEsLjYgjskBlrOzuQOzlUNySYejN3g
uwX/3icMDrurf3DsXHVGf2hUrNjqECrV+gGRnfmVBvtBs5hswIbLe4sdnARd5w0D
+2PGOwjpYQKBgGQ5W6X+v7eHHaYPqW5Xp5Nx2rURdJr++RkfuCdsqkqgx2NtYMAD
xzrVdULC2rY+xVh2d+T8t1Q7jAb/bmAr2kim+8JZ2aAfPd84Rqkw3mAcGqfFXukb
C4vWhKNoz1HPLB86cttxU4TBdSlrrCdoobENShX3i9PuR++Dcz2Oul8BAoGAHYbC
F/slY0Kw2B6lMvj4W8VVjAvuEevJb2dnmyfvAeemKnC+W67S1Cf81d38sUdZrNV8
3gJl4hXflFvCneEwnrmdlJ1s3iIp8Hea8cy/xwbw1YbnRQsWWJvJGEzNZCoeaQ7f
uDkTEeTLfrZsxBlSjPzw3BmHbLMCsuj+7q+AGMECgYEAp95XH5iyxxsqDRbtsDD3
xcOguBDd9RXvGAI0v2ZdC1CF/MLsHRT8+COxaVb60NRRiyVTbyQ2BDfR5B6Tb8L1
qv0DqY+7mrGKzFdVr3Uf2lJJYFipabaL1GrD5FkZYPW21eNfBj2QRorgLf70n7L+
VPNvPfZ32mJr45YUokuFcGU=
-----END PRIVATE KEY-----";

/// A JWT signed with [`TEST_RSA_PRIVATE_KEY`] under key ID `kid`.
#[cfg(feature = "id-token-verify")]
pub(crate) fn signed_jwt(kid: &str, claims: &serde_json::Value) -> String {
    let header = Header {
        kid: Some(kid.to_string()),
        ..Header::new(Algorithm::RS256)
    };
    crate::jwt_crypto::ensure_provider();
    let key = EncodingKey::from_rsa_pem(TEST_RSA_PRIVATE_KEY.as_bytes()).unwrap();
    encode(&header, claims, &key).unwrap()
}

/// A JWK set that publishes the public half of [`TEST_RSA_PRIVATE_KEY`] under each of
/// `key_ids`.
#[cfg(feature = "id-token-verify")]
pub(crate) fn test_jwk_set(key_ids: &[&str]) -> jsonwebtoken::jwk::JwkSet {
    crate::jwt_crypto::ensure_provider();
    let encoding_key = EncodingKey::from_rsa_pem(TEST_RSA_PRIVATE_KEY.as_bytes()).unwrap();
    let keys = key_ids
        .iter()
        .map(|kid| {
            let mut jwk =
                jsonwebtoken::jwk::Jwk::from_encoding_key(&encoding_key, Algorithm::RS256).unwrap();
            jwk.common.key_id = Some(kid.to_string());
            jwk
        })
        .collect();
    jsonwebtoken::jwk::JwkSet { keys }
}

/// One response of a [`StubServer`].
pub(crate) struct StubResponse {
    status_line: &'static str,
    headers: Vec<(&'static str, String)>,
    body: String,
}

impl StubResponse {
    pub(crate) fn json(status_line: &'static str, body: impl Into<String>) -> Self {
        Self {
            status_line,
            headers: vec![("content-type", "application/json".to_string())],
            body: body.into(),
        }
    }

    pub(crate) fn text(status_line: &'static str, body: impl Into<String>) -> Self {
        Self {
            status_line,
            headers: vec![("content-type", "application/text".to_string())],
            body: body.into(),
        }
    }

    pub(crate) fn with_header(mut self, name: &'static str, value: impl Into<String>) -> Self {
        self.headers.push((name, value.into()));
        self
    }
}

/// A request a [`StubServer`] received.
#[derive(Debug, Clone)]
pub(crate) struct ReceivedRequest {
    /// The request line, such as `POST /token HTTP/1.1`.
    pub(crate) request_line: String,
    pub(crate) headers: Vec<(String, String)>,
}

impl ReceivedRequest {
    pub(crate) fn header(&self, name: &str) -> Option<&str> {
        self.headers
            .iter()
            .find(|(header, _)| header.eq_ignore_ascii_case(name))
            .map(|(_, value)| value.as_str())
    }
}

/// A local HTTP/1.1 server that answers one request per connection with the given
/// responses in order, records every request, and stops accepting after the last one.
pub(crate) struct StubServer {
    pub(crate) url: String,
    received: Arc<Mutex<Vec<ReceivedRequest>>>,
}

impl StubServer {
    pub(crate) async fn start(responses: Vec<StubResponse>) -> Self {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        let received = Arc::new(Mutex::new(Vec::new()));
        let recorder = Arc::clone(&received);
        tokio::spawn(async move {
            for response in responses {
                let (mut socket, _) = listener.accept().await.unwrap();
                let request = read_request(&mut socket).await;
                recorder.lock().unwrap().push(request);
                let mut head = format!(
                    "HTTP/1.1 {}\r\ncontent-length: {}\r\nconnection: close\r\n",
                    response.status_line,
                    response.body.len()
                );
                for (name, value) in &response.headers {
                    head.push_str(&format!("{name}: {value}\r\n"));
                }
                head.push_str("\r\n");
                socket.write_all(head.as_bytes()).await.unwrap();
                socket.write_all(response.body.as_bytes()).await.unwrap();
            }
        });
        Self { url, received }
    }

    pub(crate) fn received(&self) -> Vec<ReceivedRequest> {
        self.received.lock().unwrap().clone()
    }
}

async fn read_request(socket: &mut tokio::net::TcpStream) -> ReceivedRequest {
    let mut buf = Vec::new();
    let mut chunk = [0u8; 4096];
    let head_end = loop {
        let read = socket.read(&mut chunk).await.unwrap();
        assert!(read > 0, "connection closed before the request head ended");
        buf.extend_from_slice(&chunk[..read]);
        if let Some(pos) = buf.windows(4).position(|window| window == b"\r\n\r\n") {
            break pos + 4;
        }
    };
    let head = String::from_utf8(buf[..head_end].to_vec()).unwrap();
    let mut lines = head.split("\r\n").filter(|line| !line.is_empty());
    let request_line = lines.next().unwrap().to_string();
    let headers: Vec<(String, String)> = lines
        .map(|line| {
            let (name, value) = line.split_once(':').unwrap();
            (name.trim().to_string(), value.trim().to_string())
        })
        .collect();
    let content_length: usize = headers
        .iter()
        .find(|(name, _)| name.eq_ignore_ascii_case("content-length"))
        .map(|(_, value)| value.parse().unwrap())
        .unwrap_or(0);
    // The body is read so that the client sees its request through before the response.
    while buf.len() < head_end + content_length {
        let read = socket.read(&mut chunk).await.unwrap();
        assert!(read > 0, "connection closed before the request body ended");
        buf.extend_from_slice(&chunk[..read]);
    }
    ReceivedRequest {
        request_line,
        headers,
    }
}

/// A gRPC endpoint that accepts connections and never answers, counting the connections it
/// received: a channel's eager connect succeeds against it, and a request sent to it reaches
/// the middleware's token fetch before it stalls.
pub(crate) async fn silent_endpoint() -> (String, Arc<AtomicUsize>) {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let url = format!("http://{}", listener.local_addr().unwrap());
    let connections = Arc::new(AtomicUsize::new(0));
    let counter = connections.clone();
    tokio::spawn(async move {
        let mut held = Vec::new();
        while let Ok((socket, _)) = listener.accept().await {
            counter.fetch_add(1, Ordering::SeqCst);
            held.push(socket);
        }
    });
    (url, connections)
}

/// A channel to an in-process tonic server that answers every call with a stream of empty
/// messages, one after each of `gaps`, and then keeps the stream open without sending
/// another.
///
/// The two talk through an in-memory pipe rather than a socket: with the clock paused, a
/// runtime with nothing to run polls sockets without waiting and then advances the clock to
/// the next timer, so a message sent on a socket can arrive after a timer that is due later.
pub(crate) async fn stalling_stream_channel(gaps: Vec<Duration>) -> Channel {
    let (connections, mut accepted) = tokio::sync::mpsc::unbounded_channel();
    tokio::spawn(
        tonic::transport::Server::builder()
            .add_service(StallingStream { gaps })
            .serve_with_incoming(futures::stream::poll_fn(move |cx| accepted.poll_recv(cx))),
    );
    Endpoint::from_static("http://stalling-stream.test")
        .connect_with_connector(tower::service_fn(move |_: Uri| {
            let (client, server) = tokio::io::duplex(64 * 1024);
            connections.send(Ok::<_, std::io::Error>(server)).unwrap();
            std::future::ready(Ok::<_, std::io::Error>(TokioIo::new(client)))
        }))
        .await
        .unwrap()
}

#[derive(Clone)]
struct StallingStream {
    gaps: Vec<Duration>,
}

impl NamedService for StallingStream {
    const NAME: &'static str = "test.StallingStream";
}

impl tower::Service<hyper::Request<tonic::body::Body>> for StallingStream {
    type Response = hyper::Response<tonic::body::Body>;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Self::Response, Infallible>>;

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

    fn call(&mut self, request: hyper::Request<tonic::body::Body>) -> Self::Future {
        let gaps = self.gaps.clone();
        let messages = tower::service_fn(move |_: tonic::Request<()>| {
            let stream = futures::stream::iter(gaps.clone())
                .then(|gap| async move {
                    tokio::time::sleep(gap).await;
                    Ok::<(), Status>(())
                })
                .chain(futures::stream::pending());
            std::future::ready(Ok::<_, Status>(tonic::Response::new(stream)))
        });
        Box::pin(async move {
            Ok(tonic::server::Grpc::new(ProstCodec::<(), ()>::default())
                .server_streaming(messages, request)
                .await)
        })
    }
}