1  //! Protocol messages for DHT operations
  2  //!
  3  //! These are the wire-level message types that will be serialized into AbzuFrame.
  4  //! Each message has a request and response variant.
  5  
  6  use serde::{Deserialize, Serialize};
  7  use crate::key::DhtKey;
  8  use crate::routing::NodeInfo;
  9  use crate::store::StoredValue;
 10  
 11  /// DHT protocol message — the top-level wrapper
 12  #[derive(Debug, Clone, Serialize, Deserialize)]
 13  pub enum DhtMessage {
 14      Request(Box<DhtRequest>),
 15      Response(DhtResponse),
 16  }
 17  
 18  /// Request types in the DHT protocol
 19  #[derive(Debug, Clone, Serialize, Deserialize)]
 20  pub enum DhtRequest {
 21      /// Ping: "Are you alive?"
 22      Ping {
 23          /// Sender's node info for routing table updates
 24          sender: NodeInfo,
 25      },
 26  
 27      /// FindNode: "What are the K closest nodes to this target?"
 28      FindNode {
 29          sender: NodeInfo,
 30          target: DhtKey,
 31      },
 32  
 33      /// FindValue: "Do you have values for this key? If not, give me closer nodes"
 34      FindValue {
 35          sender: NodeInfo,
 36          key: DhtKey,
 37      },
 38  
 39      /// Store: "Please store this value"
 40      Store {
 41          sender: NodeInfo,
 42          value: StoredValue,
 43      },
 44  
 45      /// GetProviders: "Who provides content with this CID?"
 46      /// (Higher-level convenience over FindValue with ContentProvider type)
 47      GetProviders {
 48          sender: NodeInfo,
 49          cid: DhtKey,
 50          /// Maximum providers to return
 51          max: usize,
 52      },
 53  
 54      /// AddProvider: "I'm providing content with this CID"
 55      AddProvider {
 56          sender: NodeInfo,
 57          cid: DhtKey,
 58          /// Provider announcement (signed)
 59          provider: StoredValue,
 60      },
 61  }
 62  
 63  /// Response types in the DHT protocol
 64  #[derive(Debug, Clone, Serialize, Deserialize)]
 65  pub enum DhtResponse {
 66      /// Response to Ping
 67      Pong {
 68          sender: NodeInfo,
 69      },
 70  
 71      /// Response to FindNode: K closest nodes
 72      Nodes {
 73          nodes: Vec<NodeInfo>,
 74      },
 75  
 76      /// Response to FindValue: either value(s) or closer nodes
 77      Value {
 78          /// Values found (may be empty)
 79          values: Vec<StoredValue>,
 80          /// Closer nodes if no values found
 81          closer_nodes: Vec<NodeInfo>,
 82      },
 83  
 84      /// Response to Store
 85      Stored {
 86          /// Whether the value was stored
 87          success: bool,
 88          /// Error message if not stored
 89          error: Option<String>,
 90      },
 91  
 92      /// Response to GetProviders
 93      Providers {
 94          providers: Vec<StoredValue>,
 95          closer_nodes: Vec<NodeInfo>,
 96      },
 97  
 98      /// Response to AddProvider
 99      ProviderAdded {
100          success: bool,
101          error: Option<String>,
102      },
103  
104      /// Error response (for any request)
105      Error {
106          code: ErrorCode,
107          message: String,
108      },
109  }
110  
111  /// Error codes for DHT operations
112  #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
113  #[repr(u8)]
114  pub enum ErrorCode {
115      /// Unknown or unspecified error
116      Unknown = 0,
117      /// Request was malformed
118      BadRequest = 1,
119      /// Signature verification failed
120      InvalidSignature = 2,
121      /// Value has expired
122      Expired = 3,
123      /// Storage limit reached
124      StorageFull = 4,
125      /// Rate limited
126      RateLimited = 5,
127      /// Not authorized
128      Unauthorized = 6,
129      /// Internal error
130      Internal = 255,
131  }
132  
133  impl DhtRequest {
134      /// Get the sender's node info from any request
135      pub fn sender(&self) -> &NodeInfo {
136          match self {
137              Self::Ping { sender } => sender,
138              Self::FindNode { sender, .. } => sender,
139              Self::FindValue { sender, .. } => sender,
140              Self::Store { sender, .. } => sender,
141              Self::GetProviders { sender, .. } => sender,
142              Self::AddProvider { sender, .. } => sender,
143          }
144      }
145  
146      /// Check if this request should update the routing table
147      pub fn should_update_routing_table(&self) -> bool {
148          // All valid requests should update the routing table
149          true
150      }
151  }
152  
153  impl DhtResponse {
154      /// Create an error response
155      pub fn error(code: ErrorCode, message: impl Into<String>) -> Self {
156          Self::Error {
157              code,
158              message: message.into(),
159          }
160      }
161  }
162  
163  /// Builder for creating DHT messages
164  pub struct MessageBuilder {
165      our_node: NodeInfo,
166  }
167  
168  impl MessageBuilder {
169      /// Create a new builder with our node info
170      pub fn new(our_node: NodeInfo) -> Self {
171          Self { our_node }
172      }
173  
174      /// Create a Ping request
175      pub fn ping(&self) -> DhtMessage {
176          DhtMessage::Request(Box::new(DhtRequest::Ping {
177              sender: self.our_node.clone(),
178          }))
179      }
180  
181      /// Create a FindNode request
182      pub fn find_node(&self, target: DhtKey) -> DhtMessage {
183          DhtMessage::Request(Box::new(DhtRequest::FindNode {
184              sender: self.our_node.clone(),
185              target,
186          }))
187      }
188  
189      /// Create a FindValue request
190      pub fn find_value(&self, key: DhtKey) -> DhtMessage {
191          DhtMessage::Request(Box::new(DhtRequest::FindValue {
192              sender: self.our_node.clone(),
193              key,
194          }))
195      }
196  
197      /// Create a Store request
198      pub fn store(&self, value: StoredValue) -> DhtMessage {
199          DhtMessage::Request(Box::new(DhtRequest::Store {
200              sender: self.our_node.clone(),
201              value,
202          }))
203      }
204  
205      /// Create a GetProviders request
206      pub fn get_providers(&self, cid: DhtKey, max: usize) -> DhtMessage {
207          DhtMessage::Request(Box::new(DhtRequest::GetProviders {
208              sender: self.our_node.clone(),
209              cid,
210              max,
211          }))
212      }
213  
214      /// Create an AddProvider request
215      pub fn add_provider(&self, cid: DhtKey, provider: StoredValue) -> DhtMessage {
216          DhtMessage::Request(Box::new(DhtRequest::AddProvider {
217              sender: self.our_node.clone(),
218              cid,
219              provider,
220          }))
221      }
222  
223      /// Create a Pong response
224      pub fn pong(&self) -> DhtMessage {
225          DhtMessage::Response(DhtResponse::Pong {
226              sender: self.our_node.clone(),
227          })
228      }
229  
230      /// Create a Nodes response
231      pub fn nodes(&self, nodes: Vec<NodeInfo>) -> DhtMessage {
232          DhtMessage::Response(DhtResponse::Nodes { nodes })
233      }
234  
235      /// Create a Value response
236      pub fn value(&self, values: Vec<StoredValue>, closer_nodes: Vec<NodeInfo>) -> DhtMessage {
237          DhtMessage::Response(DhtResponse::Value {
238              values,
239              closer_nodes,
240          })
241      }
242  }
243  
244  #[cfg(test)]
245  mod tests {
246      use super::*;
247  
248      fn our_node() -> NodeInfo {
249          let pubkey = [0x42u8; 32];
250          NodeInfo::new(pubkey, "test:1234".to_string(), 1000)
251      }
252  
253      #[test]
254      fn test_message_builder() {
255          let builder = MessageBuilder::new(our_node());
256  
257          let ping = builder.ping();
258          assert!(matches!(ping, DhtMessage::Request(ref req) if matches!(**req, DhtRequest::Ping { .. })));
259  
260          let find = builder.find_node([0xABu8; 32]);
261          assert!(matches!(find, DhtMessage::Request(ref req) if matches!(**req, DhtRequest::FindNode { .. })));
262      }
263  
264      #[test]
265      fn test_sender_extraction() {
266          let node = our_node();
267          let req = DhtRequest::Ping { sender: node.clone() };
268          assert_eq!(req.sender().id, node.id);
269      }
270  }