cookie auth: allow for custom login logic in fail handler

1 files changed, 80 insertions, 51 deletions

diff --git a/src/modules/auth/cookie.rs b/src/modules/auth/cookie.rs
index de4c76f..091d41e 100644
--- a/src/modules/auth/cookie.rs
+++ b/src/modules/auth/cookie.rs
@@ -8,6 +8,7 @@ use aes_gcm_siv::{
     Nonce,
 };
 use base64::Engine;
+use bytes::Bytes;
 use futures::Future;
 use headers::{Cookie, HeaderMapExt};
 use http_body_util::{combinators::BoxBody, BodyExt};
@@ -73,59 +74,22 @@ impl Node for CookieAuth {
                         _ => (),
                     }
                 }
-                let mut r = Response::new(BoxBody::<_, ServiceError>::new(
-                    String::new().clone().map_err(|_| unreachable!()),
-                ));
-                *r.status_mut() = StatusCode::FOUND;
                 debug!("login attempt for {username:?}");
                 if self.users.authentificate(username, password) {
-                    debug!("login success");
-                    let nonce = [(); 12].map(|_| random::<u8>());
-                    let plaintext = unix_seconds().to_be_bytes();
-                    let mut ciphertext = context
-                        .state
-                        .crypto_key
-                        .encrypt(
-                            Nonce::from_slice(&nonce),
-                            Payload {
-                                msg: &plaintext,
-                                aad: username.as_bytes(),
-                            },
-                        )
-                        .unwrap();
-
-                    ciphertext.extend(nonce);
-                    let auth = base64::engine::general_purpose::URL_SAFE.encode(ciphertext);
-
-                    let mut cookie_opts = String::new();
-                    if let Some(e) = self.expire {
-                        write!(cookie_opts, "; Max-Age={e}").unwrap();
-                    }
-                    if self.secure {
-                        write!(cookie_opts, "; Secure").unwrap();
-                    }
-
-                    r.headers_mut().append(
-                        SET_COOKIE,
-                        HeaderValue::from_str(&format!(
-                            "gnix_username={}{}",
-                            percent_encode(username.as_bytes(), NON_ALPHANUMERIC),
-                            cookie_opts
-                        ))
-                        .unwrap(),
-                    );
-                    r.headers_mut().append(
-                        SET_COOKIE,
-                        HeaderValue::from_str(&format!("gnix_auth={}{}", auth, cookie_opts))
-                            .unwrap(),
-                    );
+                    debug!("login success via creds");
+                    Ok(apply_login_success_headers(
+                        context, &self, referrer, username,
+                    ))
                 } else {
                     debug!("login fail");
+                    let mut r = Response::new(BoxBody::<_, ServiceError>::new(
+                        String::new().clone().map_err(|_| unreachable!()),
+                    ));
+                    *r.status_mut() = StatusCode::FOUND;
+                    r.headers_mut()
+                        .append(LOCATION, referrer.unwrap_or(HeaderValue::from_static("/")));
+                    Ok(r)
                 }
-                r.headers_mut()
-                    .append(LOCATION, referrer.unwrap_or(HeaderValue::from_static("/")));
-
-                Ok(r)
             } else {
                 if let Some(cookie) = request.headers().typed_get::<Cookie>() {
                     if let Some(auth) = cookie.get("gnix_auth") {
@@ -165,15 +129,80 @@ impl Node for CookieAuth {
                         debug!("no auth cookie");
                     }
                 }
-                debug!("unauthorized");
+                debug!("fail handler");
+                let referrer = request.headers().get(REFERER).cloned();
                 let mut r = self.fail.handle(context, request).await?;
-                *r.status_mut() = StatusCode::UNAUTHORIZED;
-                Ok(r)
+                if let Some(username) = r.headers_mut().remove("gnix-auth-success") {
+                    debug!("login success via fail handler");
+                    Ok(apply_login_success_headers(
+                        context,
+                        &self,
+                        referrer,
+                        username.to_str()?,
+                    ))
+                } else {
+                    debug!("unauthorized");
+                    *r.status_mut() = StatusCode::UNAUTHORIZED;
+                    Ok(r)
+                }
             }
         })
     }
 }
 
+fn apply_login_success_headers(
+    context: &mut NodeContext,
+    node: &CookieAuth,
+    referrer: Option<HeaderValue>,
+    username: &str,
+) -> Response<BoxBody<Bytes, ServiceError>> {
+    let nonce = [(); 12].map(|_| random::<u8>());
+    let plaintext = unix_seconds().to_be_bytes();
+    let mut ciphertext = context
+        .state
+        .crypto_key
+        .encrypt(
+            Nonce::from_slice(&nonce),
+            Payload {
+                msg: &plaintext,
+                aad: username.as_bytes(),
+            },
+        )
+        .unwrap();
+
+    ciphertext.extend(nonce);
+    let auth = base64::engine::general_purpose::URL_SAFE.encode(ciphertext);
+
+    let mut cookie_opts = String::new();
+    if let Some(e) = node.expire {
+        write!(cookie_opts, "; Max-Age={e}").unwrap();
+    }
+    if node.secure {
+        write!(cookie_opts, "; Secure").unwrap();
+    }
+
+    let mut r = Response::new(BoxBody::<_, ServiceError>::new(
+        String::new().clone().map_err(|_| unreachable!()),
+    ));
+    *r.status_mut() = StatusCode::FOUND;
+    r.headers_mut()
+        .append(LOCATION, referrer.unwrap_or(HeaderValue::from_static("/")));
+    r.headers_mut().append(
+        SET_COOKIE,
+        HeaderValue::from_str(&format!(
+            "gnix_username={}{}",
+            percent_encode(username.as_bytes(), NON_ALPHANUMERIC),
+            cookie_opts
+        ))
+        .unwrap(),
+    );
+    r.headers_mut().append(
+        SET_COOKIE,
+        HeaderValue::from_str(&format!("gnix_auth={}{}", auth, cookie_opts)).unwrap(),
+    );
+    r
+}
+
 fn unix_seconds() -> u64 {
     SystemTime::now()
         .duration_since(SystemTime::UNIX_EPOCH)