Harden more

1 files changed, 10 insertions, 1 deletions

diff --git a/proxy.py b/proxy.py
index 9bdb76f..e5e0916 100755
--- a/proxy.py
+++ b/proxy.py
@@ -5,10 +5,12 @@ import os
 import pathlib
 import ssl
 import socketserver
+import urllib.parse
 import urllib.request
 
 
 context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+proxied_hosts = []
 
 class Handler(socketserver.BaseRequestHandler):
     def handle(self):
@@ -19,8 +21,12 @@ class Handler(socketserver.BaseRequestHandler):
             absolute_uri = recv.removesuffix("\r\n")
             assert absolute_uri.startswith("gemini://"), f"Request for uri {absolute_uri} does not start with gemini://"
             logging.info(absolute_uri)
+            absolute_uri = urllib.parse.urlparse(absolute_uri)
+            host = absolute_uri.netloc
 
-            request = urllib.request.Request("https://" + absolute_uri.removeprefix("gemini://"))
+            global proxied_hosts
+            assert host in proxied_hosts, f"{host} not in {proxied_hosts}"
+            request = urllib.request.Request(absolute_uri._replace(scheme="https").geturl(), headers={"Host": host})
             request.add_header("Accept", "text/gemini")
             with urllib.request.urlopen(request) as f:
                 content = f.read().decode("UTF8")
@@ -35,6 +41,7 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", default="0.0.0.0")
     parser.add_argument("--port", type=int, default=1965)
+    parser.add_argument("proxied_host", nargs="+")
 
     group = parser.add_mutually_exclusive_group()
     group.add_argument("--certificates-from-path", type=pathlib.Path)
@@ -59,6 +66,8 @@ def main():
 
     context.sni_callback = sni_callback
 
+    global proxied_hosts
+    proxied_hosts = args.proxied_host
     with socketserver.TCPServer((args.host, args.port), Handler) as server:
         server.serve_forever()