Use request host for cookie domain

spec/invidious/helpers_spec.cr

@@ -1,8 +1,36 @@
 require "../spec_helper"
 
 CONFIG = Config.from_yaml(File.open("config/config.example.yml"))
+require "../../src/invidious/user/cookies"
 
 Spectator.describe "Helper" do
+  describe "Invidious::User::Cookies.domain_for_host" do
+    it "keeps the configured cookie domain for matching hosts" do
+      original_domain = CONFIG.domain
+      CONFIG.domain = "example.com"
+
+      begin
+        expect(Invidious::User::Cookies.domain_for_host("example.com")).to eq("example.com")
+        expect(Invidious::User::Cookies.domain_for_host("www.example.com:443")).to eq("example.com")
+        expect(Invidious::User::Cookies.domain_for_host("www.example.com, proxy")).to eq("example.com")
+      ensure
+        CONFIG.domain = original_domain
+      end
+    end
+
+    it "uses host-only cookies for unrelated hosts" do
+      original_domain = CONFIG.domain
+      CONFIG.domain = "example.com"
+
+      begin
+        expect(Invidious::User::Cookies.domain_for_host("example.onion")).to be_nil
+        expect(Invidious::User::Cookies.domain_for_host("example.com.evil.test")).to be_nil
+      ensure
+        CONFIG.domain = original_domain
+      end
+    end
+  end
+
   describe "#produce_channel_search_continuation" do
     it "correctly produces token for searching a specific channel" do
       expect(produce_channel_search_continuation("UCXuqSBlHAE6Xw-yeJA0Tunw", "", 100)).to eq("4qmFsgJqEhhVQ1h1cVNCbEhBRTZYdy15ZUpBMFR1bncaIEVnWnpaV0Z5WTJnd0FUZ0JZQUY2QkVkS2IxaTRBUUE9WgCaAilicm93c2UtZmVlZFVDWHVxU0JsSEFFNlh3LXllSkEwVHVud3NlYXJjaA%3D%3D")

src/invidious/routes/login.cr

@@ -57,7 +57,8 @@ module Invidious::Routes::Login
           sid = Base64.urlsafe_encode(Random::Secure.random_bytes(32))
           Invidious::Database::SessionIDs.insert(sid, email)
 
-          env.response.cookies["SID"] = Invidious::User::Cookies.sid(CONFIG.domain, sid)
+          cookie_domain = Invidious::User::Cookies.domain_for_request(env)
+          env.response.cookies["SID"] = Invidious::User::Cookies.sid(cookie_domain, sid)
         else
           return error_template(401, "Wrong username or password")
         end
@@ -123,7 +124,8 @@ module Invidious::Routes::Login
         view_name = "subscriptions_#{sha256(user.email)}"
         PG_DB.exec("CREATE MATERIALIZED VIEW #{view_name} AS #{MATERIALIZED_VIEW_SQL.call(user.email)}")
 
-        env.response.cookies["SID"] = Invidious::User::Cookies.sid(CONFIG.domain, sid)
+        cookie_domain = Invidious::User::Cookies.domain_for_request(env)
+        env.response.cookies["SID"] = Invidious::User::Cookies.sid(cookie_domain, sid)
 
         if env.request.cookies["PREFS"]?
           user.preferences = env.get("preferences").as(Preferences)

src/invidious/routes/preferences.cr

@@ -226,7 +226,8 @@ module Invidious::Routes::PreferencesRoute
         File.write("config/config.yml", CONFIG.to_yaml)
       end
     else
-      env.response.cookies["PREFS"] = Invidious::User::Cookies.prefs(CONFIG.domain, preferences)
+      cookie_domain = Invidious::User::Cookies.domain_for_request(env)
+      env.response.cookies["PREFS"] = Invidious::User::Cookies.prefs(cookie_domain, preferences)
     end
 
     env.redirect referer
@@ -261,7 +262,8 @@ module Invidious::Routes::PreferencesRoute
         preferences.dark_mode = "dark"
       end
 
-      env.response.cookies["PREFS"] = Invidious::User::Cookies.prefs(CONFIG.domain, preferences)
+      cookie_domain = Invidious::User::Cookies.domain_for_request(env)
+      env.response.cookies["PREFS"] = Invidious::User::Cookies.prefs(cookie_domain, preferences)
     end
 
     if redirect

src/invidious/user/cookies.cr

@@ -8,6 +8,40 @@ struct Invidious::User
     # used in here are not booleans.
     SECURE = (Kemal.config.ssl || CONFIG.https_only) ? true : false
 
+    def domain_for_request(env) : String?
+      return domain_for_host(env.request.headers["X-Forwarded-Host"]? || env.request.headers["Host"]?)
+    end
+
+    def domain_for_host(request_host : String?) : String?
+      configured_domain = CONFIG.domain
+
+      return nil if configured_domain.nil?
+      return configured_domain if request_host.nil? || request_host.empty?
+
+      normalized_config = configured_domain.lchop(".").rchop(".").downcase
+      normalized_host = normalize_host(request_host)
+
+      if normalized_host == normalized_config || normalized_host.ends_with?(".#{normalized_config}")
+        return configured_domain
+      end
+
+      return nil
+    end
+
+    private def normalize_host(request_host : String) : String
+      host = request_host.split(",").first.strip.downcase
+
+      if host.starts_with?("[")
+        if bracket_index = host.index(']')
+          host = host[1, bracket_index - 1]
+        end
+      elsif colon_index = host.index(':')
+        host = host[0, colon_index]
+      end
+
+      return host.rchop(".")
+    end
+
     # Session ID (SID) cookie
     # Parameter "domain" comes from the global config
     def sid(domain : String?, sid) : HTTP::Cookie