From 3235a62c71f3be0b686628ddcd6ba4cea00b62a7 Mon Sep 17 00:00:00 2001
From: Mistyttm <contact@mistyttm.dev>
Date: Mon, 13 Oct 2025 09:41:14 +1000
Subject: [PATCH] nixos/mautrix-discord: Fix registration token generarion

---
 .../services/matrix/mautrix-discord.nix       | 106 +++++++++---------
 nixos/tests/matrix/mautrix-discord.nix        |  48 +++++++-
 2 files changed, 94 insertions(+), 60 deletions(-)

diff --git a/nixos/modules/services/matrix/mautrix-discord.nix b/nixos/modules/services/matrix/mautrix-discord.nix
index f14e7d90196d..f0796ca1df84 100644
--- a/nixos/modules/services/matrix/mautrix-discord.nix
+++ b/nixos/modules/services/matrix/mautrix-discord.nix
@@ -13,6 +13,7 @@ let
 
   settingsFile = "${dataDir}/config.yaml";
   settingsFileUnformatted = format.generate "discord-config-unsubstituted.yaml" cfg.settings;
+  default_token = "This value is generated when generating the registration";
 in
 {
   options = {
@@ -415,66 +416,61 @@ in
           rm -f '${settingsFile}'
           old_umask=$(umask)
           umask 0177
-          envsubst \
-            -o '${settingsFile}' \
-            -i '${settingsFileUnformatted}'
-          config_has_tokens=$(yq '.appservice | has("as_token") and has("hs_token")' '${settingsFile}')
-          registration_already_exists=$([[ -f '${registrationFile}' ]] && echo "true" || echo "false")
-          echo "There are tokens in the config: $config_has_tokens"
-          echo "Registration already existed: $registration_already_exists"
-          # tokens not configured from config/environment file, and registration file
-          # is already generated, override tokens in config to make sure they are not lost
-          if [[ $config_has_tokens == "false" && $registration_already_exists == "true" ]]; then
-            echo "Copying as_token, hs_token from registration into configuration"
-            yq -sY '.[0].appservice.as_token = .[1].as_token
-              | .[0].appservice.hs_token = .[1].hs_token
-              | .[0]' '${settingsFile}' '${registrationFile}' \
-              > '${settingsFile}.tmp'
+
+          envsubst -o '${settingsFile}' -i '${settingsFileUnformatted}'
+
+          # Check if config has tokens or uses defaults
+          as_token=$(yq -r '.appservice.as_token' '${settingsFile}')
+          hs_token=$(yq -r '.appservice.hs_token' '${settingsFile}')
+          config_has_tokens=$([[ "$as_token" != "${default_token}" && "$as_token" != "null" && "$hs_token" != "${default_token}" && "$hs_token" != "null" ]] && echo "true" || echo "false")
+
+          if [[ -f '${registrationFile}' ]]; then
+            registration_exists="true"
+          else
+            registration_exists="false"
+          fi
+
+          echo "Config has tokens: $config_has_tokens, Registration exists: $registration_exists"
+
+          # If config has default tokens but registration exists, restore tokens from registration
+          if [[ $config_has_tokens == "false" && $registration_exists == "true" ]]; then
+            echo "Restoring tokens from existing registration"
+            yq -sY '.[0].appservice.as_token = .[1].as_token | .[0].appservice.hs_token = .[1].hs_token | .[0]' \
+              '${settingsFile}' '${registrationFile}' > '${settingsFile}.tmp'
             mv '${settingsFile}.tmp' '${settingsFile}'
           fi
-          # make sure --generate-registration does not affect config.yaml
+
+          # If config has default tokens and no registration exists, generate new tokens
+          if [[ $config_has_tokens == "false" && $registration_exists == "false" ]]; then
+            echo "Generating new tokens for first-time setup"
+            # Generate random tokens (64 character alphanumeric strings)
+            new_as_token=$(LC_ALL=C tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 64)
+            new_hs_token=$(LC_ALL=C tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 64)
+
+            # Save generated tokens to config
+            yq -Y ".appservice.as_token = \"$new_as_token\" | .appservice.hs_token = \"$new_hs_token\"" \
+              '${settingsFile}' > '${settingsFile}.tmp'
+            mv '${settingsFile}.tmp' '${settingsFile}'
+
+            # Verify tokens were replaced
+            if [[ $(yq -r '.appservice.as_token' '${settingsFile}') == "${default_token}" ]]; then
+              echo "ERROR: Failed to replace default tokens"
+              exit 1
+            fi
+            echo "Successfully generated and saved new tokens"
+          fi
+
+          # Generate registration file with tokens from config
           cp '${settingsFile}' '${settingsFile}.tmp'
           echo "Generating registration file"
-          mautrix-discord \
-            --generate-registration \
-            --config='${settingsFile}.tmp' \
-            --registration='${registrationFile}'
+          mautrix-discord --generate-registration --config='${settingsFile}.tmp' --registration='${registrationFile}'
           rm '${settingsFile}.tmp'
-          # no tokens configured, and new were just generated by generate registration for first time
-          if [[ $config_has_tokens == "false" && $registration_already_exists == "false" ]]; then
-            echo "Copying newly generated as_token, hs_token from registration into configuration"
-            yq -sY '.[0].appservice.as_token = .[1].as_token
-              | .[0].appservice.hs_token = .[1].hs_token
-              | .[0]' '${settingsFile}' '${registrationFile}' \
-              > '${settingsFile}.tmp'
-            mv '${settingsFile}.tmp' '${settingsFile}'
-          fi
-          # make sure --generate-registration does not affect config.yaml
-          cp '${settingsFile}' '${settingsFile}.tmp'
-          echo "Generating registration file"
-          mautrix-discord \
-            --generate-registration \
-            --config='${settingsFile}.tmp' \
-            --registration='${registrationFile}'
-          rm '${settingsFile}.tmp'
-          # no tokens configured, and new were just generated by generate registration for first time
-          if [[ $config_has_tokens == "false" && $registration_already_exists == "false" ]]; then
-            echo "Copying newly generated as_token, hs_token from registration into configuration"
-            yq -sY '.[0].appservice.as_token = .[1].as_token
-              | .[0].appservice.hs_token = .[1].hs_token
-              | .[0]' '${settingsFile}' '${registrationFile}' \
-              > '${settingsFile}.tmp'
-            mv '${settingsFile}.tmp' '${settingsFile}'
-          fi
-          # Make sure correct tokens are in the registration file
-          if [[ $config_has_tokens == "true" || $registration_already_exists == "true" ]]; then
-            echo "Copying as_token, hs_token from configuration to the registration file"
-            yq -sY '.[1].as_token = .[0].appservice.as_token
-              | .[1].hs_token = .[0].appservice.hs_token
-              | .[1]' '${settingsFile}' '${registrationFile}' \
-              > '${registrationFile}.tmp'
-            mv '${registrationFile}.tmp' '${registrationFile}'
-          fi
+
+          # Ensure registration file has the same tokens as config (mautrix-discord may regenerate them)
+          yq -sY '.[1].as_token = .[0].appservice.as_token | .[1].hs_token = .[0].appservice.hs_token | .[1]' \
+            '${settingsFile}' '${registrationFile}' > '${registrationFile}.tmp'
+          mv '${registrationFile}.tmp' '${registrationFile}'
+
           umask $old_umask
           chown :mautrix-discord-registration '${registrationFile}'
           chmod 640 '${registrationFile}'
diff --git a/nixos/tests/matrix/mautrix-discord.nix b/nixos/tests/matrix/mautrix-discord.nix
index 2609005c0ccc..1e53baac7af2 100644
--- a/nixos/tests/matrix/mautrix-discord.nix
+++ b/nixos/tests/matrix/mautrix-discord.nix
@@ -1,4 +1,4 @@
-{ pkgs, ... }:
+{ pkgs, lib, ... }:
 let
   homeserverUrl = "http://homeserver:8008";
 in
@@ -60,9 +60,8 @@ in
                 displayname = "Discord bridge bot";
                 avatar = "mxc://maunium.net/nIdEykemnwdisvHbpxflpDlC";
               };
-              # These will be generated automatically
-              as_token = "generate";
-              hs_token = "generate";
+              # Don't override as_token/hs_token - let them use the default placeholder
+              # which will trigger automatic generation
 
               database = {
                 type = "sqlite3";
@@ -86,6 +85,7 @@ in
 
         environment.systemPackages = [
           pkgs.nettools
+          pkgs.yq
         ];
       };
 
@@ -156,7 +156,45 @@ in
 
     with subtest("verify registration file was created"):
         homeserver.wait_until_succeeds("test -f /var/lib/mautrix-discord/discord-registration.yaml")
-        homeserver.succeed("ls -la /var/lib/mautrix-discord/")
+
+        # Verify tokens were generated and are not default values
+        config_as_token = homeserver.succeed("yq -r '.appservice.as_token' /var/lib/mautrix-discord/config.yaml").strip()
+        config_hs_token = homeserver.succeed("yq -r '.appservice.hs_token' /var/lib/mautrix-discord/config.yaml").strip()
+        reg_as_token = homeserver.succeed("yq -r '.as_token' /var/lib/mautrix-discord/discord-registration.yaml").strip()
+        reg_hs_token = homeserver.succeed("yq -r '.hs_token' /var/lib/mautrix-discord/discord-registration.yaml").strip()
+
+        print(f"Config as_token: {config_as_token[:20]}...")
+        print(f"Config hs_token: {config_hs_token[:20]}...")
+
+        # Verify tokens are not the default placeholder or "generate"
+        assert config_as_token not in ["This value is generated when generating the registration", "generate"], \
+            f"Config as_token was not replaced: {config_as_token}"
+        assert config_hs_token not in ["This value is generated when generating the registration", "generate"], \
+            f"Config hs_token was not replaced: {config_hs_token}"
+
+        # Verify tokens match between config and registration
+        assert config_as_token == reg_as_token, \
+            f"as_token mismatch: config={config_as_token[:20]}... vs reg={reg_as_token[:20]}..."
+        assert config_hs_token == reg_hs_token, \
+            f"hs_token mismatch: config={config_hs_token[:20]}... vs reg={reg_hs_token[:20]}..."
+
+        print("Tokens generated and synchronized correctly")
+
+    with subtest("verify tokens persist after service restart"):
+        # Restart the registration service to simulate rebuild
+        homeserver.succeed("systemctl restart mautrix-discord-registration.service")
+        homeserver.wait_for_unit("mautrix-discord-registration.service")
+
+        # Verify tokens were preserved
+        config_as_token_2 = homeserver.succeed("yq -r '.appservice.as_token' /var/lib/mautrix-discord/config.yaml").strip()
+        config_hs_token_2 = homeserver.succeed("yq -r '.appservice.hs_token' /var/lib/mautrix-discord/config.yaml").strip()
+
+        assert config_as_token_2 == config_as_token, \
+            f"as_token changed after restart: {config_as_token[:20]}... -> {config_as_token_2[:20]}..."
+        assert config_hs_token_2 == config_hs_token, \
+            f"hs_token changed after restart: {config_hs_token[:20]}... -> {config_hs_token_2[:20]}..."
+
+        print("Tokens persisted correctly after restart")
 
     with subtest("verify bridge connects to homeserver"):
         # Give the bridge a moment to connect