Fix Worker redirects breaking OAuth login

.github/workflows/ci.yml

@@ -20,9 +20,9 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
       # Create empty worker build artifacts
-      - run: mkdir -p worker/build/worker
-      - run: touch worker/build/worker/shim.mjs
-      - run: touch worker/build/worker/index.wasm
+      - run: mkdir -p worker/build
+      - run: touch worker/build/index.js
+      - run: touch worker/build/index_bg.wasm
       - run: cargo check --all
       - run: rustup component add clippy
       - run: cargo clippy

server-tests/tests/http_test.rs

@@ -6,7 +6,7 @@ use axum::{
 use http::{StatusCode, header::SET_COOKIE};
 use tokio::net::TcpListener;
 
-use crate::helpers::{ServerKind, seed_session, setup_server};
+use crate::helpers::{ServerKind, create_session_request, post, seed_session, setup_server};
 
 mod helpers;
 
@@ -83,9 +83,56 @@ async fn maintains_multiple_set_cookie_headers() {
     assert_eq!(cookies[1].to_str().unwrap(), "cookie2=value2; Path=/");
 }
 
+#[tokio::test]
+#[ignore = "requires running wrangler dev"]
+async fn worker_preserves_redirect_and_cookies() {
+    let (url, _) = setup_server(ServerKind::Worker).await;
+    let underlying_url = setup_underlying_server("redirect was followed".to_string()).await;
+    let session_name = "redirectcookies";
+    let response = post(
+        format!("{}/linkup/local-session", url),
+        create_session_request(session_name.to_string(), Some(underlying_url)),
+    )
+    .await;
+    assert_eq!(response.status(), reqwest::StatusCode::OK);
+
+    let response = get_session(
+        format!("{}/oauth-callback", url),
+        "example.com".to_string(),
+        session_name.to_string(),
+    )
+    .await;
+    assert_eq!(response.status(), reqwest::StatusCode::SEE_OTHER);
+    assert_eq!(response.headers().get("location").unwrap(), "/signed-in");
+    let cookies: Vec<_> = response.headers().get_all(SET_COOKIE).iter().collect();
+    assert_eq!(cookies.len(), 2);
+    assert!(
+        cookies
+            .iter()
+            .any(|value| *value == "session=authenticated; Path=/; HttpOnly")
+    );
+    assert!(
+        cookies
+            .iter()
+            .any(|value| *value == "oauth_state=; Path=/; Max-Age=0")
+    );
+}
+
 async fn setup_underlying_server(name: String) -> String {
     let app = Router::new()
         .route("/redirect", get(Redirect::temporary("/somethingelse")))
+        .route(
+            "/oauth-callback",
+            get(|| async {
+                (
+                    AppendHeaders([
+                        (SET_COOKIE, "session=authenticated; Path=/; HttpOnly"),
+                        (SET_COOKIE, "oauth_state=; Path=/; Max-Age=0"),
+                    ]),
+                    Redirect::to("/signed-in"),
+                )
+            }),
+        )
         .route(
             "/cookies",
             get(|| async {

worker/src/handlers/proxy.rs

@@ -1,7 +1,7 @@
 use axum::response::IntoResponse;
 use http::{HeaderMap, StatusCode};
 use linkup::{Session, get_additional_headers, get_target_service};
-use worker::{Fetch, RequestInit, wasm_bindgen::JsValue, worker_sys::web_sys};
+use worker::{Fetch, RequestInit, RequestRedirect, wasm_bindgen::JsValue, worker_sys::web_sys};
 
 use crate::{http_error::HttpError, worker_state::WorkerState, ws::handle_ws_resp};
 
@@ -47,6 +47,8 @@ pub async fn handle_all(
     let mut upstream_init = RequestInit::new();
     upstream_init
         .with_method(req.method())
+        // Return redirects and their cookies to the browser, including OAuth callbacks.
+        .with_redirect(RequestRedirect::Manual)
         .with_headers((&request_headers).into())
         .with_body(req.inner().body().map(JsValue::from));
 

worker/wrangler.ci.toml

@@ -1,5 +1,5 @@
 name = "linkup-worker"
-main = "build/worker/shim.mjs"
+main = "build/index.js"
 compatibility_date = "2024-05-30"
 
 kv_namespaces = [

worker/wrangler.toml.sample

@@ -1,5 +1,5 @@
 name = "linkup-worker"
-main = "build/worker/shim.mjs"
+main = "build/index.js"
 compatibility_date = "2024-05-30"
 
 kv_namespaces = [