1  //! DHT Key types and distance metrics
  2  //!
  3  //! All DHT operations use a unified 256-bit key space:
  4  //! - **Node IDs**: BLAKE3 hash of Ed25519 public key
  5  //! - **Content IDs**: BLAKE3 hash of content
  6  //! - **Mailbox IDs**: BLAKE3 hash of purpose prefix + pubkey
  7  
  8  use blake3::Hasher;
  9  use std::cmp::Ordering;
 10  
 11  /// 256-bit DHT key
 12  pub type DhtKey = [u8; 32];
 13  
 14  /// XOR distance between two DHT keys
 15  ///
 16  /// Returns the bitwise XOR of the two keys, which can be compared
 17  /// as an unsigned 256-bit integer for distance ordering.
 18  #[inline]
 19  pub fn distance(a: &DhtKey, b: &DhtKey) -> DhtKey {
 20      let mut result = [0u8; 32];
 21      for i in 0..32 {
 22          result[i] = a[i] ^ b[i];
 23      }
 24      result
 25  }
 26  
 27  /// Compare which key is closer to a target
 28  ///
 29  /// Returns `Ordering::Less` if `a` is closer to `target` than `b`
 30  #[inline]
 31  pub fn closer(target: &DhtKey, a: &DhtKey, b: &DhtKey) -> Ordering {
 32      let da = distance(target, a);
 33      let db = distance(target, b);
 34      da.cmp(&db)
 35  }
 36  
 37  /// Get bucket index for a key relative to our node ID
 38  ///
 39  /// The bucket index is the number of leading zero bits in the XOR distance.
 40  /// Higher index = closer to us (more leading zeros in XOR).
 41  /// For 256-bit keys: buckets 0-255, where bucket 255 contains our closest peers.
 42  #[inline]
 43  pub fn bucket_index(our_id: &DhtKey, their_id: &DhtKey) -> usize {
 44      let d = distance(our_id, their_id);
 45      leading_zeros(&d) as usize
 46  }
 47  
 48  /// Count leading zero bits in a 256-bit key
 49  fn leading_zeros(key: &DhtKey) -> u32 {
 50      for (i, &byte) in key.iter().enumerate() {
 51          if byte != 0 {
 52              return (i as u32) * 8 + byte.leading_zeros();
 53          }
 54      }
 55      256 // All zeros (same key)
 56  }
 57  
 58  /// Derive a node ID from an Ed25519 public key
 59  pub fn node_id(pubkey: &[u8; 32]) -> DhtKey {
 60      blake3::hash(pubkey).into()
 61  }
 62  
 63  /// Derive a content ID from raw data
 64  pub fn content_id(data: &[u8]) -> DhtKey {
 65      blake3::hash(data).into()
 66  }
 67  
 68  /// Derive a mailbox ID from a purpose prefix and public key
 69  ///
 70  /// Used for things like WebRTC signaling inboxes:
 71  /// ```ignore
 72  /// let inbox = mailbox_id(b"abzu-webrtc-inbox", &peer_pubkey);
 73  /// ```
 74  pub fn mailbox_id(prefix: &[u8], pubkey: &[u8; 32]) -> DhtKey {
 75      let mut hasher = Hasher::new();
 76      hasher.update(prefix);
 77      hasher.update(pubkey);
 78      hasher.finalize().into()
 79  }
 80  
 81  /// Helper for sorting nodes by distance to a target
 82  #[derive(Debug, Clone, Eq, PartialEq)]
 83  pub struct NodeByDistance<T> {
 84      pub node: T,
 85      pub distance: DhtKey,
 86  }
 87  
 88  impl<T: Eq> Ord for NodeByDistance<T> {
 89      fn cmp(&self, other: &Self) -> Ordering {
 90          // Reverse ordering for min-heap (closest first)
 91          other.distance.cmp(&self.distance)
 92      }
 93  }
 94  
 95  impl<T: Eq> PartialOrd for NodeByDistance<T> {
 96      fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
 97          Some(self.cmp(other))
 98      }
 99  }
100  
101  #[cfg(test)]
102  mod tests {
103      use super::*;
104  
105      #[test]
106      fn test_distance_self() {
107          let key = [0xABu8; 32];
108          let d = distance(&key, &key);
109          assert_eq!(d, [0u8; 32]);
110      }
111  
112      #[test]
113      fn test_distance_symmetric() {
114          let a = [0x11u8; 32];
115          let b = [0x22u8; 32];
116          assert_eq!(distance(&a, &b), distance(&b, &a));
117      }
118  
119      #[test]
120      fn test_bucket_index_same_key() {
121          let key = [0xFFu8; 32];
122          assert_eq!(bucket_index(&key, &key), 256);
123      }
124  
125      #[test]
126      fn test_bucket_index_one_bit_different() {
127          let a = [0u8; 32];
128          let mut b = [0u8; 32];
129          b[31] = 0x01; // Differ only in last bit
130          // XOR = 0x01 in last byte = 255 leading zeros
131          assert_eq!(bucket_index(&a, &b), 255);
132      }
133  
134      #[test]
135      fn test_bucket_index_first_bit_different() {
136          let a = [0x00u8; 32];
137          let mut b = [0x00u8; 32];
138          b[0] = 0x80; // Differ in first bit
139          // XOR = 0x80 in first byte = 0 leading zeros
140          assert_eq!(bucket_index(&a, &b), 0);
141      }
142  
143      #[test]
144      fn test_closer() {
145          let target = [0x00u8; 32];
146          let mut a = [0x00u8; 32];
147          let mut b = [0x00u8; 32];
148          a[31] = 0x01; // Distance 1
149          b[31] = 0x02; // Distance 2
150          assert_eq!(closer(&target, &a, &b), Ordering::Less);
151      }
152  
153      #[test]
154      fn test_node_id_deterministic() {
155          let pubkey = [0x42u8; 32];
156          let id1 = node_id(&pubkey);
157          let id2 = node_id(&pubkey);
158          assert_eq!(id1, id2);
159      }
160  
161      #[test]
162      fn test_mailbox_id() {
163          let pubkey = [0x42u8; 32];
164          let inbox1 = mailbox_id(b"webrtc", &pubkey);
165          let inbox2 = mailbox_id(b"content", &pubkey);
166          assert_ne!(inbox1, inbox2);
167      }
168  }