add neuron certificate discovery

bittensor/axon.py

@@ -34,6 +34,7 @@
 import warnings
 from inspect import signature, Signature, Parameter
 from typing import List, Optional, Tuple, Callable, Any, Dict, Awaitable
+from .utils import Certificate
 
 import uvicorn
 from fastapi import APIRouter, Depends, FastAPI
@@ -832,7 +833,10 @@ def stop(self) -> "bittensor.axon":
         return self
 
     def serve(
-        self, netuid: int, subtensor: Optional[bittensor.subtensor] = None
+        self,
+        netuid: int,
+        subtensor: Optional[bittensor.subtensor] = None,
+        certificate: Optional[Certificate] = None,
     ) -> "bittensor.axon":
         """
         Serves the Axon on the specified subtensor connection using the configured wallet. This method
@@ -858,7 +862,7 @@ def serve(
             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: bittensor.Synapse):

bittensor/chain_data.py

@@ -34,6 +34,7 @@
 from .utils import networking as net, RAOPERTAO, U16_NORMALIZED_FLOAT
 from .utils.balance import Balance
 from .utils.registration import torch, use_torch
+from .utils import Certificate
 
 custom_rpc_type_registry = {
     "types": {
@@ -73,6 +74,12 @@
                 ["total_daily_return", "Compact<u64>"],
             ],
         },
+        "NeuronCertificate": {
+            "type": "struct",
+            "type_mapping": [
+                ["certificate", "Vec<u8>"],
+            ],
+        },
         "NeuronInfo": {
             "type": "struct",
             "type_mapping": [
@@ -333,6 +340,7 @@ class ChainDataType(Enum):
     IPInfo = 7
     SubnetHyperparameters = 8
     ScheduledColdkeySwapInfo = 9
+    NeuronCertificate = 10
 
 
 def from_scale_encoding(
@@ -540,6 +548,21 @@ def from_weights_bonds_and_neuron_lite(
         return cls(**n_dict)
 
 
+# Dataclasses for chain data.
+@dataclass
+class NeuronCertificate:
+    r"""
+    Dataclass for neuron certificate.
+    """
+
+    certificate: Certificate
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> "NeuronCertificate":
+        r"""Returns a NeuronCertificate object from a ``vec_u8``."""
+        return from_scale_encoding(vec_u8, ChainDataType.NeuronCertificate)
+
+
 @dataclass
 class NeuronInfoLite:
     """Dataclass for neuron metadata, but without the weights and bonds."""

bittensor/extrinsics/serving.py

@@ -26,6 +26,8 @@
 import bittensor.utils.networking as net
 from bittensor.utils import format_error_message
 from ..errors import MetadataError
+from typing import Optional
+from ..utils import Certificate
 
 
 def serve_extrinsic(
@@ -40,6 +42,7 @@ def serve_extrinsic(
     wait_for_inclusion: bool = False,
     wait_for_finalization=True,
     prompt: bool = False,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     r"""Subscribes a Bittensor endpoint to the subtensor chain.
 
@@ -50,6 +53,8 @@ def serve_extrinsic(
             Endpoint host port i.e., ``192.122.31.4``.
         port (int):
             Endpoint port number i.e., ``9221``.
+        certificate (str):
+            TLS Certificate
         protocol (int):
             An ``int`` representation of the protocol.
         netuid (int):
@@ -81,6 +86,7 @@ def serve_extrinsic(
         "protocol": protocol,
         "placeholder1": placeholder1,
         "placeholder2": placeholder2,
+        "certificate": certificate,
     }
     bittensor.logging.debug("Checking axon ...")
     neuron = subtensor.get_neuron_for_pubkey_and_subnet(
@@ -148,6 +154,7 @@ def serve_axon_extrinsic(
     wait_for_inclusion: bool = False,
     wait_for_finalization: bool = True,
     prompt: bool = False,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     r"""Serves the axon to the network.
 
@@ -200,6 +207,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/subtensor.py

@@ -45,6 +45,7 @@
 from bittensor.utils import torch, weight_utils, format_error_message
 from .chain_data import (
     DelegateInfoLite,
+    NeuronCertificate,
     NeuronInfo,
     DelegateInfo,
     PrometheusInfo,
@@ -114,6 +115,7 @@
     ss58_to_vec_u8,
     U64_NORMALIZED_FLOAT,
     networking,
+    Certificate,
 )
 from .utils.balance import Balance
 from .utils.registration import POWSolution
@@ -1813,6 +1815,7 @@ def serve(
         placeholder2: int = 0,
         wait_for_inclusion: bool = False,
         wait_for_finalization=True,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers a neuron's serving endpoint on the Bittensor network. This function announces the
@@ -1849,6 +1852,7 @@ def serve(
             placeholder2,
             wait_for_inclusion,
             wait_for_finalization,
+            certificate=certificate,
         )
 
     def serve_axon(
@@ -1857,6 +1861,7 @@ def serve_axon(
         axon: "bittensor.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
@@ -1876,7 +1881,12 @@ def serve_axon(
         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=certificate,
         )
 
     def _do_serve_axon(
@@ -1903,11 +1913,17 @@ def _do_serve_axon(
         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"
+
         @retry(delay=1, tries=3, backoff=2, max_delay=4, logger=_logger)
         def make_substrate_call_with_retry():
             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(
@@ -5022,6 +5038,41 @@ def make_substrate_call_with_retry():
 
         return NeuronInfo.from_vec_u8(result)
 
+    def get_neuron_certificate(
+        self, uid: int, 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:
+            uid (int): The unique identifier of the neuron.
+            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
+        """
+
+        @retry(delay=2, tries=3, backoff=2, max_delay=4, logger=_logger)
+        def make_substrate_call_with_retry():
+            block_hash = None if block == None else self.substrate.get_block_hash(block)
+            params = [netuid, uid]
+            if block_hash:
+                params = params + [block_hash]
+            return self.substrate.rpc_request(
+                method="neuronInfo_getNeuronCertificate",
+                params=params,  # custom rpc method
+            )
+
+        json_body = make_substrate_call_with_retry()
+
+        if not (result := json_body.get("result", None)):
+            return None
+
+        return NeuronCertificate.from_vec_u8(result)
+
     def neurons(self, netuid: int, block: Optional[int] = None) -> List[NeuronInfo]:
         """
         Retrieves a list of all neurons within a specified subnet of the Bittensor network. This function

bittensor/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 .utils import Certificate
 
 
 class AxonServeCallParams(TypedDict):
@@ -28,6 +29,7 @@ class AxonServeCallParams(TypedDict):
     port: int
     ip_type: int
     netuid: int
+    certificate: Optional[Certificate]
 
 
 class PrometheusServeCallParams(TypedDict):

bittensor/utils/__init__.py

@@ -31,6 +31,8 @@
 U16_MAX = 65535
 U64_MAX = 18446744073709551615
 
+Certificate = str
+
 
 def ss58_to_vec_u8(ss58_address: str) -> List[int]:
     ss58_bytes: bytes = bittensor.utils.ss58_address_to_bytes(ss58_address)