From bd08dcda97c7af217887a75dbf160236d241fbd9 Mon Sep 17 00:00:00 2001
From: Max <max@privatevoid.net>
Date: Fri, 10 Nov 2023 00:31:41 +0100
Subject: [PATCH] packages/reflex-cache: decompress xz

---
 .../servers/reflex-cache/reflex_cache/ipfs.py |  7 +++++++
 .../reflex-cache/reflex_cache/nix_cache.py    |  2 +-
 .../reflex_cache/service_handler.py           | 21 +++++++++++++------
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/packages/servers/reflex-cache/reflex_cache/ipfs.py b/packages/servers/reflex-cache/reflex_cache/ipfs.py
index 9e68a7e..d7b855f 100644
--- a/packages/servers/reflex-cache/reflex_cache/ipfs.py
+++ b/packages/servers/reflex-cache/reflex_cache/ipfs.py
@@ -3,6 +3,7 @@ from urllib.parse import quote_plus
 
 import requests
 import requests_unixsocket
+import lzma
 
 
 class IPFSController:
@@ -16,6 +17,12 @@ class IPFSController:
         print(f"Downloading NAR: {nar}")
         code, _, content = self.__nix.try_all("get", nar, hint)
         if code == 200:
+            if nar.endswith(".nar.xz"):
+                print(f"Attempt decompression of {nar}")
+                decompressed = lzma.decompress(content)
+                print(f"Size diff: {len(content)} -> {len(decompressed)}")
+                content = decompressed
+
             upload = {"file": ("FILE", content, "application/octet-stream")}
             try:
                 rIpfs = requests_unixsocket.post(
diff --git a/packages/servers/reflex-cache/reflex_cache/nix_cache.py b/packages/servers/reflex-cache/reflex_cache/nix_cache.py
index 7be6d62..ed72fda 100644
--- a/packages/servers/reflex-cache/reflex_cache/nix_cache.py
+++ b/packages/servers/reflex-cache/reflex_cache/nix_cache.py
@@ -30,7 +30,7 @@ class NixCacheFetcher:
             caches = self.__caches
         for cache in caches:
             try:
-                rCache = fn(f"{cache}{path}")
+                rCache = fn(f"{cache}{path}", allow_redirects=True)
                 if rCache.status_code < bestState:
                     bestState = rCache.status_code
 
diff --git a/packages/servers/reflex-cache/reflex_cache/service_handler.py b/packages/servers/reflex-cache/reflex_cache/service_handler.py
index 1a24179..0fa5aca 100644
--- a/packages/servers/reflex-cache/reflex_cache/service_handler.py
+++ b/packages/servers/reflex-cache/reflex_cache/service_handler.py
@@ -50,11 +50,7 @@ class ReflexHTTPServiceHandler(BaseHTTPRequestHandler):
 
             if resultHash == None:
                 code, cache, _ = self._nix.try_all("head", self.path)
-                if code == 200:
-                    self.send_response(302)
-                    self.send_header("Location", f"{cache}{self.path}")
-                    self.end_headers()
-                else:
+                if code != 200:
                     self.send_response(404)
                     self.end_headers()
                     return
@@ -82,6 +78,13 @@ class ReflexHTTPServiceHandler(BaseHTTPRequestHandler):
                             self._ipfs.ipfs_fetch_task, cb, self.path, cache
                         )
                         self._workSet.add((self.path, f))
+                _, code, resultHash = f.result()
+            else:
+                code = 200
+
+            if code != 200:
+                self.send_response(code)
+                self.end_headers()
                 return
 
             self.send_response(302)
@@ -113,11 +116,15 @@ class ReflexHTTPServiceHandler(BaseHTTPRequestHandler):
             self.send_response(code)
             self.end_headers()
             if code == 200:
-                self.wfile.write(content)
                 if match := re.search(
                     "URL: (nar/[a-z0-9]*\\.nar.*)", content.decode("utf-8")
                 ):
                     nar = f"/{match.group(1)}"
+                    # we decompress xz, so tell Nix to except an uncompressed NAR
+                    if nar.endswith(".xz"):
+                        self.wfile.write(content.replace(b"Compression: xz\n", b"Compression: none\n"))
+                    else:
+                        self.wfile.write(content)
                     if not self._db.get_path(nar):
                         with self._workSetLock:
                             found = False
@@ -139,6 +146,8 @@ class ReflexHTTPServiceHandler(BaseHTTPRequestHandler):
                                     self._ipfs.ipfs_fetch_task, cb, nar
                                 )
                                 self._workSet.add((nar, f))
+                else:
+                    self.wfile.write(content)
             return
 
         else: