|
| 1 | +use http::{uri::Scheme, Request, Response, Uri}; |
| 2 | +use hyper::client::{connect::Connection, HttpConnector}; |
| 3 | +use hyper::Body; |
| 4 | +use std::fmt::Debug; |
| 5 | +use tokio::io::{AsyncRead, AsyncWrite}; |
| 6 | +use tower_service::Service; |
| 7 | + |
| 8 | +const USER_AGENT_HEADER: &str = "User-Agent"; |
| 9 | +const USER_AGENT: &str = concat!("aws-lambda-rust/", env!("CARGO_PKG_VERSION")); |
| 10 | + |
| 11 | +/// Error type that lambdas may result in |
| 12 | +pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>; |
| 13 | + |
| 14 | +#[derive(Debug)] |
| 15 | +pub struct Client<C = HttpConnector> { |
| 16 | + pub base: Uri, |
| 17 | + pub client: hyper::Client<C>, |
| 18 | +} |
| 19 | + |
| 20 | +impl Client { |
| 21 | + pub fn builder() -> ClientBuilder<HttpConnector> { |
| 22 | + ClientBuilder { |
| 23 | + connector: HttpConnector::new(), |
| 24 | + uri: None, |
| 25 | + } |
| 26 | + } |
| 27 | +} |
| 28 | + |
| 29 | +impl<C> Client<C> |
| 30 | +where |
| 31 | + C: hyper::client::connect::Connect + Sync + Send + Clone + 'static, |
| 32 | +{ |
| 33 | + pub async fn call(&self, req: Request<Body>) -> Result<Response<Body>, Error> { |
| 34 | + let req = self.set_origin(req)?; |
| 35 | + let response = self.client.request(req).await?; |
| 36 | + Ok(response) |
| 37 | + } |
| 38 | + |
| 39 | + pub fn with(base: Uri, connector: C) -> Self { |
| 40 | + let client = hyper::Client::builder().build(connector); |
| 41 | + Self { base, client } |
| 42 | + } |
| 43 | + |
| 44 | + fn set_origin<B>(&self, req: Request<B>) -> Result<Request<B>, Error> { |
| 45 | + let (mut parts, body) = req.into_parts(); |
| 46 | + let (scheme, authority) = { |
| 47 | + let scheme = self.base.scheme().unwrap_or(&Scheme::HTTP); |
| 48 | + let authority = self.base.authority().expect("Authority not found"); |
| 49 | + (scheme, authority) |
| 50 | + }; |
| 51 | + let path = parts.uri.path_and_query().expect("PathAndQuery not found"); |
| 52 | + |
| 53 | + let uri = Uri::builder() |
| 54 | + .scheme(scheme.clone()) |
| 55 | + .authority(authority.clone()) |
| 56 | + .path_and_query(path.clone()) |
| 57 | + .build(); |
| 58 | + |
| 59 | + match uri { |
| 60 | + Ok(u) => { |
| 61 | + parts.uri = u; |
| 62 | + Ok(Request::from_parts(parts, body)) |
| 63 | + } |
| 64 | + Err(e) => Err(Box::new(e)), |
| 65 | + } |
| 66 | + } |
| 67 | +} |
| 68 | + |
| 69 | +pub struct ClientBuilder<C: Service<http::Uri> = hyper::client::HttpConnector> { |
| 70 | + connector: C, |
| 71 | + uri: Option<http::Uri>, |
| 72 | +} |
| 73 | + |
| 74 | +impl<C> ClientBuilder<C> |
| 75 | +where |
| 76 | + C: Service<http::Uri> + Clone + Send + Sync + Unpin + 'static, |
| 77 | + <C as Service<http::Uri>>::Future: Unpin + Send, |
| 78 | + <C as Service<http::Uri>>::Error: Into<Box<dyn std::error::Error + Send + Sync>>, |
| 79 | + <C as Service<http::Uri>>::Response: AsyncRead + AsyncWrite + Connection + Unpin + Send + 'static, |
| 80 | +{ |
| 81 | + pub fn with_connector<C2>(self, connector: C2) -> ClientBuilder<C2> |
| 82 | + where |
| 83 | + C2: Service<http::Uri> + Clone + Send + Sync + Unpin + 'static, |
| 84 | + <C2 as Service<http::Uri>>::Future: Unpin + Send, |
| 85 | + <C2 as Service<http::Uri>>::Error: Into<Box<dyn std::error::Error + Send + Sync>>, |
| 86 | + <C2 as Service<http::Uri>>::Response: AsyncRead + AsyncWrite + Connection + Unpin + Send + 'static, |
| 87 | + { |
| 88 | + ClientBuilder { |
| 89 | + connector, |
| 90 | + uri: self.uri, |
| 91 | + } |
| 92 | + } |
| 93 | + |
| 94 | + pub fn with_endpoint(self, uri: http::Uri) -> Self { |
| 95 | + Self { uri: Some(uri), ..self } |
| 96 | + } |
| 97 | + |
| 98 | + pub fn build(self) -> Result<Client<C>, Error> { |
| 99 | + let uri = match self.uri { |
| 100 | + Some(uri) => uri, |
| 101 | + None => { |
| 102 | + let uri = std::env::var("AWS_LAMBDA_RUNTIME_API").expect("Missing AWS_LAMBDA_RUNTIME_API env var"); |
| 103 | + uri.clone().try_into().expect("Unable to convert to URL") |
| 104 | + } |
| 105 | + }; |
| 106 | + Ok(Client::with(uri, self.connector)) |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +pub fn build_request() -> http::request::Builder { |
| 111 | + http::Request::builder().header(USER_AGENT_HEADER, USER_AGENT) |
| 112 | +} |
0 commit comments