From 61d66a5d45a2dd589119a00fefd6c056945010c0 Mon Sep 17 00:00:00 2001
From: irl <iain@learmonth.me>
Date: Sun, 4 May 2025 15:34:48 +0100
Subject: [PATCH] feat: perform just one pass to replace pool mirrors

---
 src/lua/body_filter.lua   | 27 +++++++++++----------------
 src/lua/header_filter.lua | 24 ++++++++++--------------
 src/lua/utils.lua         |  7 +++++++
 3 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/lua/body_filter.lua b/src/lua/body_filter.lua
index 5e35be4..6563aba 100644
--- a/src/lua/body_filter.lua
+++ b/src/lua/body_filter.lua
@@ -1,3 +1,5 @@
+local utils = require "utils"
+
 local function matomo_tracking_code(site_id)
     return [=[
 <!-- Matomo -->
@@ -21,21 +23,14 @@ local function matomo_tracking_code(site_id)
     ]=]
 end
 
-local function rewrite_body(body, eof)
-    if not ngx.ctx.jasima_pool_map or ngx.ctx.jasima_config.rewrite_disable then
-        return body
+local function rewrite_body(body)
+    -- Rewrite links for assets and outbound links to other mirrored sites
+    local pool_map = ngx.ctx.jasima_pool_map
+    if pool_map then
+        body = body:gsub("//([a-zA-Z0-9%.%-]+%.[a-zA-Z0-9]+)/", utils.get_mirror)
     end
-    for from, to in pairs(ngx.ctx.jasima_pool_map) do
-        if ngx.ctx.jasima_config.rewrite_case_insensitive then
-            local pattern = ngx.re.escape(from)
-            body = ngx.re.gsub(body, pattern, to, "ijo")
-        else
-            -- We expect that str:match("^[%w%-%.]+$") ~= nil
-            local pattern = from:gsub("([%^%$%(%)%%%.%[%]%*%+%-%?])", "%%%1") -- escape Lua patterns
-            body = body:gsub(pattern, to)
-        end
-    end
-    if eof and ngx.ctx.jasima_config.matomo_site_id then
+    -- Add Matomo tracking code
+    if ngx.ctx.jasima_config.matomo_site_id then
         body = body:gsub("</body>", matomo_tracking_code(ngx.ctx.jasima_config.matomo_site_id))
         -- TODO: Ensure that tracking code was added when it's HTML, but only for HTML
     end
@@ -65,13 +60,13 @@ if ngx.ctx.rewriting then
 
     if #ngx.ctx.buffered > 5 * 1024 * 1024 and not eof then
         -- Don't just consume memory forever
-        ngx.arg[1] = rewrite_body(ngx.ctx.buffered, eof)  -- We still do our best
+        ngx.arg[1] = rewrite_body(ngx.ctx.buffered)  -- We still do our best
         ngx.ctx.rewriting = false
         return
     end
 
     if eof then
-        ngx.arg[1] = rewrite_body(ngx.ctx.buffered, eof)
+        ngx.arg[1] = rewrite_body(ngx.ctx.buffered)
     else
         ngx.arg[1] = nil
     end
diff --git a/src/lua/header_filter.lua b/src/lua/header_filter.lua
index f1735f2..7c91bcc 100644
--- a/src/lua/header_filter.lua
+++ b/src/lua/header_filter.lua
@@ -1,4 +1,9 @@
--- Determine if we rewrite the body of this type of content
+local utils = require "utils"
+
+if ngx.ctx.jasima_config.rewrite_disable then
+    return
+end
+
 if ngx.header["Content-Type"] then
     local content_type = ngx.header["Content-Type"]
     if content_type:find("text/html") or
@@ -14,19 +19,10 @@ if ngx.header["Content-Type"] then
     end
 end
 
--- Rewrite any redirects from the origin to point at a mirror
 if ngx.header["Location"] then
-    if ngx.ctx.jasima_pool_map and not ngx.ctx.jasima_config.rewrite_disable then
+    local pool_map = ngx.ctx.jasima_pool_map
+    if ngx.ctx.rewriting and pool_map then
         local location = ngx.header["Location"]
-        for from, to in pairs(ngx.ctx.jasima_pool_map) do
-            if ngx.ctx.jasima_config.rewrite_case_insensitive then
-                local pattern = ngx.re.escape(from)
-                location = ngx.re.gsub(location, pattern, to, "ijo")
-            else
-                local pattern = from:gsub("([%^%$%(%)%%%.%[%]%*%+%-%?])", "%%%1") -- escape Lua patterns
-                location = location:gsub(pattern, to)
-            end
-        end
-        ngx.header["Location"] = location
+        ngx.header["Location"] = location:gsub("//([a-zA-Z0-9%.%-]+%.[a-zA-Z0-9]+)/", utils.get_mirror)
     end
-end
\ No newline at end of file
+end
diff --git a/src/lua/utils.lua b/src/lua/utils.lua
index b6f2ae0..827315c 100644
--- a/src/lua/utils.lua
+++ b/src/lua/utils.lua
@@ -72,4 +72,11 @@ function _M.filter_bogons(ip_list)
     return filtered
 end
 
+function _M.get_mirror(host)
+    ngx.log(ngx.DEBUG, "Looking up mirror for " .. host)
+    local host_lower = host:lower()
+    local mirror = ngx.ctx.jasima_pool_map[host_lower]
+    return "//" .. (mirror or host) .. "/"
+end
+
 return _M