Add additional partial headers

chia/_tests/farmer_harvester/test_farmer.py

@@ -7,19 +7,23 @@
 from time import time
 from types import TracebackType
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
+from unittest.mock import ANY
 
 import pytest
 from chia_rs import AugSchemeMPL, G1Element, G2Element, PrivateKey
 from pytest_mock import MockerFixture
 from yarl import URL
 
+from chia import __version__
 from chia._tests.conftest import HarvesterFarmerEnvironment
+from chia._tests.plot_sync.test_delta import dummy_plot
 from chia._tests.util.misc import DataCase, Marks, datacases
 from chia.consensus.default_constants import DEFAULT_CONSTANTS
 from chia.farmer.farmer import UPDATE_POOL_FARMER_INFO_INTERVAL, Farmer, increment_pool_stats, strip_old_entries
+from chia.plot_sync.receiver import Receiver
 from chia.pools.pool_config import PoolWalletConfig
 from chia.protocols import farmer_protocol, harvester_protocol
-from chia.protocols.harvester_protocol import NewProofOfSpace, RespondSignatures
+from chia.protocols.harvester_protocol import NewProofOfSpace, Plot, RespondSignatures
 from chia.protocols.pool_protocol import PoolErrorCode
 from chia.server.ws_connection import WSChiaConnection
 from chia.simulator.block_tools import BlockTools
