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:
+    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)

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
 
@@ -40,6 +45,7 @@ def do_serve_axon(
     call_params: "AxonServeCallParams",
     wait_for_inclusion: bool = False,
     wait_for_finalization: bool = True,
+    certificate: Optional[Certificate] = None,
 ) -> tuple[bool, Optional[dict]]:
     """
     Internal method to submit a serve axon transaction to the Bittensor blockchain. This method creates and submits a transaction, enabling a neuron's ``Axon`` to serve requests on the network.
@@ -57,9 +63,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)
@@ -68,6 +80,7 @@ def do_serve_axon(
         extrinsic=extrinsic,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
+        certificate=certificate,
     )
     if wait_for_inclusion or wait_for_finalization:
         response.process_events()
@@ -90,6 +103,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 +138,7 @@ def serve_extrinsic(
         "protocol": protocol,
         "placeholder1": placeholder1,
         "placeholder2": placeholder2,
+        "certificate": certificate,
     }
     logging.debug("Checking axon ...")
     neuron = subtensor.get_neuron_for_pubkey_and_subnet(
@@ -159,6 +174,7 @@ def serve_extrinsic(
         call_params=params,
         wait_for_finalization=wait_for_finalization,
         wait_for_inclusion=wait_for_inclusion,
+        certificate=certificate,
     )
 
     if wait_for_inclusion or wait_for_finalization:
@@ -182,6 +198,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 +241,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

@@ -43,6 +43,7 @@
     DelegateInfo,
     NeuronInfo,
     NeuronInfoLite,
+    NeuronCertificate,
     PrometheusInfo,
     SubnetHyperparameters,
     SubnetInfo,
@@ -77,6 +78,7 @@
     ss58_to_vec_u8,
     u16_normalized_float,
     hex_to_bytes,
+    Certificate,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -1013,6 +1015,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.
@@ -1029,7 +1032,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
@@ -1170,6 +1173,37 @@ 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],
+        )
+        if not hasattr(certificate, "serialize"):
+            return None
+        certificate = certificate.serialize()
+        if not certificate:
+            return None
+        return certificate.get("certificate", 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

@@ -47,6 +47,7 @@ class AxonServeCallParams(TypedDict):
     port: int
     ip_type: int
     netuid: int
+    certificate: Optional[str]
 
 
 class PrometheusServeCallParams(TypedDict):