use amberfork_record::{Body, CapturedRequest, CapturedResponse, retain_response_headers};
use std::collections::VecDeque;
use std::fmt;
use std::future::Future;
use std::sync::{Mutex, PoisonError};
pub trait Upstream: Send + Sync {
fn send(
&self,
request: &CapturedRequest,
) -> impl Future<Output = Result<CapturedResponse, UpstreamError>> + Send;
}
#[derive(Debug)]
#[non_exhaustive]
pub enum UpstreamError {
Exhausted,
Transport(reqwest::Error),
}
impl fmt::Display for UpstreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Exhausted => f.write_str("scripted upstream has no more responses"),
Self::Transport(source) => write!(f, "live relay to the provider failed: {source}"),
}
}
}
impl std::error::Error for UpstreamError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Exhausted => None,
Self::Transport(source) => Some(source),
}
}
}
#[derive(Debug, Clone)]
pub struct LiveUpstream {
client: reqwest::Client,
origin: String,
}
impl LiveUpstream {
#[must_use]
pub fn new(client: reqwest::Client, origin: impl Into<String>) -> Self {
Self {
client,
origin: origin.into().trim_end_matches('/').to_string(),
}
}
}
impl Upstream for LiveUpstream {
fn send(
&self,
request: &CapturedRequest,
) -> impl Future<Output = Result<CapturedResponse, UpstreamError>> + Send {
let method =
reqwest::Method::from_bytes(request.method.as_bytes()).unwrap_or(reqwest::Method::POST);
let mut outbound = self
.client
.request(method, format!("{}{}", self.origin, request.path));
for (name, value) in &request.headers {
if name.eq_ignore_ascii_case("host") {
continue;
}
outbound = outbound.header(name, value);
}
let body = match &request.body {
Body::Json(value) => serde_json::to_vec(value).unwrap_or_default(),
Body::Text(text) => text.clone().into_bytes(),
};
let outbound = outbound.body(body);
async move {
let response = outbound.send().await.map_err(UpstreamError::Transport)?;
let status = response.status().as_u16();
let headers = retain_response_headers(
response
.headers()
.iter()
.map(|(name, value)| (name.as_str(), value.as_bytes())),
);
let body = response.bytes().await.map_err(UpstreamError::Transport)?;
Ok(CapturedResponse {
status,
headers,
body: Body::from_bytes(&body),
})
}
}
}
#[derive(Debug)]
pub struct ScriptedUpstream {
responses: Mutex<VecDeque<CapturedResponse>>,
}
impl ScriptedUpstream {
#[must_use]
pub fn new(responses: impl IntoIterator<Item = CapturedResponse>) -> Self {
Self {
responses: Mutex::new(responses.into_iter().collect()),
}
}
}
impl Upstream for ScriptedUpstream {
fn send(
&self,
_request: &CapturedRequest,
) -> impl Future<Output = Result<CapturedResponse, UpstreamError>> + Send {
let next = self
.responses
.lock()
.unwrap_or_else(PoisonError::into_inner)
.pop_front()
.ok_or(UpstreamError::Exhausted);
async move { next }
}
}