add neuron certificate discoverability

bittensor/axon.py

@@ -42,6 +42,7 @@
 from starlette.requests import Request
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from typing import List, Optional, Tuple, Callable, Any, Dict
+from .utils import Certificate
 
 from bittensor.errors import (
     InvalidRequestNameError,
@@ -822,7 +823,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
@@ -848,7 +852,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

@@ -28,6 +28,7 @@
 
 from .utils import networking as net, U16_MAX, U16_NORMALIZED_FLOAT
 from .utils.balance import Balance
+from .utils import Certificate
 
 custom_rpc_type_registry = {
     "types": {
@@ -67,6 +68,12 @@
                 ["total_daily_return", "Compact<u64>"],
             ],
         },
+        "NeuronCertificate": {
+            "type": "struct",
+            "type_mapping": [
+                ["certificate", "Vec<u8>"],
+            ],
+        },
         "NeuronInfo": {
             "type": "struct",
             "type_mapping": [
@@ -286,6 +293,7 @@ class ChainDataType(Enum):
     StakeInfo = 6
     IPInfo = 7
     SubnetHyperparameters = 8
+    NeuronCertificate = 9
 
 
 # Constants
@@ -514,6 +522,22 @@ def _neuron_dict_to_namespace(neuron_dict) -> "NeuronInfo":
             return neuron
 
 
+# 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:
     r"""

bittensor/extrinsics/serving.py

@@ -20,6 +20,8 @@
 import bittensor.utils.networking as net
 from rich.prompt import Confirm
 from ..errors import MetadataError
+from typing import Optional
+from ..utils import Certificate
 
 
 def serve_extrinsic(
@@ -34,6 +36,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.
 
@@ -44,6 +47,8 @@ def serve_extrinsic(
             Endpoint host port i.e., ``192.122.31.4``.
         port (int):
             Endpoint port number i.e., ``9221``.
+        certificate (str):
+            Certificate.
         protocol (int):
             An ``int`` representation of the protocol.
         netuid (int):
@@ -75,6 +80,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(
@@ -144,6 +150,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.
 
@@ -197,6 +204,7 @@ def serve_axon_extrinsic(
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
         prompt=prompt,
+        certificate=certificate,
     )
     return serve_success
 
@@ -264,7 +272,6 @@ def publish_metadata(
 
 
 from retry import retry
-from typing import Optional
 
 
 def get_metadata(self, netuid: int, hotkey: str, block: Optional[int] = None) -> str:

bittensor/subtensor.py

@@ -36,6 +36,7 @@
 # Local imports.
 from .btlogging.defines import BITTENSOR_LOGGER_NAME
 from .chain_data import (
+    NeuronCertificate,
     NeuronInfo,
     DelegateInfo,
     PrometheusInfo,
@@ -82,7 +83,12 @@
 )
 from .extrinsics.root import root_register_extrinsic, set_root_weights_extrinsic
 from .types import AxonServeCallParams, PrometheusServeCallParams
-from .utils import U16_NORMALIZED_FLOAT, ss58_to_vec_u8, U64_NORMALIZED_FLOAT
+from .utils import (
+    U16_NORMALIZED_FLOAT,
+    ss58_to_vec_u8,
+    U64_NORMALIZED_FLOAT,
+    Certificate,
+)
 from .utils.balance import Balance
 from .utils.registration import POWSolution
 
@@ -1328,6 +1334,7 @@ def serve(
         wait_for_inclusion: bool = False,
         wait_for_finalization=True,
         prompt: bool = False,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers a neuron's serving endpoint on the Bittensor network. This function announces the
@@ -1362,6 +1369,7 @@ def serve(
             placeholder2,
             wait_for_inclusion,
             wait_for_finalization,
+            certificate=certificate,
         )
 
     def serve_axon(
@@ -1371,6 +1379,7 @@ def serve_axon(
         wait_for_inclusion: bool = False,
         wait_for_finalization: bool = True,
         prompt: bool = False,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers an Axon serving endpoint on the Bittensor network for a specific neuron. This function
@@ -1391,7 +1400,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(
@@ -1418,12 +1432,18 @@ 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=2, tries=3, backoff=2, max_delay=4)
         def make_substrate_call_with_retry():
             with self.substrate as substrate:
                 call = substrate.compose_call(
                     call_module="SubtensorModule",
-                    call_function="serve_axon",
+                    call_function=call_function,
                     call_params=call_params,
                 )
                 extrinsic = substrate.create_signed_extrinsic(
@@ -3952,6 +3972,43 @@ 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 detailed information about 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)
+        def make_substrate_call_with_retry():
+            with self.substrate as substrate:
+                block_hash = None if block == None else substrate.get_block_hash(block)
+                params = [netuid, uid]
+                if block_hash:
+                    params = params + [block_hash]
+                return substrate.rpc_request(
+                    method="neuronInfo_getNeuronCertificate",
+                    params=params,  # custom rpc method
+                )
+
+        json_body = make_substrate_call_with_retry()
+        result = json_body["result"]
+
+        if result in (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

@@ -30,6 +30,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)