feat: initial implementation of sticky cookie for session persistance when load-balancing

src/handler/handler_main.rs

@@ -1,7 +1,7 @@
 // Highly motivated by https://github.com/felipenoris/hyper-reverse-proxy
-use super::{utils_headers::*, utils_request::*, utils_synth_response::*};
+use super::{utils_headers::*, utils_request::*, utils_synth_response::*, HandlerContext};
 use crate::{
-  backend::{Backend, UpstreamGroup},
+  backend::{Backend, LoadBalance, UpstreamGroup},
   error::*,
   globals::Globals,
   log::*,
@@ -91,7 +91,7 @@ where
     let request_upgraded = req.extensions_mut().remove::<hyper::upgrade::OnUpgrade>();
 
     // Build request from destination information
-    if let Err(e) = self.generate_request_forwarded(
+    let context = match self.generate_request_forwarded(
       &client_addr,
       &listen_addr,
       &mut req,
@@ -99,8 +99,11 @@ where
       upstream_group,
       tls_enabled,
     ) {
-      error!("Failed to generate destination uri for reverse proxy: {}", e);
-      return self.return_with_error_log(StatusCode::SERVICE_UNAVAILABLE, &mut log_data);
+      Err(e) => {
+        error!("Failed to generate destination uri for reverse proxy: {}", e);
+        return self.return_with_error_log(StatusCode::SERVICE_UNAVAILABLE, &mut log_data);
+      }
+      Ok(v) => v,
     };
     debug!("Request to be forwarded: {:?}", req);
     log_data.xff(&req.headers().get("x-forwarded-for"));
@@ -123,6 +126,15 @@ where
       }
     };
 
+    // Process reverse proxy context generated during the forwarding request generation.
+    if let Some(context_from_lb) = context.context_lb {
+      let res_headers = res_backend.headers_mut();
+      if let Err(e) = set_sticky_cookie_lb_context(res_headers, &context_from_lb) {
+        error!("Failed to append context to the response given from backend: {}", e);
+        return self.return_with_error_log(StatusCode::BAD_GATEWAY, &mut log_data);
+      }
+    }
+
     if res_backend.status() != StatusCode::SWITCHING_PROTOCOLS {
       // Generate response to client
       if self.generate_response_forwarded(&mut res_backend, backend).is_ok() {
@@ -229,7 +241,7 @@ where
     upgrade: &Option<String>,
     upstream_group: &UpstreamGroup,
     tls_enabled: bool,
-  ) -> Result<()> {
+  ) -> Result<HandlerContext> {
     debug!("Generate request to be forwarded");
 
     // Add te: trailer if contained in original request
@@ -265,10 +277,19 @@ where
         .insert(header::HOST, HeaderValue::from_str(&org_host)?);
     };
 
+    /////////////////////////////////////////////
     // Fix unique upstream destination since there could be multiple ones.
-    // TODO: StickyならCookieをここでgetに与える必要
-    // TODO: Stickyで、Cookieが与えられなかったらset-cookie向けにcookieを返す必要。upstreamオブジェクトに含めるのも手。
-    let upstream_chosen = upstream_group.get().ok_or_else(|| anyhow!("Failed to get upstream"))?;
+    let context_to_lb = if let LoadBalance::StickyRoundRobin(lb) = &upstream_group.lb {
+      takeout_sticky_cookie_lb_context(req.headers_mut(), &lb.sticky_config.name)?
+    } else {
+      None
+    };
+    let (upstream_chosen_opt, context_from_lb) = upstream_group.get(&context_to_lb);
+    let upstream_chosen = upstream_chosen_opt.ok_or_else(|| anyhow!("Failed to get upstream"))?;
+    let context = HandlerContext {
+      context_lb: context_from_lb,
+    };
+    /////////////////////////////////////////////
 
     // apply upstream-specific headers given in upstream_option
     let headers = req.headers_mut();
@@ -321,6 +342,6 @@ where
       *req.version_mut() = Version::HTTP_2;
     }
 
-    Ok(())
+    Ok(context)
   }
 }

src/handler/mod.rs

