bittensor_rs/
utils.rs

1//! # Bittensor Utilities
2//!
3//! Helper functions for common Bittensor operations including:
4//! - Weight normalization and payload creation
5//! - Cryptographic signature operations
6//! - Unit conversions (TAO/RAO)
7
8use crate::error::BittensorError;
9use crate::types::Hotkey;
10use crate::AccountId;
11use std::str::FromStr;
12use sp_core::{sr25519, Pair};
13
14// Weight-related types
15
16/// Represents a normalized weight for a neuron
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct NormalizedWeight {
19    /// The neuron's UID
20    pub uid: u16,
21    /// The normalized weight value (0 to u16::MAX)
22    pub weight: u16,
23}
24
25/// Normalize weights to sum to u16::MAX
26///
27/// # Arguments
28///
29/// * `weights` - Vector of (uid, weight) pairs
30///
31/// # Returns
32///
33/// Vector of `NormalizedWeight` where weights sum to approximately u16::MAX
34///
35/// # Example
36///
37/// ```
38/// use bittensor_rs::utils::{normalize_weights, NormalizedWeight};
39///
40/// let weights = vec![(0, 100), (1, 100)];
41/// let normalized = normalize_weights(&weights);
42/// assert_eq!(normalized.len(), 2);
43/// ```
44pub fn normalize_weights(weights: &[(u16, u16)]) -> Vec<NormalizedWeight> {
45    if weights.is_empty() {
46        return vec![];
47    }
48
49    let total: u64 = weights.iter().map(|(_, w)| *w as u64).sum();
50    if total == 0 {
51        return weights
52            .iter()
53            .map(|(uid, _)| NormalizedWeight {
54                uid: *uid,
55                weight: 0,
56            })
57            .collect();
58    }
59
60    let target = u16::MAX as u64;
61    weights
62        .iter()
63        .map(|(uid, weight)| {
64            let normalized = ((*weight as u64 * target) / total) as u16;
65            NormalizedWeight {
66                uid: *uid,
67                weight: normalized,
68            }
69        })
70        .collect()
71}
72
73/// Create a set_weights payload for submission to the chain
74///
75/// # Arguments
76///
77/// * `netuid` - The subnet UID
78/// * `weights` - Vector of normalized weights
79/// * `version_key` - Version key for the weights
80///
81/// # Returns
82///
83/// A payload that can be submitted to the chain
84pub fn set_weights_payload(
85    netuid: u16,
86    weights: Vec<NormalizedWeight>,
87    version_key: u64,
88) -> impl subxt::tx::Payload {
89    use crate::api::api;
90
91    let (dests, values): (Vec<u16>, Vec<u16>) =
92        weights.into_iter().map(|w| (w.uid, w.weight)).unzip();
93
94    api::tx()
95        .subtensor_module()
96        .set_weights(netuid, dests, values, version_key)
97}
98
99/// Verify a Bittensor signature
100///
101/// # Arguments
102///
103/// * `hotkey` - The hotkey that supposedly signed the data
104/// * `signature_hex` - Hex-encoded signature
105/// * `data` - The data that was signed
106///
107/// # Returns
108///
109/// * `Ok(())` if the signature is valid
110/// * `Err(BittensorError)` if verification fails
111///
112/// # Example
113///
114/// ```rust,no_run
115/// use bittensor_rs::types::Hotkey;
116/// use bittensor_rs::utils::verify_bittensor_signature;
117///
118/// let hotkey = Hotkey::new("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY".to_string()).unwrap();
119/// let result = verify_bittensor_signature(&hotkey, "abcd...", b"message");
120/// ```
121pub fn verify_bittensor_signature(
122    hotkey: &Hotkey,
123    signature_hex: &str,
124    data: &[u8],
125) -> Result<(), BittensorError> {
126    if signature_hex.is_empty() {
127        return Err(BittensorError::AuthError {
128            message: "Empty signature".to_string(),
129        });
130    }
131
132    if data.is_empty() {
133        return Err(BittensorError::AuthError {
134            message: "Empty data".to_string(),
135        });
136    }
137
138    let signature_bytes = hex::decode(signature_hex).map_err(|e| BittensorError::AuthError {
139        message: format!("Invalid hex signature format: {e}"),
140    })?;
141
142    let account_id =
143        AccountId::from_str(hotkey.as_str()).map_err(|_| BittensorError::InvalidHotkey {
144            hotkey: hotkey.as_str().to_string(),
145        })?;
146
147    if signature_bytes.len() != 64 {
148        return Err(BittensorError::AuthError {
149            message: format!(
150                "Invalid signature length: expected 64 bytes, got {}",
151                signature_bytes.len()
152            ),
153        });
154    }
155
156    let mut signature_array = [0u8; 64];
157    signature_array.copy_from_slice(&signature_bytes);
158
159    let signature = sr25519::Signature::from_raw(signature_array);
160
161    use sp_runtime::traits::Verify;
162
163    let public_key = sr25519::Public::from_raw(account_id.0);
164    let is_valid = signature.verify(data, &public_key);
165
166    if is_valid {
167        Ok(())
168    } else {
169        Err(BittensorError::AuthError {
170            message: "Signature verification failed".to_string(),
171        })
172    }
173}
174
175/// Signature type used by Bittensor (sr25519)
176pub type BittensorSignature = sr25519::Signature;
177
178/// Sign a message with a keypair
179///
180/// # Arguments
181///
182/// * `keypair` - The sr25519 keypair to sign with
183/// * `message` - The message bytes to sign
184///
185/// # Returns
186///
187/// The signature
188pub fn sign_with_keypair(keypair: &sr25519::Pair, message: &[u8]) -> BittensorSignature {
189    keypair.sign(message)
190}
191
192/// Sign a message and return hex-encoded signature
193///
194/// # Arguments
195///
196/// * `keypair` - The sr25519 keypair to sign with
197/// * `message` - The message bytes to sign
198///
199/// # Returns
200///
201/// Hex-encoded signature string
202pub fn sign_message_hex(keypair: &sr25519::Pair, message: &[u8]) -> String {
203    let signature = sign_with_keypair(keypair, message);
204    hex::encode(signature.0)
205}
206
207/// Create a signature using a subxt signer
208///
209/// # Arguments
210///
211/// * `signer` - The signer to use
212/// * `data` - The data to sign
213///
214/// # Returns
215///
216/// Hex-encoded signature string
217pub fn create_signature<T>(signer: &T, data: &[u8]) -> String
218where
219    T: subxt::tx::Signer<subxt::PolkadotConfig>,
220{
221    let signature = signer.sign(data);
222
223    match signature {
224        subxt::utils::MultiSignature::Sr25519(sig) => hex::encode(sig),
225        _ => hex::encode([0u8; 64]),
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    #[test]
234    fn test_normalize_weights_empty() {
235        let result = normalize_weights(&[]);
236        assert!(result.is_empty());
237    }
238
239    #[test]
240    fn test_normalize_weights_zero_weights() {
241        let weights = vec![(0, 0), (1, 0)];
242        let result = normalize_weights(&weights);
243        assert_eq!(result.len(), 2);
244        assert_eq!(result[0].weight, 0);
245        assert_eq!(result[1].weight, 0);
246    }
247
248    #[test]
249    fn test_normalize_weights_equal() {
250        let weights = vec![(0, 100), (1, 100)];
251        let result = normalize_weights(&weights);
252        assert_eq!(result.len(), 2);
253        // Each should be approximately half of u16::MAX
254        assert!(result[0].weight > 30000);
255        assert!(result[1].weight > 30000);
256    }
257
258    #[test]
259    fn test_normalize_weights_unequal() {
260        let weights = vec![(0, 75), (1, 25)];
261        let result = normalize_weights(&weights);
262        assert_eq!(result.len(), 2);
263        // First should be ~3x the second
264        assert!(result[0].weight > result[1].weight * 2);
265    }
266
267    #[test]
268    fn test_signature_verification_empty_signature() {
269        let hotkey =
270            Hotkey::new("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY".to_string()).unwrap();
271        let result = verify_bittensor_signature(&hotkey, "", b"data");
272        assert!(result.is_err());
273    }
274
275    #[test]
276    fn test_signature_verification_empty_data() {
277        let hotkey =
278            Hotkey::new("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY".to_string()).unwrap();
279        let result = verify_bittensor_signature(&hotkey, &"ab".repeat(64), b"");
280        assert!(result.is_err());
281    }
282
283    #[test]
284    fn test_signature_verification_invalid_hex() {
285        let hotkey =
286            Hotkey::new("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY".to_string()).unwrap();
287        let result = verify_bittensor_signature(&hotkey, "not_hex!", b"data");
288        assert!(result.is_err());
289    }
290
291    #[test]
292    fn test_signature_verification_wrong_length() {
293        let hotkey =
294            Hotkey::new("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY".to_string()).unwrap();
295        let result = verify_bittensor_signature(&hotkey, "abcd", b"data");
296        assert!(result.is_err());
297        assert!(result.unwrap_err().to_string().contains("length"));
298    }
299
300    #[test]
301    fn test_sign_and_verify() {
302        use sp_core::Pair;
303
304        // Generate a keypair
305        let (pair, _) = sr25519::Pair::generate();
306        let message = b"test message";
307
308        // Sign the message
309        let signature = sign_with_keypair(&pair, message);
310
311        // Verify using sp_runtime
312        use sp_runtime::traits::Verify;
313        let public = pair.public();
314        assert!(signature.verify(message.as_slice(), &public));
315    }
316
317    #[test]
318    fn test_sign_message_hex() {
319        use sp_core::Pair;
320
321        let (pair, _) = sr25519::Pair::generate();
322        let message = b"test message";
323
324        let hex_sig = sign_message_hex(&pair, message);
325
326        // Hex signature should be 128 characters (64 bytes * 2)
327        assert_eq!(hex_sig.len(), 128);
328
329        // Should be valid hex
330        assert!(hex::decode(&hex_sig).is_ok());
331    }
332}
bittensor_rs/utils.rs

bittensor_rs/
utils.rs
