Provide ProxyHttp method to take ownership of the downstream_session in case of skipping proxying (response_sent) #647
Description
What is the problem your feature solves, or the need it fulfills?
Some of the requests handled by a proxy service often should not be sent upstream, but handled by the service itself. Currently, this seems to be mostly done in the ProxyHttp::request_filter, which, in such a case, is to return Ok(true). However, for more complex request handling, I would prefer to use a conveniet web framework (poem in my case).
For that, I need to translate the session.downstream_session into a poem::Request and handle the poem::Response the other way around. I would like to avoid reading the full request body before passing it to the web framework. Also, e.g. in case of websockets, I would need to control the lifetime of the session.downstream_session.
Describe the solution you'd like
It would be great to have a method as part of the ProxyHttp trait to actually take over the ownership of the session.downstream_session like the following. The return type could be the same as from downstream_session.finish().await.ok().flatten() to 1) encourage calling that finish function and 2) to signal if the stream can be reused.
diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs
index 4c8d616..99ed82f 100644
--- a/pingora-proxy/src/proxy_trait.rs
+++ b/pingora-proxy/src/proxy_trait.rs
@@ -515,6 +515,14 @@ pub trait ProxyHttp {
) -> Result<()> {
Ok(())
}
+
+ async fn consume_downstream(
+ &self,
+ downstream_session: Box<HttpSession>,
+ _ctx: &mut Self::CTX,
+ ) -> Option<Stream> {
+ downstream_session.finish().await.ok().flatten()
+ }
}
/// Context struct returned by `fail_to_proxy`.This trait method would only be called if the ProxyHttp returns Ok(true) from request_filter indicating that no further proxying is to be done:
diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
index 695c088..55b8a02 100644
--- a/pingora-proxy/src/lib.rs
+++ b/pingora-proxy/src/lib.rs
@@ -541,20 +541,18 @@ impl<SV> HttpProxy<SV> {
match self.inner.request_filter(&mut session, &mut ctx).await {
Ok(response_sent) => {
if response_sent {
// TODO: log error
self.inner.logging(&mut session, None, &mut ctx).await;
self.cleanup_sub_req(&mut session);
let persistent_settings = HttpPersistentSettings::for_session(&session);
- return session
- .downstream_session
- .finish()
- .await
- .ok()
- .flatten()
+ return self
+ .inner
+ .consume_downstream(session.downstream_session, &mut ctx)
+ .await
.map(|s| ReusedHttpStream::new(s, Some(persistent_settings)));
}
/* else continue */
}
Err(e) => {
return self
.handle_error(session, &mut ctx, e, "Fail to filter request:")Regarding the call to logging: The library user would, in that function, need to account for the fact that the session is to be handled differently (e.g. there likely isn't a status code set yet).
I would be happy to open a PR for such a change including changes to the documentation. Thanks!