@@ -4,3 +4,10 @@ mod utils_request;
 mod utils_synth_response;
 
 pub use handler_main::{HttpMessageHandler, HttpMessageHandlerBuilder, HttpMessageHandlerBuilderError};
+
+use crate::backend::LbContext;
+
+#[derive(Debug)]
+struct HandlerContext {
+  context_lb: Option<LbContext>,
+}

src/handler/utils_headers.rs

@@ -1,5 +1,5 @@
 use crate::{
-  backend::{UpstreamGroup, UpstreamOption},
+  backend::{LbContext, StickyCookie, StickyCookieValue, UpstreamGroup, UpstreamOption},
   error::*,
   log::*,
   utils::*,
@@ -14,6 +14,74 @@ use std::net::SocketAddr;
 ////////////////////////////////////////////////////
 // Functions to manipulate headers
 
+/// Take sticky cookie header value from request header,
+/// and returns LbContext to be forwarded to LB if exist and if needed.
+/// Removing sticky cookie is needed and it must not be passed to the upstream.
+pub(super) fn takeout_sticky_cookie_lb_context(
+  headers: &mut HeaderMap,
+  expected_cookie_name: &str,
+) -> Result<Option<LbContext>> {
+  let mut headers_clone = headers.clone();
+
+  match headers_clone.entry(hyper::header::COOKIE) {
+    header::Entry::Vacant(_) => Ok(None),
+    header::Entry::Occupied(entry) => {
+      let cookies_iter = entry
+        .iter()
+        .flat_map(|v| v.to_str().unwrap_or("").split(';').map(|v| v.trim()));
+      let (sticky_cookies, without_sticky_cookies): (Vec<_>, Vec<_>) = cookies_iter
+        .into_iter()
+        .partition(|v| v.starts_with(expected_cookie_name));
+      if sticky_cookies.is_empty() {
+        return Ok(None);
+      }
+      if sticky_cookies.len() > 1 {
+        error!("Multiple sticky cookie values in request");
+        return Err(RpxyError::Other(anyhow!(
+          "Invalid cookie: Multiple sticky cookie values"
+        )));
+      }
+      let cookies_passed_to_upstream = without_sticky_cookies.join("; ");
+      let cookie_passed_to_lb = sticky_cookies.first().unwrap();
+      headers.remove(hyper::header::COOKIE);
+      headers.insert(hyper::header::COOKIE, cookies_passed_to_upstream.parse()?);
+
+      let sticky_cookie = StickyCookie {
+        value: StickyCookieValue::try_from(cookie_passed_to_lb, expected_cookie_name)?,
+        info: None,
+      };
+      Ok(Some(LbContext { sticky_cookie }))
+    }
+  }
+}
+
+/// Set-Cookie if LB Sticky is enabled and if cookie is newly created/updated.
+/// Set-Cookie response header could be in multiple lines.
+/// https://developer.mozilla.org/ja/docs/Web/HTTP/Headers/Set-Cookie
+pub(super) fn set_sticky_cookie_lb_context(headers: &mut HeaderMap, context_from_lb: &LbContext) -> Result<()> {
+  let sticky_cookie_string: String = context_from_lb.sticky_cookie.clone().try_into()?;
+  let new_header_val: HeaderValue = sticky_cookie_string.parse()?;
+  let expected_cookie_name = &context_from_lb.sticky_cookie.value.name;
+  match headers.entry(hyper::header::SET_COOKIE) {
+    header::Entry::Vacant(entry) => {
+      entry.insert(new_header_val);
+    }
+    header::Entry::Occupied(mut entry) => {
+      let mut flag = false;
+      for e in entry.iter_mut() {
+        if e.to_str().unwrap_or("").starts_with(expected_cookie_name) {
+          *e = new_header_val.clone();
+          flag = true;
+        }
+      }
+      if !flag {
+        entry.append(new_header_val);
+      }
+    }
+  };
+  Ok(())
+}
+
 pub(super) fn apply_upstream_options_to_header(
   headers: &mut HeaderMap,
   _client_addr: &SocketAddr,