add neuron certificate discovery

bittensor/core/axon.py

@@ -59,7 +59,7 @@
 from bittensor.core.stream import StreamingSynapse
 from bittensor.core.synapse import Synapse, TerminalInfo
 from bittensor.core.threadpool import PriorityThreadPoolExecutor
-from bittensor.utils import networking
+from bittensor.utils import networking, Certificate
 from bittensor.utils.axon_utils import allowed_nonce_window_ns, calculate_diff_seconds
 from bittensor.utils.btlogging import logging
 
@@ -807,7 +807,12 @@ def stop(self) -> "Axon":
         self.started = False
         return self
 
-    def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon":
+    def serve(
+        self,
+        netuid: int,
+        subtensor: Optional["Subtensor"] = None,
+        certificate: Optional[Certificate] = None,
+    ) -> "Axon":
         """
         Serves the Axon on the specified subtensor connection using the configured wallet. This method
         registers the Axon with a specific subnet within the Bittensor network, identified by the ``netuid``.
@@ -832,7 +837,7 @@ def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon":
             to start receiving and processing requests from other neurons.
         """
         if subtensor is not None and hasattr(subtensor, "serve_axon"):
-            subtensor.serve_axon(netuid=netuid, axon=self)
+            subtensor.serve_axon(netuid=netuid, axon=self, certificate=certificate)
         return self
 
     async def default_verify(self, synapse: "Synapse"):

bittensor/core/chain_data/__init__.py

@@ -11,6 +11,7 @@
 from .ip_info import IPInfo
 from .neuron_info import NeuronInfo
 from .neuron_info_lite import NeuronInfoLite
+from .neuron_certificate import NeuronCertificate
 from .prometheus_info import PrometheusInfo
 from .proposal_vote_data import ProposalVoteData
 from .scheduled_coldkey_swap_info import ScheduledColdkeySwapInfo

bittensor/core/chain_data/neuron_certificate.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import List
+
+from bittensor.core.chain_data.utils import from_scale_encoding, ChainDataType
+from bittensor.utils import Certificate
+
+
+# Dataclasses for chain data.
+@dataclass
+class NeuronCertificate:
+    """
+    Dataclass for neuron certificate.
+    """
+
+    certificate: Certificate
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> "NeuronCertificate":
+        """Returns a NeuronCertificate object from a ``vec_u8``."""
+        return from_scale_encoding(vec_u8, ChainDataType.NeuronCertificate)

bittensor/core/chain_data/utils.py

@@ -22,6 +22,7 @@ class ChainDataType(Enum):
     SubnetHyperparameters = 8
     ScheduledColdkeySwapInfo = 9
     AccountId = 10
+    NeuronCertificate = 11
 
 
 def from_scale_encoding(
@@ -178,6 +179,12 @@ def from_scale_encoding_using_type_string(
                 ["pruning_score", "Compact<u16>"],
             ],
         },
