1  // Copyright (c) 2025-2026 ACDC Network
  2  // This file is part of the alphaos library.
  3  //
  4  // Alpha Chain | Delta Chain Protocol
  5  // International Monetary Graphite.
  6  //
  7  // Derived from Aleo (https://aleo.org) and ProvableHQ (https://provable.com).
  8  // They built world-class ZK infrastructure. We installed the EASY button.
  9  // Their cryptography: elegant. Our modifications: bureaucracy-compatible.
 10  // Original brilliance: theirs. Robert's Rules: ours. Bugs: definitely ours.
 11  //
 12  // Original Aleo/ProvableHQ code subject to Apache 2.0 https://www.apache.org/licenses/LICENSE-2.0
 13  // All modifications and new work: CC0 1.0 Universal Public Domain Dedication.
 14  // No rights reserved. No permission required. No warranty. No refunds.
 15  //
 16  // https://creativecommons.org/publicdomain/zero/1.0/
 17  // SPDX-License-Identifier: CC0-1.0
 18  
 19  use crate::{bft::events::Event, bootstrap_client::network::MessageOrEvent, router::messages::Message};
 20  use alphavm::prelude::{FromBytes, Network, ToBytes};
 21  
 22  use bytes::{BufMut, BytesMut};
 23  use core::marker::PhantomData;
 24  use tokio_util::codec::{Decoder, Encoder, LengthDelimitedCodec};
 25  
 26  /// The maximum size of a message that can be transmitted during the handshake.
 27  const MAX_HANDSHAKE_SIZE: usize = 1024 * 1024; // 1 MiB
 28  /// The maximum size of a post-handshake message that can be obtained from the network.
 29  const MAX_POST_HANDSHAKE_SIZE: usize = 2 * 1024 * 1024; // 2 MiB
 30  
 31  /// The codec used to decode and encode network messages.
 32  pub struct BootstrapClientCodec<N: Network> {
 33      codec: LengthDelimitedCodec,
 34      _phantom: PhantomData<N>,
 35  }
 36  
 37  impl<N: Network> BootstrapClientCodec<N> {
 38      pub fn handshake() -> Self {
 39          let mut codec = Self::default();
 40          codec.codec.set_max_frame_length(MAX_HANDSHAKE_SIZE);
 41          codec
 42      }
 43  }
 44  
 45  impl<N: Network> Default for BootstrapClientCodec<N> {
 46      fn default() -> Self {
 47          Self {
 48              codec: LengthDelimitedCodec::builder()
 49                  .max_frame_length(MAX_POST_HANDSHAKE_SIZE)
 50                  .little_endian()
 51                  .new_codec(),
 52              _phantom: Default::default(),
 53          }
 54      }
 55  }
 56  
 57  impl<N: Network> Encoder<Message<N>> for BootstrapClientCodec<N> {
 58      type Error = std::io::Error;
 59  
 60      fn encode(&mut self, message: Message<N>, dst: &mut BytesMut) -> Result<(), Self::Error> {
 61          // Serialize the payload directly into dst.
 62          message
 63              .write_le(&mut dst.writer())
 64              // This error should never happen, the conversion is for greater compatibility.
 65              .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "serialization error"))?;
 66  
 67          let serialized_message = dst.split_to(dst.len()).freeze();
 68  
 69          self.codec.encode(serialized_message, dst)
 70      }
 71  }
 72  
 73  impl<N: Network> Encoder<Event<N>> for BootstrapClientCodec<N> {
 74      type Error = std::io::Error;
 75  
 76      fn encode(&mut self, event: Event<N>, dst: &mut BytesMut) -> Result<(), Self::Error> {
 77          // Serialize the payload directly into dst.
 78          event
 79              .write_le(&mut dst.writer())
 80              // This error should never happen, the conversion is for greater compatibility.
 81              .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "serialization error"))?;
 82  
 83          let serialized_event = dst.split_to(dst.len()).freeze();
 84  
 85          self.codec.encode(serialized_event, dst)
 86      }
 87  }
 88  
 89  impl<N: Network> Encoder<MessageOrEvent<N>> for BootstrapClientCodec<N> {
 90      type Error = std::io::Error;
 91  
 92      fn encode(&mut self, item: MessageOrEvent<N>, dst: &mut BytesMut) -> Result<(), Self::Error> {
 93          // Serialize the payload directly into dst.
 94          match item {
 95              MessageOrEvent::Message(message) => self.encode(message, dst),
 96              MessageOrEvent::Event(event) => self.encode(event, dst),
 97          }
 98      }
 99  }
100  
101  impl<N: Network> Decoder for BootstrapClientCodec<N> {
102      type Error = std::io::Error;
103      type Item = MessageOrEvent<N>;
104  
105      fn decode(&mut self, source: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
106          // Decode a frame containing bytes belonging to a message.
107          let bytes = match self.codec.decode(source)? {
108              Some(bytes) => bytes,
109              None => return Ok(None),
110          };
111  
112          // Reject invalid/truncated messages.
113          if bytes.len() < 2 {
114              warn!("Failed to deserialize a message: too short");
115              return Err(std::io::ErrorKind::InvalidData.into());
116          }
117  
118          // Check the ID of the serialized Message or Event.
119          let message_id = u16::from_le_bytes(bytes[..2].try_into().unwrap());
120  
121          // Discard messages that aren't of interest to a bootstrapper node.
122          match message_id {
123              2..=5 => match Message::read_le(&bytes[..]) {
124                  Ok(message) => Ok(Some(MessageOrEvent::Message(message))),
125                  Err(error) => {
126                      warn!("Failed to deserialize a message: {error}");
127                      Err(std::io::ErrorKind::InvalidData.into())
128                  }
129              },
130              7..=9 | 13 => match Event::read_le(&bytes[..]) {
131                  Ok(event) => Ok(Some(MessageOrEvent::Event(event))),
132                  Err(error) => {
133                      warn!("Failed to deserialize a message: {error}");
134                      Err(std::io::ErrorKind::InvalidData.into())
135                  }
136              },
137              id => {
138                  trace!("Ignoring an unhandled message (ID {id})");
139                  Ok(None)
140              }
141          }
142      }
143  }