bittensor_rs/extrinsics/
registration.rs

1//! # Registration Extrinsics
2//!
3//! Extrinsics for neuron registration on the Bittensor network:
4//! - `serve_axon`: Register an axon endpoint
5//! - `serve_prometheus`: Register a prometheus endpoint
6//! - `burned_register`: Register by burning TAO
7
8use crate::api::api;
9use crate::error::BittensorError;
10use crate::extrinsics::ExtrinsicResponse;
11use subxt::OnlineClient;
12use subxt::PolkadotConfig;
13
14/// Parameters for serving an axon
15#[derive(Debug, Clone)]
16pub struct ServeAxonParams {
17    /// Subnet netuid
18    pub netuid: u16,
19    /// IP version (4 or 6)
20    pub version: u32,
21    /// IP address as a u128
22    pub ip: u128,
23    /// Port number
24    pub port: u16,
25    /// IP type (4 for IPv4, 6 for IPv6)
26    pub ip_type: u8,
27    /// Protocol (e.g., 0 for gRPC)
28    pub protocol: u8,
29    /// Placeholder 1
30    pub placeholder1: u8,
31    /// Placeholder 2
32    pub placeholder2: u8,
33}
34
35impl ServeAxonParams {
36    /// Create params for an IPv4 axon
37    ///
38    /// # Example
39    ///
40    /// ```
41    /// use bittensor_rs::extrinsics::ServeAxonParams;
42    ///
43    /// let params = ServeAxonParams::ipv4(1, "192.168.1.1", 8080);
44    /// assert!(params.is_ok());
45    /// ```
46    pub fn ipv4(netuid: u16, ip: &str, port: u16) -> Result<Self, BittensorError> {
47        let ip_parts: Vec<u8> = ip
48            .split('.')
49            .map(|s| s.parse::<u8>())
50            .collect::<Result<Vec<_>, _>>()
51            .map_err(|_| BittensorError::ConfigError {
52                field: "ip".to_string(),
53                message: format!("Invalid IPv4 address: {}", ip),
54            })?;
55
56        if ip_parts.len() != 4 {
57            return Err(BittensorError::ConfigError {
58                field: "ip".to_string(),
59                message: format!("Invalid IPv4 address: {}", ip),
60            });
61        }
62
63        let ip_u128 = ((ip_parts[0] as u128) << 24)
64            | ((ip_parts[1] as u128) << 16)
65            | ((ip_parts[2] as u128) << 8)
66            | (ip_parts[3] as u128);
67
68        Ok(Self {
69            netuid,
70            version: 4,
71            ip: ip_u128,
72            port,
73            ip_type: 4,
74            protocol: 0,
75            placeholder1: 0,
76            placeholder2: 0,
77        })
78    }
79
80    /// Create params with raw values
81    pub fn new(netuid: u16, version: u32, ip: u128, port: u16, ip_type: u8, protocol: u8) -> Self {
82        Self {
83            netuid,
84            version,
85            ip,
86            port,
87            ip_type,
88            protocol,
89            placeholder1: 0,
90            placeholder2: 0,
91        }
92    }
93}
94
95/// Parameters for serving prometheus
96#[derive(Debug, Clone)]
97pub struct ServePrometheusParams {
98    /// Subnet netuid
99    pub netuid: u16,
100    /// IP version
101    pub version: u32,
102    /// IP address as u128
103    pub ip: u128,
104    /// Port number
105    pub port: u16,
106    /// IP type
107    pub ip_type: u8,
108}
109
110impl ServePrometheusParams {
111    /// Create params for IPv4 prometheus
112    pub fn ipv4(netuid: u16, ip: &str, port: u16) -> Result<Self, BittensorError> {
113        let ip_parts: Vec<u8> = ip
114            .split('.')
115            .map(|s| s.parse::<u8>())
116            .collect::<Result<Vec<_>, _>>()
117            .map_err(|_| BittensorError::ConfigError {
118                field: "ip".to_string(),
119                message: format!("Invalid IPv4 address: {}", ip),
120            })?;
121
122        if ip_parts.len() != 4 {
123            return Err(BittensorError::ConfigError {
124                field: "ip".to_string(),
125                message: format!("Invalid IPv4 address: {}", ip),
126            });
127        }
128
129        let ip_u128 = ((ip_parts[0] as u128) << 24)
130            | ((ip_parts[1] as u128) << 16)
131            | ((ip_parts[2] as u128) << 8)
132            | (ip_parts[3] as u128);
133
134        Ok(Self {
135            netuid,
136            version: 4,
137            ip: ip_u128,
138            port,
139            ip_type: 4,
140        })
141    }
142}
143
144/// Serve an axon endpoint on the network
145///
146/// This registers your neuron's axon endpoint so other neurons can connect to it.
147pub async fn serve_axon<S>(
148    client: &OnlineClient<PolkadotConfig>,
149    signer: &S,
150    params: ServeAxonParams,
151) -> Result<ExtrinsicResponse<()>, BittensorError>
152where
153    S: subxt::tx::Signer<PolkadotConfig>,
154{
155    let call = api::tx().subtensor_module().serve_axon(
156        params.netuid,
157        params.version,
158        params.ip,
159        params.port,
160        params.ip_type,
161        params.protocol,
162        params.placeholder1,
163        params.placeholder2,
164    );
165
166    let tx_hash = client
167        .tx()
168        .sign_and_submit_default(&call, signer)
169        .await
170        .map_err(|e| BittensorError::TxSubmissionError {
171            message: format!("Failed to submit serve_axon: {}", e),
172        })?;
173
174    Ok(ExtrinsicResponse::success()
175        .with_message("Axon served successfully")
176        .with_extrinsic_hash(&format!("{:?}", tx_hash))
177        .with_data(()))
178}
179
180/// Serve a prometheus endpoint on the network
181pub async fn serve_prometheus<S>(
182    client: &OnlineClient<PolkadotConfig>,
183    signer: &S,
184    params: ServePrometheusParams,
185) -> Result<ExtrinsicResponse<()>, BittensorError>
186where
187    S: subxt::tx::Signer<PolkadotConfig>,
188{
189    let call = api::tx().subtensor_module().serve_prometheus(
190        params.netuid,
191        params.version,
192        params.ip,
193        params.port,
194        params.ip_type,
195    );
196
197    let tx_hash = client
198        .tx()
199        .sign_and_submit_default(&call, signer)
200        .await
201        .map_err(|e| BittensorError::TxSubmissionError {
202            message: format!("Failed to submit serve_prometheus: {}", e),
203        })?;
204
205    Ok(ExtrinsicResponse::success()
206        .with_message("Prometheus served successfully")
207        .with_extrinsic_hash(&format!("{:?}", tx_hash))
208        .with_data(()))
209}
210
211/// Register a neuron by burning TAO
212///
213/// This is an alternative to the POW-based registration.
214pub async fn burned_register<S>(
215    client: &OnlineClient<PolkadotConfig>,
216    signer: &S,
217    netuid: u16,
218) -> Result<ExtrinsicResponse<()>, BittensorError>
219where
220    S: subxt::tx::Signer<PolkadotConfig>,
221{
222    let call = api::tx()
223        .subtensor_module()
224        .burned_register(netuid, signer.account_id());
225
226    let tx_hash = client
227        .tx()
228        .sign_and_submit_default(&call, signer)
229        .await
230        .map_err(|e| BittensorError::TxSubmissionError {
231            message: format!("Failed to submit burned_register: {}", e),
232        })?;
233
234    Ok(ExtrinsicResponse::success()
235        .with_message("Burned registration successful")
236        .with_extrinsic_hash(&format!("{:?}", tx_hash))
237        .with_data(()))
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn test_serve_axon_ipv4() {
246        let params = ServeAxonParams::ipv4(1, "192.168.1.1", 8080).unwrap();
247        assert_eq!(params.netuid, 1);
248        assert_eq!(params.port, 8080);
249        assert_eq!(params.ip_type, 4);
250        // 192.168.1.1 = (192 << 24) | (168 << 16) | (1 << 8) | 1
251        let expected_ip = (192u128 << 24) | (168u128 << 16) | (1u128 << 8) | 1u128;
252        assert_eq!(params.ip, expected_ip);
253    }
254
255    #[test]
256    fn test_serve_axon_invalid_ip() {
257        let result = ServeAxonParams::ipv4(1, "invalid", 8080);
258        assert!(result.is_err());
259    }
260
261    #[test]
262    fn test_serve_axon_too_few_octets() {
263        let result = ServeAxonParams::ipv4(1, "192.168.1", 8080);
264        assert!(result.is_err());
265    }
266
267    #[test]
268    fn test_serve_axon_too_many_octets() {
269        let result = ServeAxonParams::ipv4(1, "192.168.1.1.1", 8080);
270        assert!(result.is_err());
271    }
272
273    #[test]
274    fn test_serve_axon_new() {
275        let params = ServeAxonParams::new(1, 4, 0x7f000001, 8080, 4, 0);
276        assert_eq!(params.netuid, 1);
277        assert_eq!(params.version, 4);
278        assert_eq!(params.ip, 0x7f000001);
279        assert_eq!(params.port, 8080);
280        assert_eq!(params.ip_type, 4);
281        assert_eq!(params.protocol, 0);
282        assert_eq!(params.placeholder1, 0);
283        assert_eq!(params.placeholder2, 0);
284    }
285
286    #[test]
287    fn test_serve_axon_debug() {
288        let params = ServeAxonParams::ipv4(1, "127.0.0.1", 8080).unwrap();
289        let debug = format!("{:?}", params);
290        assert!(debug.contains("ServeAxonParams"));
291        assert!(debug.contains("netuid: 1"));
292    }
293
294    #[test]
295    fn test_serve_prometheus_ipv4() {
296        let params = ServePrometheusParams::ipv4(1, "10.0.0.1", 9090).unwrap();
297        assert_eq!(params.netuid, 1);
298        assert_eq!(params.port, 9090);
299    }
300
301    #[test]
302    fn test_serve_prometheus_invalid_ip() {
303        let result = ServePrometheusParams::ipv4(1, "not-an-ip", 9090);
304        assert!(result.is_err());
305    }
306
307    #[test]
308    fn test_serve_prometheus_debug() {
309        let params = ServePrometheusParams::ipv4(1, "0.0.0.0", 9090).unwrap();
310        let debug = format!("{:?}", params);
311        assert!(debug.contains("ServePrometheusParams"));
312    }
313
314    #[test]
315    fn test_serve_axon_clone() {
316        let params = ServeAxonParams::ipv4(1, "192.168.1.1", 8080).unwrap();
317        let cloned = params.clone();
318        assert_eq!(params.netuid, cloned.netuid);
319        assert_eq!(params.ip, cloned.ip);
320    }
321}