Implement cookie rewriting for proxy context

public/sw.js

@@ -61,6 +61,104 @@ function isHeavyCookieSite(url) {
     return HEAVY_COOKIE_DOMAINS.some((domain) => urlStr.includes(domain));
 }
 
+// Helper function to rewrite cookies in responses
+function rewriteResponseCookies(response, requestUrl) {
+    // Get Set-Cookie headers from the response
+    const setCookieHeader = response.headers.get("set-cookie");
+
+    if (!setCookieHeader) {
+        return response;
+    }
+
+    try {
+        // Parse the URL to get the proxy prefix
+        const url = new URL(requestUrl);
+        const pathname = url.pathname;
+        let proxyPrefix = "/";
+
+        // Determine which proxy is being used
+        if (pathname.startsWith("/~/uv/")) {
+            proxyPrefix = "/~/uv";
+        } else if (pathname.startsWith("/~/scramjet/")) {
+            proxyPrefix = "/~/scramjet";
+        }
+
+        // Rewrite the cookie
+        const rewrittenCookie = rewriteCookie(
+            setCookieHeader,
+            requestUrl,
+            location.host,
+            proxyPrefix
+        );
+
+        // Create a new response with rewritten cookies
+        const newHeaders = new Headers(response.headers);
+        newHeaders.set("set-cookie", rewrittenCookie);
+
+        return new Response(response.body, {
+            status: response.status,
+            statusText: response.statusText,
+            headers: newHeaders
+        });
+    } catch (error) {
+        console.error("Error rewriting response cookies:", error);
+        return response;
+    }
+}
+
+// Helper function to rewrite a single cookie
+function rewriteCookie(cookieString, targetUrl, proxyHost, proxyPrefix) {
+    try {
+        // Parse the cookie
+        const parts = cookieString.split(";").map((p) => p.trim());
+        const rewrittenParts = [];
+
+        for (let i = 0; i < parts.length; i++) {
+            const part = parts[i];
+            const lowerPart = part.toLowerCase();
+
+            // Rewrite domain attribute
+            if (lowerPart.startsWith("domain=")) {
+                rewrittenParts.push(`Domain=${proxyHost}`);
+            }
+            // Rewrite path attribute
+            else if (lowerPart.startsWith("path=")) {
+                const pathPrefix = "path=";
+                const pathValue = part.substring(pathPrefix.length);
+                const newPath = pathValue.startsWith(proxyPrefix)
+                    ? pathValue
+                    : proxyPrefix + pathValue;
+                rewrittenParts.push(`Path=${newPath}`);
+            }
+            // Handle SameSite attribute
+            else if (lowerPart.startsWith("samesite=")) {
+                // Set to None for cross-origin contexts
+                rewrittenParts.push("SameSite=None");
+            }
+            // Keep other attributes as-is
+            else {
+                rewrittenParts.push(part);
+            }
+        }
+
+        // Ensure SameSite=None and Secure are set for proxied cookies
+        const hasSameSite = rewrittenParts.some((p) => p.toLowerCase().startsWith("samesite="));
+        const hasSecure = rewrittenParts.some((p) => p.toLowerCase() === "secure");
+
+        if (!hasSameSite) {
+            rewrittenParts.push("SameSite=None");
+        }
+        if (!hasSecure) {
+            rewrittenParts.push("Secure");
+        }
+
+        return rewrittenParts.join("; ");
+    } catch (error) {
+        console.error("Error parsing/rewriting cookie:", error);
+        return cookieString;
+    }
+}
+
 // Helper function to ensure proper CAPTCHA handling
 function enhanceCaptchaRequest(request) {
     // Clone the request to ensure all headers and properties are preserved
@@ -125,6 +223,11 @@ self.addEventListener("fetch", function (event) {
                     response = await fetch(request);
                 }
 
+                // Rewrite cookies in proxied responses
+                if (isProxiedRequest) {
+                    response = rewriteResponseCookies(response, url);
+                }
+
                 // Inject interceptor script into proxied HTML responses
                 if (isProxiedRequest) {
                     response = await injectInterceptorScript(response);

server/cookie-rewrite.ts

@@ -0,0 +1,174 @@
+/**
+ * Cookie Rewrite Middleware for Proxy
+ *
+ * This middleware rewrites cookies from proxied sites to work correctly
+ * in the proxy context by:
+ * 1. Rewriting cookie domains to the proxy domain
+ * 2. Rewriting cookie paths to include the proxy prefix
+ * 3. Setting proper SameSite attributes for cross-origin contexts
+ * 4. Handling special cases for CAPTCHA and verification domains
+ */
+
+import type { FastifyRequest, FastifyReply } from "fastify";
+
+/**
+ * Special domains that require specific cookie handling
+ */
+const CAPTCHA_DOMAINS = [
+    "google.com",
+    "recaptcha.net",
+    "gstatic.com",
+    "hcaptcha.com",
+    "cloudflare.com",
+    "turnstile.cloudflare.com"
+];
+
+/**
+ * Parse a Set-Cookie header into its components
+ */
+interface CookieAttributes {
+    name: string;
+    value: string;
+    domain?: string;
+    path?: string;
+    expires?: string;
+    maxAge?: string;
+    secure?: boolean;
+    httpOnly?: boolean;
+    sameSite?: string;
+}
+
+function parseCookie(cookieString: string): CookieAttributes {
+    const parts = cookieString.split(";").map((part) => part.trim());
+    const [nameValue, ...attributes] = parts;
+    const [name, value] = nameValue.split("=");
+
+    const cookie: CookieAttributes = {
+        name: name.trim(),
+        value: value || ""
+    };
+
+    for (const attr of attributes) {
+        const [key, val] = attr.split("=").map((s) => s.trim());
+        const lowerKey = key.toLowerCase();
+
+        switch (lowerKey) {
+            case "domain":
+                cookie.domain = val;
+                break;
+            case "path":
+                cookie.path = val;
+                break;
+            case "expires":
+                cookie.expires = val;
+                break;
+            case "max-age":
+                cookie.maxAge = val;
+                break;
+            case "secure":
+                cookie.secure = true;
+                break;
+            case "httponly":
+                cookie.httpOnly = true;
+                break;
+            case "samesite":
+                cookie.sameSite = val;
+                break;
+        }
+    }
+
+    return cookie;
+}
+
+/**
+ * Serialize a cookie object back to a Set-Cookie string
+ */
+function serializeCookie(cookie: CookieAttributes): string {
+    let result = `${cookie.name}=${cookie.value}`;
+
+    if (cookie.domain) {
+        result += `; Domain=${cookie.domain}`;
+    }
+    if (cookie.path) {
+        result += `; Path=${cookie.path}`;
+    }
+    if (cookie.expires) {
+        result += `; Expires=${cookie.expires}`;
+    }
+    if (cookie.maxAge) {
+        result += `; Max-Age=${cookie.maxAge}`;
+    }
+    if (cookie.secure) {
+        result += "; Secure";
+    }
+    if (cookie.httpOnly) {
+        result += "; HttpOnly";
+    }
+    if (cookie.sameSite) {
+        result += `; SameSite=${cookie.sameSite}`;
+    }
+
+    return result;
+}
+
+/**
+ * Check if a URL is a CAPTCHA-related domain
+ */
+function isCaptchaDomain(url: string): boolean {
+    const urlLower = url.toLowerCase();
+    return CAPTCHA_DOMAINS.some((domain) => urlLower.includes(domain));
+}
+
+/**
+ * Rewrite a single cookie for the proxy context
+ */
+export function rewriteCookie(
+    cookieString: string,
+    targetUrl: string,
+    proxyHost: string,
+    proxyPrefix: string
+): string {
+    const cookie = parseCookie(cookieString);
+
+    // Rewrite domain to the proxy domain
+    // Always set to proxy domain for proper cookie scoping
+    cookie.domain = proxyHost;
+
+    // Rewrite path to include proxy prefix
+    if (cookie.path) {
+        // Only rewrite if not already prefixed
+        if (!cookie.path.startsWith(proxyPrefix)) {
+            cookie.path = proxyPrefix + cookie.path;
+        }
+    } else {
+        cookie.path = proxyPrefix + "/";
+    }
+
+    // Set SameSite attribute for cross-origin contexts
+    // For proxied content, we need SameSite=None to allow cross-origin cookies
+    if (!cookie.sameSite || cookie.sameSite.toLowerCase() !== "none") {
+        cookie.sameSite = "None";
+        // SameSite=None requires Secure flag
+        cookie.secure = true;
+    }
+
+    return serializeCookie(cookie);
+}
+
+/**
+ * Rewrite all Set-Cookie headers in a response
+ */
+export function rewriteSetCookieHeaders(
+    setCookieHeaders: string | string[] | undefined,
+    targetUrl: string,
+    proxyHost: string,
+    proxyPrefix: string
+): string[] {
+    if (!setCookieHeaders) {
+        return [];
+    }
+
+    const cookies = Array.isArray(setCookieHeaders) ? setCookieHeaders : [setCookieHeaders];
+
+    return cookies.map((cookie) => rewriteCookie(cookie, targetUrl, proxyHost, proxyPrefix));
+}

server/index.ts

@@ -16,6 +16,7 @@ import { createBareServer } from "@tomphttp/bare-server-node";
 import { handler as astroHandler } from "../dist/server/entry.mjs";
 import { createServer, IncomingMessage, ServerResponse } from "node:http";
 import { Socket } from "node:net";
+import { rewriteSetCookieHeaders } from "./cookie-rewrite.js";
 
 const bareServer = createBareServer("/bare/", {
     connectionLimiter: {
@@ -118,6 +119,42 @@ await app.register(fastifyStatic, {
 
 await app.register(fastifyMiddie);
 
+// Add cookie rewrite hook for proxy responses
+app.addHook("onSend", async (request: FastifyRequest, reply: FastifyReply, payload) => {
+    // Only apply to proxy routes
+    const isProxyRoute = request.url.startsWith("/~/uv/") || request.url.startsWith("/~/scramjet/");
+
+    if (isProxyRoute) {
+        const setCookieHeaders = reply.getHeader("set-cookie");
+
+        if (setCookieHeaders) {
+            try {
+                // Determine proxy prefix
+                const proxyPrefix = request.url.startsWith("/~/uv/") ? "/~/uv" : "/~/scramjet";
+                const proxyHost = request.hostname;
+
+                // Rewrite the cookies
+                const rewrittenCookies = rewriteSetCookieHeaders(
+                    setCookieHeaders as string | string[],
+                    request.url,
+                    proxyHost,
+                    proxyPrefix
+                );
+
+                // Set the rewritten cookies
+                if (rewrittenCookies.length > 0) {
+                    reply.header("set-cookie", rewrittenCookies);
+                }
+            } catch (error) {
+                console.error("Error rewriting cookies:", error);
+                // Don't fail the request if cookie rewriting fails
+            }
+        }
+    }
+
+    return payload;
+});
+
 await app.use(astroHandler);
 
 app.setNotFoundHandler((req, res) => {