Feature suggestion for trillium: Persistent state for a persistent HTTP connection

[joshtriplett]

I have an internal API that uses the peer IP of an HTTP connection to look up information about the system making the request. I'd like to look up that information once per persistent connection, rather than once per request. Is there any convenient way I could cache that information and have it persist through any subsequent requests made on the same Transport, and then get dropped the when the Transport disconnects? (I can't just cache this in a table by peer IP, because peer IPs can be reused if there isn't an active connection from them.)

[joshtriplett]

I realize this is a fairly unusual use case, and I don't mind having to jump through some hoops to make it work. I'd just love somewhere I can hang some state off a Transport and get to it from any Conn based on that Transport.

[jbr]

I'd like to frame this sort of feature "what's an escape hatch for advanced users to implement nonstandard behavior without imposing a performance, security, or complexity cost on other users?"

It's really important to me that trillium be able to support idiosyncratic/advanced usage, but the less common the need, the more likely it is to require "ejecting" from an existing crate by forking it (building it into the app if it's only used in one place) to add behavior. This applies to my own usage of trillium as well, for what it's worth. I try really hard to ask "is this something most app authors are going to want, or is this something I should implement in my app."

In this case, I think the right escape hatch is to implement your own runtime adapter, either wrapping another runtime adapter or just starting from one and modifying it directly. You'd then be able to add additional state to the Transport. The reason the runtime adapters are built in terms of trillium-server-common is to make that sort of ejection relatively painless; a lot of the shared logic implementation is in ConfigExt and in default implementations on Server.

I think the remaining change needed to trillium would be a Conn::downcast_transport_ref<T: Transport>(&self) -> Option<&T> and Conn::downcast_transport_mut<T: Transport>(&mut self) -> Option<&mut T>.

My concern here is that for some (many) applications, subsequent Conns on the transport may be from different http clients on the other side of a reverse proxy or load balancer, and it's an unnecessary security hazard to introduce an api that is primarily a footgun for that common use case. I don't want someone to be like "ooh I can use this transport_state feature to cache stuff" and then when that turns into a vulnerability ask why trillium let them do that.

I've thought about adding logging hooks of some sort for transport events to be able to track keepalive usage, but am apprehensive of anything that breaks the illusion/abstraction of http events being fully independent from each other.

I'll post a gist with an example app with an "ejected" smol runtime adapter that has a StateSet on Transport in order to confirm that this would work. You'd probably want to store the types directly instead of in a StateSet if you know what they are at compile time.

[joshtriplett]

Having to have my own Transport and then downcast to that Transport sounds completely reasonable to me as the right level of ergonomics for this unusual application. And I agree with you that having this too easy to access would make people likely to share state when it shouldn't be. Having the two downcast_transport functions would make this relatively doable.

That said, based on your description, I don't think I'd need to completely replace the runtime adapter in order to do that. There's already a built-in mechanism for wrapping a Transport: the Acceptor. I think I can use an existing runtime, and add an acceptor that wraps the Transport in my own wrapper that passes through all the Transport methods and adds methods to get at the state.

Given that, don't worry about demoing a forked runtime adapter. I'd really appreciate the downcast_transport_ref and downcast_transport_mut functions, though.

[joshtriplett]

#432

[joshtriplett]

Using the implementation in #432, here's an example of attaching state to a custom transport:

`main.rs`

use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

use trillium::Conn;
use trillium_http::transport::BoxedTransport;
use trillium_router::{router, RouterConnExt};
use trillium_server_common::{Acceptor, Transport};

#[derive(Debug)]
struct TransportWrapper(BoxedTransport, Option<String>);

impl trillium_server_common::AsyncRead for TransportWrapper {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_read(cx, buf)
    }

    fn poll_read_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &mut [io::IoSliceMut<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_read_vectored(cx, bufs)
    }
}

impl trillium_server_common::AsyncWrite for TransportWrapper {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_flush(cx)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_close(cx)
    }

    fn poll_write_vectored(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        bufs: &[io::IoSlice<'_>],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_write_vectored(cx, bufs)
    }
}

