1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267
#![forbid(unsafe_code)]
#![deny(
clippy::dbg_macro,
missing_copy_implementations,
rustdoc::missing_crate_level_docs,
missing_debug_implementations,
missing_docs,
nonstandard_style,
unused_qualifications
)]
/*!
http reverse proxy trillium handler
*/
use full_duplex_async_copy::full_duplex_copy;
use size::{Base, Size};
use std::convert::TryInto;
use trillium::{
async_trait, conn_try, Conn, Handler, KnownHeaderName,
Status::{NotFound, SwitchingProtocols},
};
use trillium_http::{transport::BoxedTransport, Upgrade};
use url::Url;
pub use trillium_client::{Client, Connector};
/**
the proxy handler
*/
#[derive(Debug)]
pub struct Proxy<C: Connector> {
target: Url,
client: Client<C>,
pass_through_not_found: bool,
halt: bool,
}
impl<C: Connector> Proxy<C> {
/**
construct a new proxy handler that sends all requests to the url
provided. if the url contains a path, the inbound request path
will be joined onto the end.
```
use trillium_smol::{TcpConnector, ClientConfig};
use trillium_proxy::Proxy;
let proxy = Proxy::<TcpConnector>::new("http://docs.trillium.rs/trillium_proxy");
```
*/
pub fn new(target: impl TryInto<Url>) -> Self {
let url = match target.try_into() {
Ok(url) => url,
Err(_) => panic!("could not convert proxy target into a url"),
};
assert!(!url.cannot_be_a_base(), "{} cannot be a base", url);
Self {
target: url,
client: Client::new().with_default_pool(),
pass_through_not_found: true,
halt: true,
}
}
/**
chainable constructor to specify the client Connector
configuration
```
use trillium_smol::{TcpConnector, ClientConfig};
use trillium_proxy::Proxy;
let proxy = Proxy::<TcpConnector>::new("http://trillium.rs")
.with_config(ClientConfig { //<-
nodelay: Some(true),
..Default::default()
});
```
*/
pub fn with_config(mut self, config: C::Config) -> Self {
self.client = self.client.with_config(config);
self
}
/**
chainable constructor to specfiy a [`Client`] to use. This is
useful if the application already is using trillium_client for
other requests, as it will reuse the same connection pool and
connector config.
note that this clears out any changes made with
[`Proxy::with_config`]. configure the client directly if you are
providing one
```
use trillium_smol::{TcpConnector, ClientConfig};
use trillium_proxy::{Proxy, Client};
let client = Client::new().with_default_pool();
let proxy = Proxy::<TcpConnector>::new("http://trillium.rs")
.with_client(client); //<-
```
```
// sharing a client with other trillium handlers
# use trillium_smol::{TcpConnector, ClientConfig};
# use trillium_proxy::{Proxy, Client};
use trillium::State;
let client = Client::new().with_default_pool();
let handler = (
State::new(client.clone()),
Proxy::<TcpConnector>::new("http://trillium.rs").with_client(client)
);
```
*/
pub fn with_client(mut self, client: Client<C>) -> Self {
self.client = client;
self
}
/**
chainable constructor to set the 404 Not Found handling
behavior. By default, this proxy will pass through the trillium
Conn unmodified if the proxy response is a 404 not found, allowing
it to be chained in a tuple handler. To modify this behavior, call
proxy_not_found, and the full 404 response will be forwarded. The
Conn will be halted unless [`Proxy::without_halting`] was
configured
```
# use trillium_smol::TcpConnector;
# use trillium_proxy::Proxy;
let proxy = Proxy::<TcpConnector>::new("http://trillium.rs")
.proxy_not_found();
```
*/
pub fn proxy_not_found(mut self) -> Self {
self.pass_through_not_found = false;
self
}
/**
The default behavior for this handler is to halt the conn on any
response other than a 404. If [`Proxy::proxy_not_found`] has been
configured, the default behavior for all response statuses is to
halt the trillium conn. To change this behavior, call
without_halting when constructing the proxy, and it will not halt
the conn. This is useful when passing the proxy reply through
[`trillium_html_rewriter`](https://docs.trillium.rs/trillium_html_rewriter).
```
# use trillium_smol::TcpConnector;
# use trillium_proxy::Proxy;
let proxy = Proxy::<TcpConnector>::new("http://trillium.rs")
.without_halting();
```
*/
pub fn without_halting(mut self) -> Self {
self.halt = false;
self
}
}
struct UpstreamUpgrade<T>(Upgrade<T>);
#[async_trait]
impl<C: Connector> Handler for Proxy<C> {
async fn run(&self, mut conn: Conn) -> Conn {
let mut request_url = conn_try!(self.target.clone().join(conn.path()), conn);
let querystring = conn.querystring();
if !querystring.is_empty() {
request_url.set_query(Some(conn.querystring()));
}
log::debug!("proxying to {}", request_url);
let mut client_conn = self.client.build_conn(conn.method(), request_url);
client_conn.request_headers().extend(
conn.headers()
.clone()
.into_iter()
.filter(|(name, _)| name != KnownHeaderName::Host),
);
// need a better solution for streaming request bodies through
// the proxy, but http-types::Body needs to be 'static. Fixing
// this probably will entail moving away from http-types::Body
// for outbound bodies.
//
// this is very inefficient and possibly unscalable in
// situations where request bodies are large. there is no
// reason that we couldn't have another lifetime on client
// conn here, though
if let Ok(client_body_content) = conn.request_body().await.read_bytes().await {
client_conn.set_request_body(client_body_content);
}
client_conn
.request_headers()
.insert(KnownHeaderName::Connection, "keep-alive");
trillium::conn_try!(client_conn.send().await, conn);
let conn = match client_conn.status() {
Some(SwitchingProtocols) => {
conn.headers_mut()
.extend(std::mem::take(client_conn.response_headers_mut()).into_iter());
conn.with_state(UpstreamUpgrade(client_conn.into()))
.with_status(SwitchingProtocols)
}
Some(NotFound) if self.pass_through_not_found => {
client_conn.recycle().await;
return conn;
}
Some(status) => {
conn.headers_mut()
.extend(std::mem::take(client_conn.response_headers_mut()).into_iter());
conn.with_body(client_conn).with_status(status)
}
_ => unreachable!(),
};
if self.halt {
conn.halt()
} else {
conn
}
}
fn has_upgrade(&self, upgrade: &Upgrade<BoxedTransport>) -> bool {
upgrade
.state
.get::<UpstreamUpgrade<C::Transport>>()
.is_some()
}
async fn upgrade(&self, mut upgrade: trillium::Upgrade) {
let upstream = upgrade
.state
.take::<UpstreamUpgrade<C::Transport>>()
.unwrap()
.0;
let downstream = upgrade;
match full_duplex_copy(upstream, downstream).await {
Err(e) => log::error!("{}:{} {:?}", file!(), line!(), e),
Ok((up, down)) => log::debug!("wrote {} up and {} down", bytes(up), bytes(down)),
}
}
}
fn bytes(bytes: u64) -> String {
Size::from_bytes(bytes)
.format()
.with_base(Base::Base10)
.to_string()
}