v2
This commit is contained in:
lbl
@@ -0,0 +1,98 @@
|
||||
mod proto {
|
||||
include!(concat!(env!("OUT_DIR"), "/protocol.client.rs"));
|
||||
}
|
||||
|
||||
use anyhow::bail;
|
||||
use bytes::BytesMut;
|
||||
use prost::Message;
|
||||
use std::net::{Ipv4Addr, Ipv6Addr};
|
||||
|
||||
use crate::protocol::ProtoToBytesMut;
|
||||
pub use proto::*;
|
||||
|
||||
pub fn encode_nat_info(nat_info: &rust_p2p_core::nat::NatInfo) -> proto::NatInfo {
|
||||
let nat_type = match nat_info.nat_type {
|
||||
rust_p2p_core::nat::NatType::Cone => proto::NatType::Cone,
|
||||
rust_p2p_core::nat::NatType::Symmetric => proto::NatType::Symmetric,
|
||||
};
|
||||
|
||||
proto::NatInfo {
|
||||
nat_type: nat_type.into(),
|
||||
public_ips: nat_info.public_ips.iter().map(|v| (*v).into()).collect(),
|
||||
public_udp_ports: nat_info
|
||||
.public_udp_ports
|
||||
.iter()
|
||||
.map(|v| (*v).into())
|
||||
.collect(),
|
||||
public_port_range: nat_info.public_port_range.into(),
|
||||
local_ipv4s: nat_info.local_ipv4s.iter().map(|v| (*v).into()).collect(),
|
||||
ipv6: nat_info.ipv6.map(|v| v.octets().to_vec()),
|
||||
local_udp_ports: nat_info
|
||||
.local_udp_ports
|
||||
.iter()
|
||||
.map(|v| (*v).into())
|
||||
.collect(),
|
||||
local_tcp_port: nat_info.local_tcp_port.into(),
|
||||
public_tcp_port: nat_info.public_tcp_port.into(),
|
||||
}
|
||||
}
|
||||
pub fn decode_nat_info(msg: proto::NatInfo) -> anyhow::Result<rust_p2p_core::nat::NatInfo> {
|
||||
let nat_type = match msg.nat_type() {
|
||||
proto::NatType::Cone => rust_p2p_core::nat::NatType::Cone,
|
||||
proto::NatType::Symmetric => rust_p2p_core::nat::NatType::Symmetric,
|
||||
};
|
||||
let ipv6: Option<[u8; 16]> = msg.ipv6.and_then(|v| v.as_slice().try_into().ok());
|
||||
|
||||
// Validate all ports fit in u16
|
||||
let validate_port = |p: u32| -> anyhow::Result<u16> {
|
||||
u16::try_from(p).map_err(|_| anyhow::anyhow!("invalid port number: {}", p))
|
||||
};
|
||||
|
||||
let public_udp_ports: Result<Vec<_>, _> = msg
|
||||
.public_udp_ports
|
||||
.into_iter()
|
||||
.map(validate_port)
|
||||
.collect();
|
||||
let local_udp_ports: Result<Vec<_>, _> =
|
||||
msg.local_udp_ports.into_iter().map(validate_port).collect();
|
||||
|
||||
Ok(rust_p2p_core::nat::NatInfo {
|
||||
nat_type,
|
||||
public_ips: msg.public_ips.into_iter().map(|v| v.into()).collect(),
|
||||
public_udp_ports: public_udp_ports?,
|
||||
mapping_tcp_addr: vec![],
|
||||
mapping_udp_addr: vec![],
|
||||
public_port_range: validate_port(msg.public_port_range)?,
|
||||
local_ipv4: msg
|
||||
.local_ipv4s
|
||||
.first()
|
||||
.map(|v| (*v).into())
|
||||
.unwrap_or(Ipv4Addr::UNSPECIFIED),
|
||||
local_ipv4s: msg.local_ipv4s.into_iter().map(|v| v.into()).collect(),
|
||||
ipv6: ipv6.map(Ipv6Addr::from),
|
||||
local_udp_ports: local_udp_ports?,
|
||||
local_tcp_port: validate_port(msg.local_tcp_port)?,
|
||||
public_tcp_port: validate_port(msg.public_tcp_port)?,
|
||||
})
|
||||
}
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PunchInfo {
|
||||
pub nat_info: rust_p2p_core::nat::NatInfo,
|
||||
}
|
||||
|
||||
impl PunchInfo {
|
||||
pub fn from_slice(buf: &[u8]) -> anyhow::Result<Self> {
|
||||
let msg = proto::PunchInfo::decode(buf)?;
|
||||
let Some(nat_info) = msg.nat_info else {
|
||||
bail!("Punched info decode failed.");
|
||||
};
|
||||
let nat_info = decode_nat_info(nat_info)?;
|
||||
Ok(Self { nat_info })
|
||||
}
|
||||
pub fn encode(&self) -> BytesMut {
|
||||
let message = proto::PunchInfo {
|
||||
nat_info: Some(encode_nat_info(&self.nat_info)),
|
||||
};
|
||||
message.encode_bytes_mut()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
use crate::protocol::ProtoToBytesMut;
|
||||
pub(crate) use crate::protocol::control_message::proto::SelectiveBroadcast;
|
||||
use crate::protocol::control_message::proto::request_message::RequestPayload;
|
||||
use crate::protocol::control_message::proto::response_message::ResponsePayload;
|
||||
use anyhow::bail;
|
||||
use bytes::BytesMut;
|
||||
use prost::Message;
|
||||
use std::net::Ipv4Addr;
|
||||
|
||||
mod proto {
|
||||
include!(concat!(env!("OUT_DIR"), "/protocol.control_message.rs"));
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq, Default)]
|
||||
pub enum RegistrationMode {
|
||||
#[default]
|
||||
Normal = 0,
|
||||
PreRegister = 1,
|
||||
}
|
||||
|
||||
impl From<RegistrationMode> for proto::RegistrationMode {
|
||||
fn from(mode: RegistrationMode) -> Self {
|
||||
match mode {
|
||||
RegistrationMode::Normal => proto::RegistrationMode::Normal,
|
||||
RegistrationMode::PreRegister => proto::RegistrationMode::PreRegister,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<proto::RegistrationMode> for RegistrationMode {
|
||||
fn from(mode: proto::RegistrationMode) -> Self {
|
||||
match mode {
|
||||
proto::RegistrationMode::Normal => RegistrationMode::Normal,
|
||||
proto::RegistrationMode::PreRegister => RegistrationMode::PreRegister,
|
||||
}
|
||||
}
|
||||
}
|
||||
pub(crate) struct RegRequestMsg {
|
||||
pub network_code: String,
|
||||
pub device_id: String,
|
||||
pub ip: Option<Ipv4Addr>,
|
||||
pub name: String,
|
||||
pub version: String,
|
||||
pub key_sign: Option<String>,
|
||||
pub ip_variable: bool,
|
||||
pub server_id: u32,
|
||||
pub registration_mode: RegistrationMode,
|
||||
}
|
||||
impl RegRequestMsg {
|
||||
// pub fn check(&self) -> anyhow::Result<()> {
|
||||
// if self.network_code.is_empty() {
|
||||
// return Err(anyhow!("network_code cannot be empty"));
|
||||
// }
|
||||
// if self.network_code.len() > MAX_NETWORK_CODE_LEN {
|
||||
// return Err(anyhow!(
|
||||
// "network_code length exceeds {} characters (current: {})",
|
||||
// MAX_NETWORK_CODE_LEN,
|
||||
// self.network_code.len()
|
||||
// ));
|
||||
// }
|
||||
// if self.device_id.is_empty() {
|
||||
// return Err(anyhow!("device_id cannot be empty"));
|
||||
// }
|
||||
// if self.device_id.len() > MAX_DEVICE_ID_LEN {
|
||||
// return Err(anyhow!(
|
||||
// "device_id length exceeds {} characters (current: {})",
|
||||
// MAX_DEVICE_ID_LEN,
|
||||
// self.device_id.len()
|
||||
// ));
|
||||
// }
|
||||
//
|
||||
// if self.name.len() > MAX_NAME_LEN {
|
||||
// return Err(anyhow!(
|
||||
// "name length exceeds {} characters (current: {})",
|
||||
// MAX_NAME_LEN,
|
||||
// self.name.len()
|
||||
// ));
|
||||
// }
|
||||
//
|
||||
// if self.version.len() > MAX_VERSION_LEN {
|
||||
// return Err(anyhow!(
|
||||
// "version length exceeds {} characters (current: {})",
|
||||
// MAX_VERSION_LEN,
|
||||
// self.version.len()
|
||||
// ));
|
||||
// }
|
||||
//
|
||||
// Ok(())
|
||||
// }
|
||||
// pub fn from(msg: proto::RegRequestMsg) -> anyhow::Result<Self> {
|
||||
// Ok(Self {
|
||||
// network_code: msg.network_code,
|
||||
// device_id: msg.device_id,
|
||||
// ip: msg.ip.map(|ip| ip.into()),
|
||||
// name: msg.name,
|
||||
// version: msg.version,
|
||||
// key_sign: msg.key_sign,
|
||||
// ip_variable: msg.ip_variable,
|
||||
// server_id: msg.server_id,
|
||||
// })
|
||||
// }
|
||||
pub fn to(self) -> proto::RegRequestMsg {
|
||||
proto::RegRequestMsg {
|
||||
network_code: self.network_code,
|
||||
device_id: self.device_id,
|
||||
ip: self.ip.map(|ip| ip.into()),
|
||||
name: self.name,
|
||||
version: self.version,
|
||||
key_sign: self.key_sign,
|
||||
ip_variable: self.ip_variable,
|
||||
server_id: self.server_id,
|
||||
registration_mode: proto::RegistrationMode::from(self.registration_mode).into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct RegResponseMsg {
|
||||
pub ip: Ipv4Addr,
|
||||
pub prefix_len: u8,
|
||||
pub gateway: Ipv4Addr,
|
||||
pub server_version: String,
|
||||
}
|
||||
impl RegResponseMsg {
|
||||
pub fn from(msg: proto::RegResponseMsg) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
ip: msg.ip.into(),
|
||||
prefix_len: (msg.prefix_len & 0xFF) as u8,
|
||||
gateway: msg.gateway.into(),
|
||||
server_version: msg.server_version,
|
||||
})
|
||||
}
|
||||
pub fn to(self) -> proto::RegResponseMsg {
|
||||
proto::RegResponseMsg {
|
||||
ip: self.ip.into(),
|
||||
prefix_len: self.prefix_len as _,
|
||||
gateway: self.gateway.into(),
|
||||
server_version: self.server_version,
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct ErrorResponseMsg {
|
||||
pub code: u32,
|
||||
pub message: String,
|
||||
}
|
||||
impl ErrorResponseMsg {
|
||||
pub fn from(msg: proto::ErrorResponseMsg) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
code: msg.code,
|
||||
message: msg.message,
|
||||
})
|
||||
}
|
||||
pub fn to(self) -> proto::ErrorResponseMsg {
|
||||
proto::ErrorResponseMsg {
|
||||
code: self.code,
|
||||
message: self.message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub struct ConfirmRegResponseMsg {
|
||||
pub success: bool,
|
||||
}
|
||||
impl ConfirmRegResponseMsg {
|
||||
pub fn from(msg: proto::ConfirmRegResponseMsg) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
success: msg.success,
|
||||
})
|
||||
}
|
||||
pub fn to(self) -> proto::ConfirmRegResponseMsg {
|
||||
proto::ConfirmRegResponseMsg {
|
||||
success: self.success,
|
||||
}
|
||||
}
|
||||
}
|
||||
pub(crate) enum RequestMessage {
|
||||
Reg(RegRequestMsg),
|
||||
ConfirmReg,
|
||||
}
|
||||
impl RequestMessage {
|
||||
pub fn encode(self) -> BytesMut {
|
||||
let request_payload = match self {
|
||||
RequestMessage::Reg(reg) => RequestPayload::Reg(reg.to()),
|
||||
RequestMessage::ConfirmReg => RequestPayload::ConfirmReg(proto::ConfirmRegMsg {}),
|
||||
};
|
||||
proto::RequestMessage {
|
||||
request_payload: Some(request_payload),
|
||||
}
|
||||
.encode_bytes_mut()
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
pub enum ResponseMessage {
|
||||
Reg(RegResponseMsg),
|
||||
Error(ErrorResponseMsg),
|
||||
ConfirmReg(ConfirmRegResponseMsg),
|
||||
}
|
||||
impl ResponseMessage {
|
||||
pub fn from_slice(buf: &[u8]) -> anyhow::Result<Self> {
|
||||
let msg = proto::ResponseMessage::decode(buf)?;
|
||||
let Some(payload) = msg.response_payload else {
|
||||
bail!("unsupported")
|
||||
};
|
||||
match payload {
|
||||
ResponsePayload::Reg(reg) => Ok(ResponseMessage::Reg(RegResponseMsg::from(reg)?)),
|
||||
ResponsePayload::Error(e) => Ok(ResponseMessage::Error(ErrorResponseMsg::from(e)?)),
|
||||
ResponsePayload::ConfirmReg(c) => {
|
||||
Ok(ResponseMessage::ConfirmReg(ConfirmRegResponseMsg::from(c)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn encode(self) -> BytesMut {
|
||||
let response_payload = match self {
|
||||
ResponseMessage::Reg(reg) => ResponsePayload::Reg(reg.to()),
|
||||
ResponseMessage::Error(e) => ResponsePayload::Error(e.to()),
|
||||
ResponseMessage::ConfirmReg(c) => ResponsePayload::ConfirmReg(c.to()),
|
||||
};
|
||||
proto::ResponseMessage {
|
||||
response_payload: Some(response_payload),
|
||||
}
|
||||
.encode_bytes_mut()
|
||||
}
|
||||
}
|
||||
|
||||
impl SelectiveBroadcast {
|
||||
pub fn new(ips: &[Ipv4Addr], data: Vec<u8>) -> Self {
|
||||
SelectiveBroadcast {
|
||||
ips: ips.iter().map(|v| (*v).into()).collect(),
|
||||
data,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientSimpleInfo {
|
||||
pub ip: Ipv4Addr,
|
||||
pub online: bool,
|
||||
}
|
||||
impl ClientSimpleInfo {
|
||||
pub fn from(msg: proto::ClientSimpleInfo) -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
ip: msg.ip.into(),
|
||||
online: msg.online,
|
||||
})
|
||||
}
|
||||
pub fn to(self) -> proto::ClientSimpleInfo {
|
||||
proto::ClientSimpleInfo {
|
||||
ip: self.ip.into(),
|
||||
online: self.online,
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug)]
|
||||
pub struct ClientSimpleInfoList {
|
||||
pub data_version: u64,
|
||||
pub list: Vec<ClientSimpleInfo>,
|
||||
pub is_all: bool,
|
||||
pub time: i64,
|
||||
}
|
||||
impl ClientSimpleInfoList {
|
||||
pub fn from_slice(buf: &[u8]) -> anyhow::Result<Self> {
|
||||
let msg = proto::ClientSimpleInfoList::decode(buf)?;
|
||||
let mut list = Vec::with_capacity(msg.list.len());
|
||||
for x in msg.list {
|
||||
list.push(ClientSimpleInfo::from(x)?);
|
||||
}
|
||||
Ok(Self {
|
||||
data_version: msg.data_version,
|
||||
list,
|
||||
is_all: msg.is_all,
|
||||
time: msg.time,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
/*
|
||||
0 15 31
|
||||
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
| 1 | msg_type(7) |max ttl(4) |curr ttl(4)| C | G | R | reserve(13) |
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
| seq(32) |
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
| src ID(32) |
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
| dest ID(32) |
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
| payload(n) |
|
||||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||||
*/
|
||||
#![allow(dead_code)]
|
||||
use crate::protocol::transmission::TransmissionBytes;
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use std::io;
|
||||
use zerocopy::byteorder::{NetworkEndian, U32};
|
||||
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Ref, Unaligned};
|
||||
|
||||
#[derive(Debug, FromBytes, IntoBytes, Unaligned, KnownLayout, Immutable)]
|
||||
#[repr(C)]
|
||||
pub struct NetHeader {
|
||||
/// Byte 0: bit7 = 1, bit0..6 = msg_type
|
||||
pub type_byte: u8,
|
||||
/// Byte 1: high 4 = max ttl, low 4 = curr ttl
|
||||
pub ttl_byte: u8,
|
||||
/// Byte 2: C(0x80) | G(0x40) | reserve
|
||||
pub flags_byte: u8,
|
||||
/// Byte 3: reserve
|
||||
pub _reserved: u8,
|
||||
|
||||
pub seq: U32<NetworkEndian>,
|
||||
pub src_id: U32<NetworkEndian>,
|
||||
pub dest_id: U32<NetworkEndian>,
|
||||
}
|
||||
const COMPRESSED: u8 = 0x80;
|
||||
const GATEWAY: u8 = 0x40;
|
||||
const FEC: u8 = 0x20;
|
||||
impl NetHeader {
|
||||
#[inline]
|
||||
pub fn msg_type(&self) -> u8 {
|
||||
self.type_byte & 0x7F
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn set_msg_type(&mut self, msg_type: u8) {
|
||||
self.type_byte = (msg_type & 0x7F) | 0x80;
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn max_ttl(&self) -> u8 {
|
||||
self.ttl_byte >> 4
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn curr_ttl(&self) -> u8 {
|
||||
self.ttl_byte & 0x0F
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn set_ttl(&mut self, max: u8, curr: u8) {
|
||||
self.ttl_byte = (max << 4) | (curr & 0x0F);
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn decr_ttl(&mut self) {
|
||||
let curr = self.curr_ttl();
|
||||
if curr == 0 {
|
||||
return;
|
||||
}
|
||||
self.ttl_byte = (self.ttl_byte & 0xF0) | (curr - 1);
|
||||
}
|
||||
|
||||
fn set_flag(&mut self, mask: u8, val: bool) {
|
||||
if val {
|
||||
self.flags_byte |= mask;
|
||||
} else {
|
||||
self.flags_byte &= !mask;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
|
||||
pub enum MsgType {
|
||||
Turn = 1,
|
||||
Broadcast = 2,
|
||||
ExcludeBroadcast = 3,
|
||||
TargetBroadcast = 4,
|
||||
|
||||
Ping = 5,
|
||||
Pong = 6,
|
||||
PingTurn = 7,
|
||||
PongTurn = 8,
|
||||
|
||||
PunchStart1 = 9,
|
||||
PunchStart2 = 10,
|
||||
PunchReq = 11,
|
||||
PunchRes = 12,
|
||||
|
||||
PushClientIps = 13,
|
||||
|
||||
RpcReq = 14,
|
||||
RpcRes = 15,
|
||||
|
||||
Quic = 17,
|
||||
}
|
||||
impl From<MsgType> for u8 {
|
||||
fn from(val: MsgType) -> Self {
|
||||
val as u8
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<u8> for MsgType {
|
||||
type Error = io::Error;
|
||||
|
||||
fn try_from(value: u8) -> Result<Self, Self::Error> {
|
||||
let val = match value {
|
||||
1 => MsgType::Turn,
|
||||
2 => MsgType::Broadcast,
|
||||
3 => MsgType::ExcludeBroadcast,
|
||||
4 => MsgType::TargetBroadcast,
|
||||
|
||||
5 => MsgType::Ping,
|
||||
6 => MsgType::Pong,
|
||||
7 => MsgType::PingTurn,
|
||||
8 => MsgType::PongTurn,
|
||||
|
||||
9 => MsgType::PunchStart1,
|
||||
10 => MsgType::PunchStart2,
|
||||
11 => MsgType::PunchReq,
|
||||
12 => MsgType::PunchRes,
|
||||
|
||||
13 => MsgType::PushClientIps,
|
||||
|
||||
14 => MsgType::RpcReq,
|
||||
15 => MsgType::RpcRes,
|
||||
|
||||
17 => MsgType::Quic,
|
||||
_ => {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("invalid msg type:{value}"),
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(val)
|
||||
}
|
||||
}
|
||||
|
||||
pub const HEAD_LENGTH: usize = std::mem::size_of::<NetHeader>();
|
||||
|
||||
pub struct NetPacket<B> {
|
||||
buffer: B,
|
||||
}
|
||||
impl<B: AsRef<[u8]>> NetPacket<B> {
|
||||
pub fn new(buffer: B) -> io::Result<NetPacket<B>> {
|
||||
if buffer.as_ref().len() < HEAD_LENGTH {
|
||||
return Err(io::ErrorKind::InvalidInput.into());
|
||||
}
|
||||
Ok(NetPacket { buffer })
|
||||
}
|
||||
fn header(&self) -> Ref<&[u8], NetHeader> {
|
||||
// Safe: NetHeader is Unaligned and length is validated in new()
|
||||
let (header, _) = Ref::<&[u8], NetHeader>::from_prefix(self.buffer.as_ref()).unwrap();
|
||||
header
|
||||
}
|
||||
pub fn buffer(&self) -> &[u8] {
|
||||
self.buffer.as_ref()
|
||||
}
|
||||
pub fn into_buffer(self) -> B {
|
||||
self.buffer
|
||||
}
|
||||
pub fn source_buf(&self) -> &B {
|
||||
&self.buffer
|
||||
}
|
||||
pub fn msg_type(&self) -> io::Result<MsgType> {
|
||||
self.header().msg_type().try_into()
|
||||
}
|
||||
pub fn max_ttl(&self) -> u8 {
|
||||
self.header().max_ttl()
|
||||
}
|
||||
pub fn ttl(&self) -> u8 {
|
||||
self.header().curr_ttl()
|
||||
}
|
||||
|
||||
pub fn seq(&self) -> u32 {
|
||||
self.header().seq.get()
|
||||
}
|
||||
|
||||
pub fn src_id(&self) -> u32 {
|
||||
self.header().src_id.get()
|
||||
}
|
||||
|
||||
pub fn dest_id(&self) -> u32 {
|
||||
self.header().dest_id.get()
|
||||
}
|
||||
pub fn is_compressed(&self) -> bool {
|
||||
(self.header().flags_byte & COMPRESSED) != 0
|
||||
}
|
||||
pub fn is_gateway(&self) -> bool {
|
||||
(self.header().flags_byte & GATEWAY) != 0
|
||||
}
|
||||
pub fn is_fec(&self) -> bool {
|
||||
(self.header().flags_byte & FEC) != 0
|
||||
}
|
||||
pub fn head(&self) -> &[u8] {
|
||||
&self.buffer.as_ref()[..HEAD_LENGTH]
|
||||
}
|
||||
pub fn payload(&self) -> &[u8] {
|
||||
&self.buffer.as_ref()[HEAD_LENGTH..]
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: AsRef<[u8]> + AsMut<[u8]>> NetPacket<B> {
|
||||
fn header_mut(&mut self) -> Ref<&mut [u8], NetHeader> {
|
||||
// Safe: NetHeader is Unaligned and length is validated in new()
|
||||
let (header, _) = Ref::<&mut [u8], NetHeader>::from_prefix(self.buffer.as_mut()).unwrap();
|
||||
header
|
||||
}
|
||||
|
||||
pub fn set_msg_type(&mut self, msg_type: MsgType) {
|
||||
self.header_mut().set_msg_type(msg_type.into());
|
||||
}
|
||||
|
||||
pub fn decr_ttl(&mut self){
|
||||
self.header_mut().decr_ttl()
|
||||
}
|
||||
|
||||
pub fn set_ttl(&mut self, ttl: u8) {
|
||||
self.header_mut().set_ttl(ttl, ttl);
|
||||
}
|
||||
|
||||
pub fn set_seq(&mut self, seq: u32) {
|
||||
self.header_mut().seq.set(seq);
|
||||
}
|
||||
|
||||
pub fn set_src_id(&mut self, id: u32) {
|
||||
self.header_mut().src_id.set(id);
|
||||
}
|
||||
|
||||
pub fn set_dest_id(&mut self, id: u32) {
|
||||
self.header_mut().dest_id.set(id);
|
||||
}
|
||||
|
||||
pub fn set_compressed_flag(&mut self, compressed: bool) {
|
||||
self.header_mut().set_flag(COMPRESSED, compressed);
|
||||
}
|
||||
pub fn set_gateway_flag(&mut self, gateway: bool) {
|
||||
self.header_mut().set_flag(GATEWAY, gateway);
|
||||
}
|
||||
pub fn set_fec_flag(&mut self, fec: bool) {
|
||||
self.header_mut().set_flag(FEC, fec);
|
||||
}
|
||||
|
||||
pub fn set_payload(&mut self, data: &[u8]) -> io::Result<()> {
|
||||
let buf = self.buffer.as_mut();
|
||||
if buf.len() < HEAD_LENGTH + data.len() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"Invalid message length",
|
||||
));
|
||||
}
|
||||
buf[HEAD_LENGTH..HEAD_LENGTH + data.len()].copy_from_slice(data);
|
||||
Ok(())
|
||||
}
|
||||
pub fn head_mut(&mut self) -> &mut [u8] {
|
||||
&mut self.buffer.as_mut()[..HEAD_LENGTH]
|
||||
}
|
||||
pub fn payload_mut(&mut self) -> &mut [u8] {
|
||||
&mut self.buffer.as_mut()[HEAD_LENGTH..]
|
||||
}
|
||||
pub fn source_buf_mut(&mut self) -> &mut B {
|
||||
&mut self.buffer
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for NetPacket<Bytes> {
|
||||
fn clone(&self) -> Self {
|
||||
NetPacket {
|
||||
buffer: self.buffer.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NetPacket<BytesMut> {
|
||||
pub fn into_bytes(self) -> NetPacket<Bytes> {
|
||||
NetPacket {
|
||||
buffer: self.buffer.freeze(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NetPacket<TransmissionBytes> {
|
||||
pub fn into_bytes(self) -> NetPacket<Bytes> {
|
||||
NetPacket {
|
||||
buffer: self.buffer.into_bytes().freeze(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
use bytes::BytesMut;
|
||||
use prost::Message;
|
||||
|
||||
pub(crate) mod client_message;
|
||||
pub mod control_message;
|
||||
pub(crate) mod ip_packet_protocol;
|
||||
pub(crate) mod rpc_message;
|
||||
pub(crate) mod transmission;
|
||||
|
||||
pub trait ProtoToBytesMut: Message {
|
||||
fn encode_bytes_mut(&self) -> BytesMut
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
let mut bytes_mut = BytesMut::with_capacity(self.encoded_len());
|
||||
self.encode_raw(&mut bytes_mut);
|
||||
bytes_mut
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Message> ProtoToBytesMut for T {}
|
||||
@@ -0,0 +1,4 @@
|
||||
mod proto {
|
||||
include!(concat!(env!("OUT_DIR"), "/protocol.rpc.rs"));
|
||||
}
|
||||
pub use proto::*;
|
||||
@@ -0,0 +1,266 @@
|
||||
use bytes::{Buf, Bytes, BytesMut};
|
||||
use std::borrow::{Borrow, BorrowMut};
|
||||
use std::io;
|
||||
use std::ops::{Deref, DerefMut};
|
||||
|
||||
const DEFAULT_BUF_SIZE: usize = 2048;
|
||||
#[derive(Clone)]
|
||||
pub struct TransmissionBytes {
|
||||
buf: BytesMut,
|
||||
start: usize,
|
||||
end: usize,
|
||||
}
|
||||
impl From<BytesMut> for TransmissionBytes {
|
||||
fn from(buf: BytesMut) -> TransmissionBytes {
|
||||
let end = buf.len();
|
||||
Self { buf, start: 0, end }
|
||||
}
|
||||
}
|
||||
impl From<Bytes> for TransmissionBytes {
|
||||
fn from(buf: Bytes) -> TransmissionBytes {
|
||||
let end = buf.len();
|
||||
Self {
|
||||
buf: BytesMut::from(buf),
|
||||
start: 0,
|
||||
end,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl From<&[u8]> for TransmissionBytes {
|
||||
fn from(buf: &[u8]) -> TransmissionBytes {
|
||||
let end = buf.len();
|
||||
Self {
|
||||
buf: BytesMut::from(buf),
|
||||
start: 0,
|
||||
end,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TransmissionBytes {
|
||||
pub fn new_offset(start: usize) -> Self {
|
||||
TransmissionBytes {
|
||||
buf: BytesMut::zeroed(DEFAULT_BUF_SIZE),
|
||||
start,
|
||||
end: start,
|
||||
}
|
||||
}
|
||||
pub fn new_offset_zeroed(start: usize) -> Self {
|
||||
TransmissionBytes {
|
||||
buf: BytesMut::zeroed(DEFAULT_BUF_SIZE),
|
||||
start,
|
||||
end: DEFAULT_BUF_SIZE,
|
||||
}
|
||||
}
|
||||
pub fn zeroed(cap: usize) -> Self {
|
||||
TransmissionBytes {
|
||||
buf: BytesMut::zeroed(cap),
|
||||
start: 0,
|
||||
end: cap,
|
||||
}
|
||||
}
|
||||
pub fn zeroed_size(size: usize, reserve: usize) -> Self {
|
||||
TransmissionBytes {
|
||||
buf: BytesMut::zeroed(size + reserve),
|
||||
start: 0,
|
||||
end: size,
|
||||
}
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn with_capacity(head_room: usize, capacity: usize) -> Self {
|
||||
TransmissionBytes {
|
||||
buf: BytesMut::zeroed(capacity),
|
||||
start: head_room,
|
||||
end: head_room,
|
||||
}
|
||||
}
|
||||
pub fn len(&self) -> usize {
|
||||
self.end - self.start
|
||||
}
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn capacity(&self) -> usize {
|
||||
self.buf.capacity()
|
||||
}
|
||||
/// 头部可用空间（可向前扩展的字节数）
|
||||
#[inline]
|
||||
pub fn head_room(&self) -> usize {
|
||||
self.start
|
||||
}
|
||||
|
||||
/// 尾部可用空间（可向后扩展的字节数）
|
||||
#[inline]
|
||||
#[allow(dead_code)]
|
||||
pub fn tail_room(&self) -> usize {
|
||||
self.buf.capacity() - self.end
|
||||
}
|
||||
#[inline]
|
||||
fn as_slice(&self) -> &[u8] {
|
||||
&self.buf[self.start..self.end]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn as_slice_mut(&mut self) -> &mut [u8] {
|
||||
&mut self.buf[self.start..self.end]
|
||||
}
|
||||
pub fn put(&mut self, data: &[u8]) -> io::Result<()> {
|
||||
let need = data.len();
|
||||
let free = self.buf.capacity() - self.end;
|
||||
|
||||
if need > free {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("data too large:need={need},free={free}"),
|
||||
));
|
||||
}
|
||||
|
||||
self.buf[self.end..self.end + need].copy_from_slice(data);
|
||||
self.end += need;
|
||||
Ok(())
|
||||
}
|
||||
pub fn retreat_head(&mut self, len: usize) -> io::Result<()> {
|
||||
if len > self.head_room() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"retreat_head beyond start: len={len}, head_room={}",
|
||||
self.head_room()
|
||||
),
|
||||
));
|
||||
}
|
||||
self.start -= len;
|
||||
Ok(())
|
||||
}
|
||||
pub fn advance_head(&mut self, len: usize) -> io::Result<()> {
|
||||
if len > self.len() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"advance_head beyond end: len={len}, data_len={}",
|
||||
self.len()
|
||||
),
|
||||
));
|
||||
}
|
||||
self.start += len;
|
||||
Ok(())
|
||||
}
|
||||
pub fn set_len(&mut self, new_len: usize) -> io::Result<()> {
|
||||
let new_end = self.start + new_len;
|
||||
if new_end > self.buf.capacity() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"set_len exceeds capacity: new_len={new_len}, max={}",
|
||||
self.buf.capacity() - self.start
|
||||
),
|
||||
));
|
||||
}
|
||||
self.end = new_end;
|
||||
Ok(())
|
||||
}
|
||||
pub fn resize(&mut self, new_len: usize, value: u8) {
|
||||
let new_end = self.start + new_len;
|
||||
self.buf.resize(new_end, value);
|
||||
self.end = new_end;
|
||||
}
|
||||
pub fn extend_end(&mut self, n: usize) {
|
||||
if self.end + n > self.buf.len() {
|
||||
self.buf.resize(self.end + n, 0);
|
||||
}
|
||||
self.end += n;
|
||||
}
|
||||
pub fn shrink_end(&mut self, n: usize) {
|
||||
if n >= self.end - self.start {
|
||||
self.end = self.start;
|
||||
} else {
|
||||
self.end -= n;
|
||||
}
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn clear(&mut self) {
|
||||
self.start = 0;
|
||||
self.end = 0;
|
||||
}
|
||||
pub fn into_bytes(mut self) -> BytesMut {
|
||||
self.buf.truncate(self.end);
|
||||
if self.start > 0 {
|
||||
self.buf.advance(self.start);
|
||||
}
|
||||
self.buf
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<[u8]> for TransmissionBytes {
|
||||
#[inline]
|
||||
fn as_ref(&self) -> &[u8] {
|
||||
self.as_slice()
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for TransmissionBytes {
|
||||
type Target = [u8];
|
||||
|
||||
#[inline]
|
||||
fn deref(&self) -> &[u8] {
|
||||
self.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl AsMut<[u8]> for TransmissionBytes {
|
||||
#[inline]
|
||||
fn as_mut(&mut self) -> &mut [u8] {
|
||||
self.as_slice_mut()
|
||||
}
|
||||
}
|
||||
|
||||
impl DerefMut for TransmissionBytes {
|
||||
#[inline]
|
||||
fn deref_mut(&mut self) -> &mut [u8] {
|
||||
self.as_mut()
|
||||
}
|
||||
}
|
||||
|
||||
impl Borrow<[u8]> for TransmissionBytes {
|
||||
fn borrow(&self) -> &[u8] {
|
||||
self.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl BorrowMut<[u8]> for TransmissionBytes {
|
||||
fn borrow_mut(&mut self) -> &mut [u8] {
|
||||
self.as_mut()
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ShrinkEnd {
|
||||
fn shrink_end(&mut self, n: usize);
|
||||
}
|
||||
pub trait ExtendEnd {
|
||||
fn extend_end(&mut self, n: usize);
|
||||
}
|
||||
|
||||
impl ShrinkEnd for TransmissionBytes {
|
||||
fn shrink_end(&mut self, n: usize) {
|
||||
self.shrink_end(n);
|
||||
}
|
||||
}
|
||||
|
||||
impl ExtendEnd for TransmissionBytes {
|
||||
fn extend_end(&mut self, n: usize) {
|
||||
self.extend_end(n);
|
||||
}
|
||||
}
|
||||
|
||||
impl ShrinkEnd for &mut TransmissionBytes {
|
||||
fn shrink_end(&mut self, n: usize) {
|
||||
TransmissionBytes::shrink_end(self, n);
|
||||
}
|
||||
}
|
||||
|
||||
impl ExtendEnd for &mut TransmissionBytes {
|
||||
fn extend_end(&mut self, n: usize) {
|
||||
TransmissionBytes::extend_end(self, n);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user