impl Transport for TransportWrapper {
    fn set_ip_ttl(&mut self, ttl: u32) -> std::io::Result<()> {
        self.0.set_ip_ttl(ttl)
    }

    fn set_linger(&mut self, linger: Option<std::time::Duration>) -> std::io::Result<()> {
        self.0.set_linger(linger)
    }

    fn set_nodelay(&mut self, nodelay: bool) -> std::io::Result<()> {
        self.0.set_nodelay(nodelay)
    }

    fn peer_addr(&self) -> std::io::Result<Option<std::net::SocketAddr>> {
        self.0.peer_addr()
    }
}

impl TransportWrapper {
    fn set_custom(&mut self, s: String) {
        self.1 = Some(s);
    }

    fn get_custom(&self) -> Option<String> {
        self.1.clone()
    }
}

#[derive(Clone, Debug)]
struct WrappingAcceptor;

#[trillium::async_trait]
impl<Input: Transport> Acceptor<Input> for WrappingAcceptor {
    type Output = TransportWrapper;
    type Error = std::convert::Infallible;
    async fn accept(&self, input: Input) -> Result<Self::Output, Self::Error> {
        Ok(TransportWrapper(BoxedTransport::new(input), None))
    }
}

fn main() {
    trillium_smol::config().with_acceptor(WrappingAcceptor).run(
        router()
            .get("/set/*", |mut conn: Conn| async move {
                let s = conn.wildcard().unwrap_or_default().to_string();
                let Some(t): Option<&mut TransportWrapper> = conn.downcast_transport_mut() else {
                    return conn.with_status(500).with_body("Downcast error\n").halt();
                };
                t.set_custom(format!("{s}\n"));
                conn.ok("Value set\n")
            })
            .get("/get", |conn: Conn| async move {
                let Some(t): Option<&TransportWrapper> = conn.downcast_transport_ref() else {
                    return conn.with_status(500).with_body("Downcast error\n").halt();
                };
                let Some(s) = t.get_custom() else {
                    return conn.with_status(404).with_body("Value not set\n").halt();
                };
                conn.ok(s)
            }),
    );
}

$ curl http://localhost:8080/get
Value not set
$ curl http://localhost:8080/set/hello
Value set
$ curl http://localhost:8080/get
Value not set
$ curl http://localhost:8080/set/hello http://localhost:8080/get
Value set
hello
$ curl http://localhost:8080/get
Value not set

[jbr]

Great point about Acceptor! I hadn't thought of that option. I assume in practice you'd need to wrap another Acceptor? Initially I was considering making Acceptor work a bit like a "handler for transports" like a sequence of mapping transformations, but decided to let them just wrap each other if needed. It might be worth revisiting that at some point.

One small improvement to reduce the boilerplate: trillium_macros offers a derive for forwarding AsyncRead and AsyncWrite. Using it looks like this: https://gist.github.com/jbr/2a7fa1ef8e08102bfdc472a359229e20

[joshtriplett]

I assume in practice you'd need to wrap another Acceptor?

If doing HTTPS, yeah. That's pretty easily done, though.

Initially I was considering making Acceptor work a bit like a "handler for transports" like a sequence of mapping transformations, but decided to let them just wrap each other if needed. It might be worth revisiting that at some point.

I think just wrapping is quite easy, considering Acceptor only requires implementing one function and two associated types. And I think a stack of Acceptors will be rare enough that it's probably not worth making that substantially easier.

One small improvement to reduce the boilerplate: trillium_macros offers a derive for forwarding AsyncRead and AsyncWrite.

Nice! That's much better, thank you.

Having a similar macro for deriving Transport around a containing Transport would be nice, as well. I'd expect that most wrappers around a TCP transport, for instance, just want to forward all four methods.

The only other annoyance is that I'd love to avoid double-boxing (having a BoxedTransport inside my transport inside trillium::Conn's BoxedTransport), but then I'd have to name the exact transport type in a generic in order to downcast. I don't know a good way around that though.

Thanks again for suggesting the Transport approach; I think #432 (and a release with that PR in it) would be sufficient to fully solve this problem at this point. I think I'm going to build a generic library on top of that for a Transport with a StateSet.

[joshtriplett]

This has now been made available via #451; closing. Thank you!