+        "NeuronCertificate": {
+            "type": "struct",
+            "type_mapping": [
+                ["certificate", "Vec<u8>"],
+            ],
+        },
         "axon_info": {
             "type": "struct",
             "type_mapping": [

bittensor/core/extrinsics/serving.py

@@ -20,7 +20,12 @@
 from bittensor.core.errors import MetadataError
 from bittensor.core.extrinsics.utils import submit_extrinsic
 from bittensor.core.settings import version_as_int
-from bittensor.utils import format_error_message, networking as net, unlock_key
+from bittensor.utils import (
+    format_error_message,
+    networking as net,
+    unlock_key,
+    Certificate,
+)
 from bittensor.utils.btlogging import logging
 from bittensor.utils.networking import ensure_connected
 
@@ -57,9 +62,15 @@ def do_serve_axon(
     This function is crucial for initializing and announcing a neuron's ``Axon`` service on the network, enhancing the decentralized computation capabilities of Bittensor.
     """
 
+    if call_params["certificate"] is None:
+        del call_params["certificate"]
+        call_function = "serve_axon"
+    else:
+        call_function = "serve_axon_tls"
+
     call = self.substrate.compose_call(
         call_module="SubtensorModule",
-        call_function="serve_axon",
+        call_function=call_function,
         call_params=call_params,
     )
     extrinsic = self.substrate.create_signed_extrinsic(call=call, keypair=wallet.hotkey)
@@ -90,6 +101,7 @@ def serve_extrinsic(
     placeholder2: int = 0,
     wait_for_inclusion: bool = False,
     wait_for_finalization=True,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     """Subscribes a Bittensor endpoint to the subtensor chain.
 
@@ -124,6 +136,7 @@ def serve_extrinsic(
         "protocol": protocol,
         "placeholder1": placeholder1,
         "placeholder2": placeholder2,
+        "certificate": certificate,
     }
     logging.debug("Checking axon ...")
     neuron = subtensor.get_neuron_for_pubkey_and_subnet(
@@ -182,6 +195,7 @@ def serve_axon_extrinsic(
     axon: "Axon",
     wait_for_inclusion: bool = False,
     wait_for_finalization: bool = True,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     """Serves the axon to the network.
 
@@ -224,6 +238,7 @@ def serve_axon_extrinsic(
         protocol=4,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
+        certificate=certificate,
     )
     return serve_success
 

bittensor/core/subtensor.py

@@ -60,6 +60,7 @@
     ss58_to_vec_u8,
     u16_normalized_float,
     hex_to_bytes,
+    Certificate,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -992,6 +993,7 @@ def serve_axon(
         axon: "Axon",
         wait_for_inclusion: bool = False,
         wait_for_finalization: bool = True,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers an ``Axon`` serving endpoint on the Bittensor network for a specific neuron. This function is used to set up the Axon, a key component of a neuron that handles incoming queries and data processing tasks.
@@ -1008,7 +1010,7 @@ def serve_axon(
         By registering an Axon, the neuron becomes an active part of the network's distributed computing infrastructure, contributing to the collective intelligence of Bittensor.
         """
         return serve_axon_extrinsic(
-            self, netuid, axon, wait_for_inclusion, wait_for_finalization
+            self, netuid, axon, wait_for_inclusion, wait_for_finalization, certificate
         )
 
     # metagraph
@@ -1149,6 +1151,41 @@ def get_neuron_for_pubkey_and_subnet(
             block=block,
         )
 
+    def get_neuron_certificate(
+        self, hotkey: str, netuid: int, block: Optional[int] = None
+    ) -> Optional["Certificate"]:
+        """
+        Retrieves the TLS certificate for a specific neuron identified by its unique identifier (UID)
+        within a specified subnet (netuid) of the Bittensor network.
+
+        Args:
+            hotkey (str): The hotkey to query.
+            netuid (int): The unique identifier of the subnet.
+            block (Optional[int], optional): The blockchain block number for the query.
+
+        Returns:
+            Optional[Certificate]: the certificate of the neuron if found, ``None`` otherwise.
+
+        This function is used for certificate discovery for setting up mutual tls communication between neurons
+        """
+
+        certificate = self.query_module(
+            module="SubtensorModule",
+            name="NeuronCertificates",
+            block=block,
+            params=[netuid, hotkey],
+        )
+        try:
+            serialized_certificate = certificate.serialize()
+            if serialized_certificate:
+                return (
+                    chr(serialized_certificate["algorithm"])
+                    + serialized_certificate["public_key"]
+                )
+        except AttributeError:
+            return None
+        return None
+
     @networking.ensure_connected
     def neuron_for_uid(
         self, uid: Optional[int], netuid: int, block: Optional[int] = None

bittensor/core/types.py

@@ -15,7 +15,8 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-from typing import TypedDict
+from typing import TypedDict, Optional
+from bittensor.utils import Certificate
 
 
 class AxonServeCallParams(TypedDict):
@@ -26,6 +27,7 @@ class AxonServeCallParams(TypedDict):
     port: int
     ip_type: int
     netuid: int
+    certificate: Optional[Certificate]
 
 
 class PrometheusServeCallParams(TypedDict):

bittensor/utils/__init__.py

@@ -40,6 +40,8 @@
 U16_MAX = 65535
 U64_MAX = 18446744073709551615
 
+Certificate = str
+
 
 UnlockStatus = namedtuple("UnlockStatus", ["success", "message"])
 

bittensor/utils/mock/subtensor_mock.py

@@ -30,6 +30,7 @@
     PrometheusInfo,
     AxonInfo,
 )
+from bittensor.core.types import AxonServeCallParams, PrometheusServeCallParams
 from bittensor.core.errors import ChainQueryError
 from bittensor.core.subtensor import Subtensor
 from bittensor.utils import RAOPERTAO, u16_normalized_float
@@ -39,26 +40,6 @@
 __GLOBAL_MOCK_STATE__ = {}
 
 
-class AxonServeCallParams(TypedDict):
-    """Axon serve chain call parameters."""
-
-    version: int
-    ip: int
-    port: int
-    ip_type: int
-    netuid: int
-
-
-class PrometheusServeCallParams(TypedDict):
-    """Prometheus serve chain call parameters."""
-
-    version: int
-    ip: int
-    port: int
-    ip_type: int
-    netuid: int
-
-
 BlockNumber = int
 
 

tests/e2e_tests/test_neuron_certificate.py

@@ -0,0 +1,60 @@
+import pytest
+from bittensor.core.subtensor import Subtensor
+from bittensor.core.axon import Axon
+from bittensor.utils.btlogging import logging
+from tests.e2e_tests.utils.chain_interactions import (
+    wait_interval,
+    register_subnet,
+)
+from tests.e2e_tests.utils.e2e_test_utils import (
+    setup_wallet,
+)
+
+
+@pytest.mark.asyncio
+async def test_neuron_certificate(local_chain):
+    """
+    Tests the metagraph
+
+    Steps:
+        1. Register a subnet through Alice
+        2. Serve Alice axon with neuron certificate
+        3. Verify neuron certificate can be retrieved
+    Raises:
+        AssertionError: If any of the checks or verifications fail
+    """
+    logging.info("Testing neuron_certificate")
+    netuid = 1
+
+    # Register root as Alice - the subnet owner and validator
+    alice_keypair, alice_wallet = setup_wallet("//Alice")
+    register_subnet(local_chain, alice_wallet)
+
+    # Verify subnet <netuid> created successfully
+    assert local_chain.query(
+        "SubtensorModule", "NetworksAdded", [netuid]
+    ).serialize(), "Subnet wasn't created successfully"
+
+    subtensor = Subtensor(network="ws://localhost:9945")
+
+    # Register Alice as a neuron on the subnet
+    assert subtensor.burned_register(
+        alice_wallet, netuid
+    ), "Unable to register Alice as a neuron"
+
+    # Serve Alice's axon with a certificate
+    axon = Axon(wallet=alice_wallet)
+    encoded_certificate = "?FAKE_ALICE_CERT"
+    axon.serve(netuid=netuid, subtensor=subtensor, certificate=encoded_certificate)
+
+    await wait_interval(tempo=1, subtensor=subtensor, netuid=netuid)
+
+    # Verify we are getting the correct certificate
+    assert (
+        subtensor.get_neuron_certificate(
+            netuid=netuid, hotkey=alice_keypair.ss58_address
+        )
+        == encoded_certificate
+    )
+
+    logging.info("✅ Passed test_neuron_certificate")