@@ -85,10 +89,15 @@ def __init__(
 
 class DummyHarvesterPeer:
     return_invalid_response: bool
-    peer_node_id: bytes32 = std_hash(b"1")
+    peer_node_id: bytes32
+    version: str
 
-    def __init__(self, return_valid_response: bool):
-        self.return_invalid_response = return_valid_response
+    def __init__(
+        self, return_invalid_response: bool = False, peer_node_id: bytes32 = std_hash(b"1"), version: str = "1.0.0"
+    ):
+        self.return_invalid_response = return_invalid_response
+        self.peer_node_id = peer_node_id
+        self.version = version
 
     async def send_message(self, arg1: Any) -> None:
         pass
@@ -1202,3 +1211,110 @@ async def test_farmer_pool_info_config_update(
     assert len(config["pool"]["pool_list"]) == 1
     assert config["pool"]["pool_list"][0]["p2_singleton_puzzle_hash"] == p2_singleton_puzzle_hash.hex()
     assert config["pool"]["pool_list"][0]["pool_url"] == case.expected_pool_url_in_config
+
+
+@dataclass
+class DummyReceiver:
+    _total_plot_size: int
+    _total_effective_plot_size: int
+    _plots: Dict[str, Plot]
+
+    def total_plot_size(self) -> int:
+        return self._total_plot_size
+
+    def total_effective_plot_size(self) -> int:
+        return self._total_effective_plot_size
+
+    def plots(self) -> Dict[str, Plot]:
+        return self._plots
+
+
+@dataclass
+class PartialSubmitHeaderCase(DataCase):
+    _id: str
+    farmer_peer_id: bytes32
+    harvester_peer: DummyHarvesterPeer
+    harvester_receiver: Optional[DummyReceiver]
+    expected_headers: Dict[str, str]
+    marks: Marks = ()
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+
+@datacases(
+    PartialSubmitHeaderCase(
+        "without_capacities",
+        farmer_peer_id=bytes32.fromhex("b665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3"),
+        harvester_peer=DummyHarvesterPeer(
+            peer_node_id=bytes32.fromhex("a665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3"),
+            version="1.2.3.asdf42",
+        ),
+        harvester_receiver=None,
+        expected_headers={
+            "User-Agent": f"Chia Blockchain v.{__version__}",
+            "chia-farmer-peer-id": "b665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3",
+            "chia-farmer-version": __version__,
+            "chia-harvester-peer-id": "a665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3",
+            "chia-harvester-version": "1.2.3.asdf42",
+        },
+    ),
+    PartialSubmitHeaderCase(
+        "with_capacities",
+        farmer_peer_id=bytes32.fromhex("b665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3"),
+        harvester_peer=DummyHarvesterPeer(
+            peer_node_id=bytes32.fromhex("a665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3"),
+            version="1.2.3.asdf42",
+        ),
+        harvester_receiver=DummyReceiver(
+            _total_plot_size=1073741824000,
+            _total_effective_plot_size=2168958484480,
+            _plots={"some_path": dummy_plot("some_path"), "some_other_path": dummy_plot("some_other_path")},
+        ),
+        expected_headers={
+            "User-Agent": f"Chia Blockchain v.{__version__}",
+            "chia-farmer-peer-id": "b665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3",
+            "chia-farmer-version": __version__,
+            "chia-harvester-peer-id": "a665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3",
+            "chia-harvester-version": "1.2.3.asdf42",
+            "chia-harvester-raw-capacity-bytes": "1073741824000",
+            "chia-harvester-effective-capacity-bytes": "2168958484480",
+            "chia-harvester-plot-count": "2",
+        },
+    ),
+)
+@pytest.mark.anyio
+async def test_farmer_additional_headers_on_partial_submit(
+    mocker: MockerFixture,
+    farmer_one_harvester: Tuple[List[HarvesterService], FarmerService, BlockTools],
+    case: PartialSubmitHeaderCase,
+) -> None:
+    _, farmer_service, _ = farmer_one_harvester
+    assert farmer_service.rpc_server is not None
+    farmer_api = farmer_service._api
+    farmer_api.farmer.server.node_id = case.farmer_peer_id
+    if case.harvester_receiver is not None:
+        farmer_api.farmer.plot_sync_receivers[case.harvester_peer.peer_node_id] = cast(
+            Receiver, case.harvester_receiver
+        )
+
+    sp, pos, new_pos = create_valid_pos(farmer_api.farmer)
+    assert pos.pool_contract_puzzle_hash is not None
+
+    assert (
+        verify_and_get_quality_string(
+            pos, DEFAULT_CONSTANTS, sp.challenge_hash, sp.challenge_chain_sp, height=uint32(1)
+        )
+        is not None
+    )
+
+    mock_http_post = mocker.patch(
+        "aiohttp.ClientSession.post",
+        return_value=DummyPoolResponse(True, 200, new_difficulty=123),
+    )
+
+    peer = cast(WSChiaConnection, case.harvester_peer)
+    await farmer_api.new_proof_of_space(new_pos, peer)
+
+    mock_http_post.assert_called_once_with(ANY, json=ANY, ssl=ANY, headers=case.expected_headers)

chia/farmer/farmer_api.py

@@ -361,7 +361,10 @@ async def new_proof_of_space(
                             f"{pool_url}/partial",
                             json=post_partial_request.to_json_dict(),
                             ssl=ssl_context_for_root(get_mozilla_ca_crt(), log=self.farmer.log),
-                            headers={"User-Agent": f"Chia Blockchain v.{__version__}"},
+                            headers={
+                                "User-Agent": f"Chia Blockchain v.{__version__}",
+                                **self._get_additional_headers_for_partial(peer),
+                            },
                         ) as resp:
                             if not resp.ok:
                                 self.farmer.log.error(f"Error sending partial to {pool_url}, {resp.status}")
@@ -656,6 +659,24 @@ async def plot_sync_duplicates(self, message: PlotSyncPathList, peer: WSChiaConn
     async def plot_sync_done(self, message: PlotSyncDone, peer: WSChiaConnection) -> None:
         await self.farmer.plot_sync_receivers[peer.peer_node_id].sync_done(message)
 
+    def _get_additional_headers_for_partial(self, harvester_peer: WSChiaConnection) -> Dict[str, str]:
+        headers: Dict[str, str] = {
+            "chia-farmer-peer-id": self.farmer.server.node_id.hex(),
+            "chia-farmer-version": __version__,
+            "chia-harvester-peer-id": harvester_peer.peer_node_id.hex(),
+            "chia-harvester-version": harvester_peer.version,
+        }
+        receiver = self.farmer.plot_sync_receivers.get(harvester_peer.peer_node_id)
+        if receiver is None:
+            return headers
+
+        return {
+            **headers,
+            "chia-harvester-raw-capacity-bytes": f"{receiver.total_plot_size()}",
+            "chia-harvester-effective-capacity-bytes": f"{receiver.total_effective_plot_size()}",
+            "chia-harvester-plot-count": f"{len(receiver.plots())}",
+        }
+
     def _process_respond_signatures(
         self, response: harvester_protocol.RespondSignatures
     ) -> Optional[Union[DeclareProofOfSpace, SignedValues]]: