kvarn_extensions/
reverse-proxy.rs

1use crate::connection::{Connection, EstablishedConnection};
2use kvarn::prelude::{internals::*, *};
3use std::net::{Ipv4Addr, SocketAddrV4};
4
5#[path = "url-rewrite.rs"]
6pub mod url_rewrite;
7
8pub use async_bits::CopyBuffer;
9#[macro_use]
10pub mod async_bits {
11    use kvarn::prelude::*;
12    macro_rules! ready {
13        ($poll: expr) => {
14            match $poll {
15                Poll::Ready(v) => v,
16                Poll::Pending => return Poll::Pending,
17            }
18        };
19    }
20
21    #[derive(Debug)]
22    pub struct CopyBuffer {
23        read_done: bool,
24        pos: usize,
25        cap: usize,
26        buf: Box<[u8]>,
27    }
28
29    impl CopyBuffer {
30        pub fn new() -> Self {
31            Self {
32                read_done: false,
33                pos: 0,
34                cap: 0,
35                buf: std::vec::from_elem(0, 2048).into_boxed_slice(),
36            }
37        }
38        pub fn with_capacity(initialized: usize) -> Self {
39            Self {
40                read_done: false,
41                pos: 0,
42                cap: 0,
43                buf: std::vec::from_elem(0, initialized).into_boxed_slice(),
44            }
45        }
46
47        /// Returns Ok(true) if it's done reading.
48        pub fn poll_copy<R, W>(
49            &mut self,
50            cx: &mut Context<'_>,
51            mut reader: Pin<&mut R>,
52            mut writer: Pin<&mut W>,
53        ) -> Poll<io::Result<bool>>
54        where
55            R: AsyncRead + ?Sized,
56            W: AsyncWrite + ?Sized,
57        {
58            loop {
59                // If our buffer is empty, then we need to read some data to
60                // continue.
61                if self.pos == self.cap && !self.read_done {
62                    let me = &mut *self;
63                    let mut buf = ReadBuf::new(&mut me.buf);
64                    ready!(reader.as_mut().poll_read(cx, &mut buf))?;
65                    let n = buf.filled().len();
66                    if n == 0 {
67                        self.read_done = true;
68                    } else {
69                        self.pos = 0;
70                        self.cap = n;
71                    }
72                }
73
74                // If our buffer has some data, let's write it out!
75                while self.pos < self.cap {
76                    let i = ready!(writer
77                        .as_mut()
78                        .poll_write(cx, &self.buf[self.pos..self.cap]))?;
79                    if i == 0 {
80                        return Poll::Ready(Err(io::Error::new(
81                            io::ErrorKind::WriteZero,
82                            "write zero byte into writer",
83                        )));
84                    } else {
85                        self.pos += i;
86                    }
87                    if self.pos >= self.cap {
88                        return Poll::Ready(Ok(false));
89                    }
90                }
91
92                // If we've written all the data and we've seen EOF, flush out the
93                // data and finish the transfer.
94                if self.pos == self.cap && self.read_done {
95                    ready!(writer.as_mut().poll_flush(cx))?;
96                    return Poll::Ready(Ok(true));
97                }
98            }
99        }
100    }
101    impl Default for CopyBuffer {
102        fn default() -> Self {
103            Self::new()
104        }
105    }
106}
107
108#[doc(hidden)]
109pub mod chain {
110    use super::*;
111
112    pub struct Chain<T, U> {
113        first: T,
114        second: U,
115        done_first: bool,
116    }
117
118    pub(super) fn chain<T, U>(first: T, second: U) -> Chain<T, U>
119    where
120        T: AsyncRead,
121        U: AsyncRead,
122    {
123        Chain {
124            first,
125            second,
126            done_first: false,
127        }
128    }
129
130    impl<T, U> fmt::Debug for Chain<T, U>
131    where
132        T: fmt::Debug,
133        U: fmt::Debug,
134    {
135        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
136            f.debug_struct("Chain")
137                .field("t", &self.first)
138                .field("u", &self.second)
139                .finish()
140        }
141    }
142
143    impl<T, U> AsyncRead for Chain<T, U>
144    where
145        T: AsyncRead + Unpin,
146        U: AsyncRead + Unpin,
147    {
148        fn poll_read(
149            mut self: Pin<&mut Self>,
150            cx: &mut Context<'_>,
151            buf: &mut ReadBuf<'_>,
152        ) -> Poll<io::Result<()>> {
153            let me = &mut *self;
154
155            if !me.done_first {
156                let rem = buf.remaining();
157                ready!(Pin::new(&mut me.first).poll_read(cx, buf))?;
158                if buf.remaining() == rem {
159                    me.done_first = true;
160                } else {
161                    return Poll::Ready(Ok(()));
162                }
163            }
164            Pin::new(&mut me.second).poll_read(cx, buf)
165        }
166    }
167}
168
169pub enum MaybeChunked<R1, R2> {
170    No(R1),
171    Yes(async_chunked_transfer::Decoder<R2>),
172}
173impl<R1: AsyncRead + Unpin, R2: AsyncRead + Unpin> AsyncRead for MaybeChunked<R1, R2> {
174    fn poll_read(
175        mut self: Pin<&mut Self>,
176        cx: &mut Context<'_>,
177        buf: &mut ReadBuf<'_>,
178    ) -> Poll<io::Result<()>> {
179        match &mut *self {
180            Self::No(reader) => Pin::new(reader).poll_read(cx, buf),
181            Self::Yes(reader) => Pin::new(reader).poll_read(cx, buf),
182        }
183    }
184}
185
186pub enum ConnectionResponse {
187    Whole(Response<Bytes>),
188    /// First part of body is in [`Response::body`], the remaining, `len - response.body().len()`
189    /// bytes, are readable from the [`EstablishedConnection`]. If the transfer-encoding is
190    /// chunked, the [`Response::body`] is also chunked so you can just continue relaying the data.
191    Partial {
192        len: Option<usize>,
193        response: Response<Bytes>,
194    },
195}
196impl EstablishedConnection {
197    /// If the body is `max_len` or greater, you need to continue streaming the body. See
198    /// [`ConnectionResponse::Partial`].
199    pub async fn request<'a, T: Debug>(
200        &'a mut self,
201        request: &Request<T>,
202        body: &[u8],
203        timeout: Duration,
204        max_len: usize,
205    ) -> Result<ConnectionResponse, GatewayError> {
206        let mut buffered = tokio::io::BufWriter::new(&mut *self);
207        debug!("Sending request");
208        write::request(request, body, &mut buffered).await?;
209
210        debug!("Sent reverse-proxy request. Reading response.");
211
212        let response = match tokio::time::timeout(
213            timeout,
214            kvarn::prelude::async_bits::read::response(&mut *self, 4 * 1024 * 1024, timeout),
215        )
216        .await
217        {
218            Ok(result) => match result {
219                Err(err) => return Err(err.into()),
220                Ok(response) => {
221                    let chunked =
222                        utils::header_eq(response.headers(), "transfer-encoding", "chunked");
223                    let len = if chunked {
224                        usize::MAX
225                    } else if body.is_empty() {
226                        utils::get_body_length_response(&response, Some(request.method()))
227                    } else {
228                        utils::get_body_length_response(&response, None)
229                    };
230
231                    if !chunked && len > max_len {
232                        return Ok(ConnectionResponse::Partial {
233                            len: Some(len),
234                            // it isn't chunked, so this is fine!
235                            response,
236                        });
237                    }
238
239                    let (mut head, body) = utils::split_response(response);
240
241                    let body = if len == 0 || len <= body.len() {
242                        body
243                    } else {
244                        let mut buffer = BytesMut::with_capacity(body.len() + 512);
245                        let original_body_len = body.len();
246
247                        let mut reader = if chunked {
248                            let reader = chain::chain(std::io::Cursor::new(body), self);
249                            let decoder = async_chunked_transfer::Decoder::new(reader);
250                            MaybeChunked::Yes(decoder)
251                        } else {
252                            buffer.extend(&body);
253                            MaybeChunked::No(self)
254                        };
255
256                        if let Ok(result) = tokio::time::timeout(
257                            timeout,
258                            read_to_end_or_max(&mut buffer, &mut reader, len.min(max_len)),
259                        )
260                        .await
261                        {
262                            result?;
263                            if buffer.len() >= max_len && chunked {
264                                use std::fmt::Write;
265
266                                // we're over max len and chunked, resort to write what we've read
267                                // thus far as a chunk.
268                                let buf_cap = buffer.len() + 4 /* \r\n ×2 */ + 10 /* length string */;
269                                let mut new_buf = BytesMut::with_capacity(buf_cap);
270                                write!(new_buf, "{}\r\n", buffer.len())
271                                    .expect("writing an integer to a BytesMut!!");
272                                new_buf.extend_from_slice(b"\r\n");
273                                drop(reader);
274                                return Ok(ConnectionResponse::Partial {
275                                    len: None,
276                                    response: head.map(|()| new_buf.freeze()),
277                                });
278                            }
279                        } else {
280                            warn!("Remote read timed out.");
281                            unsafe { buffer.set_len(if chunked { 0 } else { original_body_len }) };
282                        }
283
284                        if chunked {
285                            utils::remove_all_headers(head.headers_mut(), "transfer-encoding");
286                            debug!("Decoding chunked transfer-encoding.");
287                        }
288                        buffer.freeze()
289                    };
290
291                    head.map(|()| body)
292                }
293            },
294            Err(_) => return Err(GatewayError::Timeout),
295        };
296        Ok(ConnectionResponse::Whole(response))
297    }
298}
299
300#[derive(Debug)]
301pub enum GatewayError {
302    Io(io::Error),
303    Timeout,
304    Parse(parse::Error),
305}
306impl From<io::Error> for GatewayError {
307    fn from(err: io::Error) -> Self {
308        Self::Io(err)
309    }
310}
311impl From<parse::Error> for GatewayError {
312    fn from(err: parse::Error) -> Self {
313        Self::Parse(err)
314    }
315}
316
317#[derive(Debug)]
318pub enum OpenBackError {
319    Front(io::Error),
320    Back(io::Error),
321    Closed,
322}
323impl OpenBackError {
324    pub fn get_io(&self) -> Option<&io::Error> {
325        match self {
326            Self::Front(e) | Self::Back(e) => Some(e),
327            Self::Closed => None,
328        }
329    }
330    pub fn get_io_kind(&self) -> io::ErrorKind {
331        match self {
332            Self::Front(e) | Self::Back(e) => e.kind(),
333            Self::Closed => io::ErrorKind::BrokenPipe,
334        }
335    }
336}
337pub struct ByteProxy<'a, B: AsyncRead + AsyncWrite + Unpin> {
338    front: &'a mut ResponseBodyPipe,
339    back: &'a mut B,
340    front_buf: Vec<u8>,
341    back_buf: Vec<u8>,
342}
343impl<'a, B: AsyncRead + AsyncWrite + Unpin> ByteProxy<'a, B> {
344    pub fn new(front: &'a mut ResponseBodyPipe, back: &'a mut B) -> Self {
345        Self {
346            front,
347            back,
348            front_buf: Vec::with_capacity(16 * 1024),
349            back_buf: Vec::with_capacity(16 * 1024),
350        }
351    }
352    pub async fn channel(&mut self) -> Result<(), OpenBackError> {
353        let mut front_done = false;
354        let mut back_done = false;
355        loop {
356            match (front_done, back_done) {
357                // if any one is done, close the other!
358                (true, false) => {
359                    self.back.shutdown().await.map_err(OpenBackError::Back)?;
360                    break;
361                }
362                (false, true) => {
363                    if let ResponseBodyPipe::Http1(h1) = self.front {
364                        h1.lock()
365                            .await
366                            .shutdown()
367                            .await
368                            .map_err(OpenBackError::Front)?;
369                        break;
370                    } else {
371                        front_done = true;
372                    }
373                }
374                (false, false) => {
375                    let ResponseBodyPipe::Http1(h1) = self.front else {
376                        // won't ever go into this branch again
377                        front_done = true;
378                        continue;
379                    };
380                    let front_read = async {
381                        unsafe { self.front_buf.set_len(self.front_buf.capacity()) };
382                        let read = h1
383                            .lock()
384                            .await
385                            .read(&mut self.front_buf)
386                            .await
387                            .map_err(OpenBackError::Front)?;
388                        if read == 0 {
389                            front_done = true;
390                        }
391                        unsafe { self.front_buf.set_len(read) };
392                        Ok::<(), OpenBackError>(())
393                    };
394
395                    let back_read = async {
396                        unsafe { self.back_buf.set_len(self.back_buf.capacity()) };
397                        let read = self
398                            .back
399                            .read(&mut self.back_buf)
400                            .await
401                            .map_err(OpenBackError::Back)?;
402                        if read == 0 {
403                            back_done = true;
404                        }
405                        unsafe { self.back_buf.set_len(read) };
406                        Ok::<(), OpenBackError>(())
407                    };
408                    tokio::select! {
409                        r = front_read => {
410                            r?;
411                            self.back
412                                .write_all(&self.front_buf)
413                                .await
414                                .map_err(OpenBackError::Back)?;
415                            self.back
416                                .flush()
417                                .await
418                                .map_err(OpenBackError::Back)?;
419                        }
420                        r = back_read => {
421                            r?;
422                            self.front
423                                .send(Bytes::copy_from_slice(&self.back_buf))
424                                .await
425                                .map_err(io::Error::from)
426                                .map_err(OpenBackError::Front)?;
427                            self.front
428                                .flush()
429                                .await
430                                .map_err(io::Error::from)
431                                .map_err(OpenBackError::Front)?;
432                        }
433                    }
434                }
435                (true, true) => {
436                    break;
437                }
438            }
439        }
440        Ok(())
441    }
442}
443
444pub type ModifyRequestFn = Arc<dyn Fn(&mut Request<()>, &mut Bytes, SocketAddr) + Send + Sync>;
445pub type GetConnectionFn = Arc<dyn (Fn(&FatRequest, &Bytes) -> Option<Connection>) + Send + Sync>;
446
447/// Creates a new [`GetConnectionFn`] which always returns `kind`
448pub fn static_connection(kind: Connection) -> GetConnectionFn {
449    Arc::new(move |_, _| Some(kind.clone()))
450}
451
452#[must_use = "mount the reverse proxy manager"]
453pub struct Manager {
454    when: extensions::If,
455    connection: GetConnectionFn,
456    modify: Vec<ModifyRequestFn>,
457    timeout: Duration,
458    rewrite_url: bool,
459    priority: i32,
460}
461impl Manager {
462    /// Consider using [`static_connection`] if your connection type is not dependent of the request.
463    pub fn new(when: extensions::If, connection: GetConnectionFn, timeout: Duration) -> Self {
464        Self {
465            when,
466            connection,
467            modify: vec![],
468            timeout,
469            rewrite_url: true,
470            priority: -128,
471        }
472    }
473    /// Disables the built-in feature of rewriting the relative URLs so they point to the forwarded
474    /// site.
475    ///
476    /// **NOTE** that rewrite doesn't work when the response body is streamed.
477    pub fn disable_url_rewrite(mut self) -> Self {
478        self.rewrite_url = false;
479        self
480    }
481    /// Set the priority of the extension. The default is `-128`.
482    pub fn with_priority(mut self, priority: i32) -> Self {
483        self.priority = priority;
484        self
485    }
486    /// Add a function to run before the request is sent.
487    /// These are ran in the order they are added in.
488    pub fn add_modify_fn(mut self, modify: ModifyRequestFn) -> Self {
489        self.modify.push(modify);
490        self
491    }
492    /// [Add a modify fn](Self::add_modify_fn) which adds the IP of the request as the header
493    /// `x-real-ip`.
494    pub fn with_x_real_ip(self) -> Self {
495        self.add_modify_fn(Arc::new(|req, _, addr| {
496            req.headers_mut().insert(
497                "x-real-ip",
498                HeaderValue::try_from(addr.ip().to_string()).unwrap(),
499            );
500        }))
501    }
502    /// Consider using [`static_connection`] if your connection type is not dependent of the request.
503    pub fn base(base_path: &str, connection: GetConnectionFn, timeout: Duration) -> Self {
504        assert_eq!(base_path.chars().next(), Some('/'));
505        let path = if base_path.ends_with('/') {
506            base_path.to_owned()
507        } else {
508            let mut s = String::with_capacity(base_path.len() + 1);
509            s.push_str(base_path);
510            s.push('/');
511            s
512        };
513        let path = Arc::new(path);
514
515        let when_path = Arc::clone(&path);
516        let when = Box::new(move |request: &FatRequest, _host: &Host| {
517            request.uri().path().starts_with(when_path.as_str())
518        });
519
520        let modify: ModifyRequestFn = Arc::new(move |request, _, _| {
521            let path = Arc::clone(&path);
522
523            let mut parts = request.uri().clone().into_parts();
524
525            if let Some(path_and_query) = parts.path_and_query.as_ref() {
526                if let Some(s) = path_and_query.as_str().get(path.as_str().len() - 1..) {
527                    // We know this is a good path and query; we've just removed the first x bytes.
528                    // The -1 will always be on a char boundary; the last character is always '/'
529                    let short =
530                        uri::PathAndQuery::from_maybe_shared(Bytes::copy_from_slice(s.as_bytes()))
531                            .unwrap();
532                    parts.path_and_query = Some(short);
533                    parts.scheme = Some(uri::Scheme::HTTP);
534                    // For unwrap, see ↑
535                    let uri = Uri::from_parts(parts).unwrap();
536                    *request.uri_mut() = uri;
537                } else {
538                    error!("We didn't get the expected path string from Kvarn. We asked for one which started with `base_path`");
539                }
540            }
541        });
542
543        Self::new(when, connection, timeout).add_modify_fn(modify)
544    }
545    /// Attach this reverse proxy to `extensions`.
546    ///
547    /// !!Please!! use a [`crate::force_cache`] extension on the paths where the reverse proxy
548    /// acts, with [`comprash::ClientCachePreference::Ignore`].
549    /// ALSO, disable server cache!
550    pub fn mount(self, extensions: &mut Extensions) {
551        let connection = self.connection;
552        let modify = self.modify;
553
554        macro_rules! return_status {
555            ($result:expr, $status:expr, $host:expr) => {
556                match $result {
557                    Some(v) => v,
558                    None => {
559                        return default_error_response($status, $host, None).await;
560                    }
561                }
562            };
563        }
564
565        let timeout = self.timeout;
566        let rewrite_url = self.rewrite_url;
567
568        extensions.add_prepare_fn(
569            self.when,
570            prepare!(
571                req,
572                host,
573                _path,
574                addr,
575                move |connection: GetConnectionFn,
576                      modify: Vec<ModifyRequestFn>,
577                      timeout: Duration,
578                      rewrite_url: bool| {
579                    let mut empty_req = utils::empty_clone_request(req);
580                    let mut bytes = return_status!(
581                        req.body_mut().read_to_bytes(1024 * 1024 * 16).await.ok(),
582                        StatusCode::BAD_GATEWAY,
583                        host
584                    );
585
586                    let connection =
587                        return_status!(connection(req, &bytes), StatusCode::BAD_REQUEST, host);
588                    let mut connection = return_status!(
589                        connection.establish().await.ok(),
590                        StatusCode::GATEWAY_TIMEOUT,
591                        host
592                    );
593
594                    empty_req
595                        .headers_mut()
596                        .insert("accept-encoding", HeaderValue::from_static("identity"));
597
598                    if utils::header_eq(empty_req.headers(), "connection", "keep-alive") {
599                        empty_req
600                            .headers_mut()
601                            .insert("connection", HeaderValue::from_static("close"));
602                    }
603
604                    *empty_req.version_mut() = Version::HTTP_11;
605
606                    if let Ok(value) = host.name.parse() {
607                        empty_req.headers_mut().insert("host", value);
608                    }
609
610                    let wait = matches!(empty_req.method(), &Method::CONNECT)
611                        || empty_req.headers().get("upgrade")
612                            == Some(&HeaderValue::from_static("websocket"));
613
614                    let path = empty_req.uri().path().to_owned();
615
616                    for modify in modify {
617                        modify(&mut empty_req, &mut bytes, addr);
618                    }
619
620                    // limit cached responses to 10 MB, otherwise stream body
621                    let max_len = 10 * 1024 * 1024;
622                    let result = connection
623                        .request(&empty_req, &bytes, *timeout, max_len)
624                        .await;
625                    let mut response = match result {
626                        Ok(response) => {
627                            let (mut response, others) = match response {
628                                ConnectionResponse::Whole(r) => (r, None),
629                                ConnectionResponse::Partial { len, response } => {
630                                    (response, Some(len))
631                                }
632                            };
633
634                            // if we plan to stream the body, don't rewrite parts of it.
635                            if *rewrite_url && others.is_none() {
636                                let content_type = response
637                                    .headers()
638                                    .get("content-type")
639                                    .and_then(|ct| ct.to_str().ok())
640                                    .and_then(|ct| ct.parse::<Mime>().ok());
641                                if let Some(
642                                    (mime::TEXT, mime::HTML | mime::CSS)
643                                    | (mime::APPLICATION, mime::JAVASCRIPT),
644                                ) = content_type.as_ref().map(|ct| (ct.type_(), ct.subtype()))
645                                {
646                                    if let Some(prefix) = path.strip_suffix(empty_req.uri().path())
647                                    {
648                                        // Since we strip `.path` (which starts with `/`, Kvarn denies requests with more than one `/`),
649                                        // prefix is guaranteed not to end with `/`.
650                                        response = response.map(|body| {
651                                            url_rewrite::absolute(&body, prefix).freeze()
652                                        });
653                                    }
654                                }
655                            }
656                            let headers = response.headers_mut();
657                            if let None | Some(None) = &others {
658                                utils::remove_all_headers(headers, "content-length");
659                            }
660                            utils::remove_all_headers(headers, "keep-alive");
661                            if headers
662                                .get("connection")
663                                .and_then(|h| h.to_str().ok())
664                                .map(|h| !h.trim().eq_ignore_ascii_case("upgrade"))
665                                .unwrap_or(true)
666                            {
667                                utils::remove_all_headers(headers, "connection");
668                            }
669
670                            if let Some(len) = others {
671                                if let Some(len) = len {
672                                    response.headers_mut().insert(
673                                        "content-length",
674                                        HeaderValue::from_bytes(len.to_string().as_bytes())
675                                            .expect("integer isn't HeaderValue?"),
676                                    );
677                                }
678
679                                if wait {
680                                    error!(
681                                        "Waiting for websocket but also \
682                                        streaming body of length > {max_len}"
683                                    );
684                                }
685                                // trust me bro, it's stringently read
686                                #[allow(clippy::uninit_vec)]
687                                return FatResponse::no_cache(response).with_future(
688                                    response_pipe_fut!(
689                                        response_pipe,
690                                        _,
691                                        move |connection: EstablishedConnection| {
692                                            let mut buf = Vec::with_capacity(1024 * 64);
693                                            unsafe { buf.set_len(buf.capacity()) };
694                                            let mut i = 0u32;
695                                            loop {
696                                                // add 1 at the top to skip waiting for connection on first iter
697                                                i = i.wrapping_add(1);
698                                                let r = connection.read(&mut buf).await;
699                                                // okey this is crazy deep nesting i'm sorry
700                                                match r {
701                                                    Ok(read) => {
702                                                        if read == 0 {
703                                                            break;
704                                                        }
705                                                        // one chunk is max 64kB (see buffer above)
706                                                        // we want to check connection status every, say, 10MB, to not
707                                                        // exhaust resources.
708                                                        // 10MB/64kB = 160
709                                                        let data =
710                                                            Bytes::copy_from_slice(&buf[..read]);
711                                                        let r = if i % 160 == 0 {
712                                                            // to not just spew data in HTTP/2,
713                                                            // growing memory size to infinity!!!
714                                                            response_pipe
715                                                                .send_with_wait(
716                                                                    data,
717                                                                    10 * 1024 * 1024,
718                                                                )
719                                                                .await
720                                                        } else {
721                                                            response_pipe.send(data).await
722                                                        };
723                                                        match r {
724                                                            Ok(()) => {}
725                                                            Err(_) => {
726                                                                break;
727                                                            }
728                                                        }
729                                                    }
730                                                    Err(err) => {
731                                                        warn!(
732                                                            "Failed to stream body \
733                                                            from reverse connection: {err}"
734                                                        );
735                                                        break;
736                                                    }
737                                                }
738                                            }
739                                            let _ = response_pipe.close().await;
740                                        }
741                                    ),
742                                );
743                            } else {
744                                FatResponse::cache(response)
745                            }
746                        }
747                        Err(err) => {
748                            warn!("Got error {:?}", err);
749                            default_error_response(
750                                match err {
751                                    GatewayError::Io(_) | GatewayError::Parse(_) => {
752                                        StatusCode::BAD_GATEWAY
753                                    }
754                                    GatewayError::Timeout => StatusCode::GATEWAY_TIMEOUT,
755                                },
756                                host,
757                                None,
758                            )
759                            .await
760                        }
761                    };
762
763                    if wait {
764                        debug!("Keeping the pipe open!");
765                        let future = response_pipe_fut!(
766                            response_pipe,
767                            _,
768                            move |connection: EstablishedConnection| {
769                                let udp_connection =
770                                    matches!(connection, EstablishedConnection::Udp(_));
771
772                                let mut open_back = ByteProxy::new(response_pipe, connection);
773                                debug!("Created open back!");
774
775                                // Add 90 second timeout to UDP connections.
776                                let timeout_result = if udp_connection {
777                                    tokio::time::timeout(
778                                        Duration::from_secs(90),
779                                        open_back.channel(),
780                                    )
781                                    .await
782                                } else {
783                                    Ok(open_back.channel().await)
784                                };
785
786                                if let Ok(r) = timeout_result {
787                                    debug!("Open back responded! {:?}", r);
788                                    if let Err(err) = r {
789                                        if !matches!(
790                                            err.get_io_kind(),
791                                            io::ErrorKind::ConnectionAborted
792                                                | io::ErrorKind::ConnectionReset
793                                                | io::ErrorKind::BrokenPipe
794                                        ) {
795                                            warn!("Reverse proxy io error: {:?}", err);
796                                        }
797                                    }
798                                }
799                            }
800                        );
801
802                        response = response
803                            .with_future(future)
804                            .with_compress(comprash::CompressPreference::None);
805                    } else {
806                        drop(connection.shutdown().await);
807                        drop(connection);
808                    }
809
810                    response
811                }
812            ),
813            extensions::Id::new(self.priority, "Reverse proxy").no_override(),
814        );
815    }
816}
817
818pub fn localhost(port: u16) -> SocketAddr {
819    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port))
820}
kvarn_extensions/reverse-proxy.rs

kvarn_extensions/
reverse-proxy.rs
