Update pool url on redirect

chia/farmer/farmer.py

@@ -6,6 +6,7 @@
 import logging
 import time
 import traceback
+from dataclasses import dataclass
 from math import floor
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, AsyncIterator, ClassVar, Dict, List, Optional, Set, Tuple, Union, cast
@@ -17,7 +18,7 @@
 from chia.daemon.keychain_proxy import KeychainProxy, connect_to_keychain_and_validate, wrap_local_keychain
 from chia.plot_sync.delta import Delta
 from chia.plot_sync.receiver import Receiver
-from chia.pools.pool_config import PoolWalletConfig, add_auth_key, load_pool_config
+from chia.pools.pool_config import PoolWalletConfig, add_auth_key, load_pool_config, update_pool_url
 from chia.protocols import farmer_protocol, harvester_protocol
 from chia.protocols.pool_protocol import (
     AuthenticationPayload,
@@ -64,6 +65,12 @@
 UPDATE_POOL_FARMER_INFO_INTERVAL: int = 300
 
 
+@dataclass(frozen=True)
+class GetPoolInfoResult:
+    pool_info: Dict[str, Any]
+    new_pool_url: Optional[str]
+
+
 def strip_old_entries(pairs: List[Tuple[float, Any]], before: float) -> List[Tuple[float, Any]]:
     for index, [timestamp, points] in enumerate(pairs):
         if timestamp >= before:
@@ -327,16 +334,19 @@ async def plot_sync_callback(self, peer_id: bytes32, delta: Optional[Delta]) ->
         if receiver.initial_sync() or harvester_updated:
             self.state_changed("harvester_update", receiver.to_dict(True))
 
-    async def _pool_get_pool_info(self, pool_config: PoolWalletConfig) -> Optional[Dict[str, Any]]:
+    async def _pool_get_pool_info(self, pool_config: PoolWalletConfig) -> Optional[GetPoolInfoResult]:
         try:
             async with aiohttp.ClientSession(trust_env=True) as session:
-                async with session.get(
-                    f"{pool_config.pool_url}/pool_info", ssl=ssl_context_for_root(get_mozilla_ca_crt(), log=self.log)
-                ) as resp:
+                url = f"{pool_config.pool_url}/pool_info"
+                async with session.get(url, ssl=ssl_context_for_root(get_mozilla_ca_crt(), log=self.log)) as resp:
                     if resp.ok:
                         response: Dict[str, Any] = json.loads(await resp.text())
                         self.log.info(f"GET /pool_info response: {response}")
-                        return response
+                        new_pool_url: Optional[str] = None
+                        if resp.url != url and all(r.status in {301, 308} for r in resp.history):
+                            new_pool_url = f"{resp.url}".replace("/pool_info", "")
+
+                        return GetPoolInfoResult(pool_info=response, new_pool_url=new_pool_url)
                     else:
                         self.handle_failed_pool_response(
                             pool_config.p2_singleton_puzzle_hash,
@@ -560,15 +570,19 @@ async def update_pool_state(self) -> None:
                 if time.time() >= pool_state["next_pool_info_update"]:
                     pool_state["next_pool_info_update"] = time.time() + UPDATE_POOL_INFO_INTERVAL
                     # Makes a GET request to the pool to get the updated information
-                    pool_info = await self._pool_get_pool_info(pool_config)
-                    if pool_info is not None and "error_code" not in pool_info:
+                    pool_info_result = await self._pool_get_pool_info(pool_config)
+                    if pool_info_result is not None and "error_code" not in pool_info_result.pool_info:
+                        pool_info = pool_info_result.pool_info
                         pool_state["authentication_token_timeout"] = pool_info["authentication_token_timeout"]
                         # Only update the first time from GET /pool_info, gets updated from GET /farmer later
                         if pool_state["current_difficulty"] is None:
                             pool_state["current_difficulty"] = pool_info["minimum_difficulty"]
                     else:
                         pool_state["next_pool_info_update"] = time.time() + UPDATE_POOL_INFO_FAILURE_RETRY_INTERVAL
 
+                    if pool_info_result is not None and pool_info_result.new_pool_url is not None:
+                        update_pool_url(self._root_path, pool_config, pool_info_result.new_pool_url)
+
                 if time.time() >= pool_state["next_farmer_update"]:
                     pool_state["next_farmer_update"] = time.time() + UPDATE_POOL_FARMER_INFO_INTERVAL
                     authentication_token_timeout = pool_state["authentication_token_timeout"]

chia/pools/pool_config.py

@@ -3,7 +3,7 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from typing import List
+from typing import Any, Callable, Dict, List
 
 from chia_rs import G1Element
 
@@ -62,22 +62,62 @@ def load_pool_config(root_path: Path) -> List[PoolWalletConfig]:
 
 # TODO: remove this a few versions after 1.3, since authentication_public_key is deprecated. This is here to support
 # downgrading to versions older than 1.3.
-def add_auth_key(root_path: Path, config_entry: PoolWalletConfig, auth_key: G1Element) -> None:
+def add_auth_key(root_path: Path, pool_wallet_config: PoolWalletConfig, auth_key: G1Element) -> None:
+    def update_auth_pub_key_for_entry(config_entry: Dict[str, Any]) -> bool:
+        auth_key_hex = bytes(auth_key).hex()
+        if config_entry.get("authentication_public_key", "") != auth_key_hex:
+            config_entry["authentication_public_key"] = auth_key_hex
+
+            return True
+
+        return False
+
+    update_pool_config_entry(
+        root_path=root_path,
+        pool_wallet_config=pool_wallet_config,
+        update_closure=update_auth_pub_key_for_entry,
+        update_log_message=f"Updating pool config for auth key: {auth_key}",
+    )
+
+
+def update_pool_url(root_path: Path, pool_wallet_config: PoolWalletConfig, pool_url: str) -> None:
+    def update_pool_url_for_entry(config_entry: Dict[str, Any]) -> bool:
+        if config_entry.get("pool_url", "") != pool_url:
+            config_entry["pool_url"] = pool_url
+
+            return True
+
+        return False
+
+    update_pool_config_entry(
+        root_path=root_path,
+        pool_wallet_config=pool_wallet_config,
+        update_closure=update_pool_url_for_entry,
+        update_log_message=f"Updating pool config for pool_url change: {pool_wallet_config.pool_url} -> {pool_url}",
+    )
+
+
+def update_pool_config_entry(
+    root_path: Path,
+    pool_wallet_config: PoolWalletConfig,
+    update_closure: Callable[[Dict[str, Any]], bool],
+    update_log_message: str,
+) -> None:
     with lock_and_load_config(root_path, "config.yaml") as config:
         pool_list = config["pool"].get("pool_list", [])
         updated = False
         if pool_list is not None:
             for pool_config_dict in pool_list:
                 try:
-                    if hexstr_to_bytes(pool_config_dict["owner_public_key"]) == bytes(config_entry.owner_public_key):
-                        auth_key_hex = bytes(auth_key).hex()
-                        if pool_config_dict.get("authentication_public_key", "") != auth_key_hex:
-                            pool_config_dict["authentication_public_key"] = auth_key_hex
+                    if hexstr_to_bytes(pool_config_dict["owner_public_key"]) == bytes(
+                        pool_wallet_config.owner_public_key
+                    ):
+                        if update_closure(pool_config_dict):
                             updated = True
                 except Exception as e:
                     log.error(f"Exception updating config: {pool_config_dict} {e}")
         if updated:
-            log.info(f"Updating pool config for auth key: {auth_key}")
+            log.info(update_log_message)
             config["pool"]["pool_list"] = pool_list
             save_config(root_path, "config.yaml", config)