Rewrite Location header in redirect responses.

Preserve the proxy prefix in the path info when the proxied server sends the client to a location without a scheme or host.

Author: ryanlovett

jupyter_server_proxy/handlers.py

@@ -257,6 +257,42 @@ def _get_context_path(self, host, port):
         else:
             return url_path_join(self.base_url, "proxy", host_and_port)
 
+    def _rewrite_location_header(self, location, host, port, proxied_path):
+        """
+        Rewrite Location header in redirect responses to preserve the proxy prefix.
+
+        When a backend server issues a redirect, the Location header typically contains
+        a path relative to the backend's root. We need to prepend the proxy prefix so
+        the browser navigates to the correct proxied URL.
+
+        For example:
+        - Original Location: /subdir/
+        - Proxy context: /user/{username}/proxy/9000
+        - Rewritten Location: /user/{username}/proxy/9000/subdir/
+        """
+        # Parse the location header
+        parsed = urlparse(location)
+
+        # Only rewrite if the location is a relative path (no scheme or host)
+        if parsed.scheme or parsed.netloc:
+            # Absolute URL - leave as is
+            self.log.debug(f"Not rewriting absolute Location header: {location}")
+            return location
+
+        # Get the proxy context path
+        context_path = self._get_context_path(host, port)
+
+        # Rewrite the path to include the proxy prefix
+        new_path = url_path_join(context_path, parsed.path)
+
+        # Reconstruct the location with the rewritten path
+        rewritten = parsed._replace(path=new_path)
+
+        rewritten_location = urlunparse(rewritten)
+        self.log.info(f"Rewrote Location header: {location} -> {rewritten_location}")
+
+        return rewritten_location
+
     def get_client_uri(self, protocol, host, port, proxied_path):
         if self.absolute_url:
             context_path = self._get_context_path(host, port)
@@ -542,6 +578,15 @@ def rewrite_pe(rewritable_response: RewritableResponse):
             self._headers = httputil.HTTPHeaders()
             for header, v in rewritten_response.headers.get_all():
                 if header not in ("Content-Length", "Transfer-Encoding", "Connection"):
+                    # Rewrite Location header in redirects to preserve proxy prefix
+                    if header == "Location" and rewritten_response.code in (
+                        301,
+                        302,
+                        303,
+                        307,
+                        308,
+                    ):
+                        v = self._rewrite_location_header(v, host, port, proxied_path)
                     # some header appear multiple times, eg 'Set-Cookie'
                     self.add_header(header, v)
 

tests/resources/jupyter_server_config.py

@@ -143,6 +143,9 @@ def my_env():
         "unix_socket": True,
         "raw_socket_proxy": True,
     },
+    "python-redirect": {
+        "command": [sys.executable, _get_path("redirectserver.py"), "--port={port}"],
+    },
 }
 
 c.ServerProxy.non_service_rewrite_response = hello_to_foo

tests/resources/redirectserver.py

@@ -0,0 +1,62 @@
+"""
+Simple webserver that returns 301 redirects to test Location header rewriting.
+"""
+
+import argparse
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from urllib.parse import urlparse, urlunparse
+
+
+class RedirectHandler(BaseHTTPRequestHandler):
+    """Handler that returns 301 redirects with relative Location headers."""
+    
+    def do_GET(self):
+        """
+        Handle GET requests:
+        - Requests without trailing slash: 301 redirect to path with trailing slash
+        - Requests with trailing slash: 200 OK
+        - /redirect-to/target: 301 redirect to /target
+        """
+        # Parse the path to separate path and query string
+        parsed = urlparse(self.path)
+        path = parsed.path
+        query = parsed.query
+        
+        if path.startswith("/redirect-to/"):
+            # Extract the target path (remove /redirect-to prefix)
+            target = path[len("/redirect-to"):]
+            # Preserve query string if present
+            if query:
+                target = f"{target}?{query}"
+            self.send_response(301)
+            self.send_header("Location", target)
+            self.send_header("Content-type", "text/plain")
+            self.end_headers()
+            self.wfile.write(b"Redirecting...\n")
+        elif not path.endswith("/"):
+            # Add trailing slash, preserve query string
+            new_path = path + "/"
+            if query:
+                new_location = f"{new_path}?{query}"
+            else:
+                new_location = new_path
+            self.send_response(301)
+            self.send_header("Location", new_location)
+            self.send_header("Content-type", "text/plain")
+            self.end_headers()
+            self.wfile.write(b"Redirecting...\n")
+        else:
+            # Normal response
+            self.send_response(200)
+            self.send_header("Content-type", "text/plain")
+            self.end_headers()
+            self.wfile.write(f"Success: {self.path}\n".encode())
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--port", type=int, required=True)
+    args = ap.parse_args()
+    
+    httpd = HTTPServer(("127.0.0.1", args.port), RedirectHandler)
+    httpd.serve_forever()

tests/test_proxies.py

@@ -528,3 +528,33 @@ async def test_server_proxy_rawsocket(
         await conn.write_message(msg)
         res = await conn.read_message()
         assert res == msg.swapcase()
+
+
+def test_server_proxy_redirect_location_header_rewrite(
+    a_server_port_and_token: Tuple[int, str],
+) -> None:
+    """
+    Test that Location headers in redirect responses are rewritten to include
+    the proxy prefix.
+
+    This can happen when servers like python's http.server issue 301
+    redirects with relative Location headers (e.g., /subdir/) that don't
+    include the proxy prefix, causing 404 errors.
+    """
+    PORT, TOKEN = a_server_port_and_token
+
+    # Test 1: Named server proxy - redirect without trailing slash
+    r = request_get(PORT, "/python-redirect/mydir", TOKEN)
+    assert r.code == 301
+    location = r.headers.get("Location")
+    # Should be rewritten to include the proxy prefix
+    # The token query parameter should be preserved in the redirect
+    assert location == f"/python-redirect/mydir/?token={TOKEN}"
+
+    # Test 2: Named server proxy - explicit redirect-to endpoint
+    r = request_get(PORT, "/python-redirect/redirect-to/target/path", TOKEN)
+    assert r.code == 301
+    location = r.headers.get("Location")
+    # Should be rewritten to include the proxy prefix
+    # The token query parameter should be preserved in the redirect
+    assert location == f"/python-redirect/target/path?token={TOKEN}"