From b4acd09865cdb5475bbabce87fe68e942e69cbd3 Mon Sep 17 00:00:00 2001
From: Oliver Stenbom <oliver@stenbom.eu>
Date: Tue, 29 Sep 2026 08:51:12 +0200
Subject: [PATCH] Fix Worker redirect handling for OAuth callbacks

---
 .github/workflows/ci.yml        |  6 ++--
 server-tests/tests/http_test.rs | 49 ++++++++++++++++++++++++++++++++-
 worker/src/handlers/proxy.rs    |  4 ++-
 worker/wrangler.ci.toml         |  2 +-
 worker/wrangler.toml.sample     |  2 +-
 5 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d2f69fd1..8b903769 100644
--- a/.github/workflows/ci.yml
+++ b/.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
diff --git a/server-tests/tests/http_test.rs b/server-tests/tests/http_test.rs
index 3787c2a2..414abb6b 100644
--- a/server-tests/tests/http_test.rs
+++ b/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 {
diff --git a/worker/src/handlers/proxy.rs b/worker/src/handlers/proxy.rs
index 94f25b00..09a94a5a 100644
--- a/worker/src/handlers/proxy.rs
+++ b/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));
 
diff --git a/worker/wrangler.ci.toml b/worker/wrangler.ci.toml
index bc69ec8c..9859bf35 100644
--- a/worker/wrangler.ci.toml
+++ b/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 = [
diff --git a/worker/wrangler.toml.sample b/worker/wrangler.toml.sample
index fe060643..217506b6 100644
--- a/worker/wrangler.toml.sample
+++ b/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 = [