refactor: remove panic-prone unwrap and expect calls

This commit is contained in:
lbl
2026-08-23 12:04:40 +08:00
parent 2100ec1ca5
commit c2fd1ff870
18 changed files with 161 additions and 128 deletions
+11 -12
View File
@@ -1,4 +1,4 @@
fn main() {
fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut config = prost_build::Config::new();
match protoc_bin_vendored::protoc_bin_path() {
@@ -14,15 +14,14 @@ fn main() {
config.protoc_arg("--experimental_allow_proto3_optional");
config
.compile_protos(
&[
"proto/control_message.proto",
"proto/rpc.proto",
"proto/client.proto",
"proto/fec.proto",
],
&["proto"],
)
.unwrap();
config.compile_protos(
&[
"proto/control_message.proto",
"proto/rpc.proto",
"proto/client.proto",
"proto/fec.proto",
],
&["proto"],
)?;
Ok(())
}
+2 -2
View File
@@ -79,7 +79,7 @@ impl NetworkManager {
log::info!("绑定出口网卡: {name}");
}
let mtu = config.mtu.unwrap_or(DEFAULT_MTU);
let packet_crypto = PacketCrypto::new_from_str(config.password.as_deref());
let packet_crypto = PacketCrypto::new_from_str(config.password.as_deref())?;
let packet_compression = PacketCompression::new(config.compress);
let (server_manager_list, tunnel_to_server, server_rpc) = create_server_tunnel(
app_state.clone(),
@@ -163,7 +163,7 @@ impl NetworkManager {
EnhancedTunInbound::Nat(
internal_nat_inbound
.clone()
.expect("internal_nat_inbound must be Some in no-device mode"),
.context("internal NAT is unavailable in no-device mode")?,
),
None,
),
+14 -13
View File
@@ -33,15 +33,16 @@ impl PacketCrypto {
.map(|b| format!("{:02x}", b))
.collect::<String>()
}
pub fn new(key_bytes: [u8; 32]) -> Self {
let unbound = UnboundKey::new(&CHACHA20_POLY1305, &key_bytes).unwrap();
pub fn new(key_bytes: [u8; 32]) -> io::Result<Self> {
let unbound = UnboundKey::new(&CHACHA20_POLY1305, &key_bytes)
.map_err(|_| io::Error::other("failed to initialize ChaCha20-Poly1305 key"))?;
let key = LessSafeKey::new(unbound);
Self {
Ok(Self {
key,
seq: Arc::new(AtomicU32::new(rand::random())),
}
})
}
pub fn new_from_str(s: &str) -> Self {
pub fn new_from_str(s: &str) -> io::Result<Self> {
let hash = ring::digest::digest(&ring::digest::SHA256, s.as_bytes());
let mut key_bytes = [0u8; 32];
key_bytes.copy_from_slice(hash.as_ref());
@@ -161,7 +162,7 @@ mod tests {
#[test]
fn test_encrypt_decrypt_in_place() {
let key = [7u8; 32];
let crypto = PacketCrypto::new(key);
let crypto = PacketCrypto::new(key).unwrap();
let payload_len = 20;
let mut pkt = build_test_packet(payload_len);
@@ -198,7 +199,7 @@ mod tests {
#[test]
fn test_nonce_unique_per_packet() {
let crypto = PacketCrypto::new([7u8; 32]);
let crypto = PacketCrypto::new([7u8; 32]).unwrap();
let mut pkt1 = build_test_packet(20);
let mut pkt2 = build_test_packet(20);
@@ -227,7 +228,7 @@ mod tests {
#[test]
fn test_clone_shares_seq_counter() {
let crypto = PacketCrypto::new([9u8; 32]);
let crypto = PacketCrypto::new([9u8; 32]).unwrap();
let cloned = crypto.clone();
let mut pkt1 = build_test_packet(8);
@@ -243,9 +244,9 @@ mod tests {
#[test]
fn test_cross_version_compat() {
let key = [7u8; 32];
let crypto = PacketCrypto::new(key);
let crypto = PacketCrypto::new(key).unwrap();
// 用相同密钥的另一个实例模拟对端
let peer = PacketCrypto::new(key);
let peer = PacketCrypto::new(key).unwrap();
// 模拟旧版本发包:seq 固定为 0,nonce 直接由头部计算
let mut pkt = build_test_packet(20);
@@ -281,7 +282,7 @@ mod tests {
/// 必须导致解密失败,而不是被静默接受。
#[test]
fn test_tampered_flags_rejected() {
let crypto = PacketCrypto::new([7u8; 32]);
let crypto = PacketCrypto::new([7u8; 32]).unwrap();
let mut pkt = build_test_packet(20);
crypto.encrypt_in_place(&mut pkt).expect("encrypt failed");
@@ -298,7 +299,7 @@ mod tests {
/// AAD 覆盖 msg_type(byte0):中间人篡改消息类型必须导致解密失败。
#[test]
fn test_tampered_msg_type_rejected() {
let crypto = PacketCrypto::new([7u8; 32]);
let crypto = PacketCrypto::new([7u8; 32]).unwrap();
let mut pkt = build_test_packet(20);
crypto.encrypt_in_place(&mut pkt).expect("encrypt failed");
@@ -315,7 +316,7 @@ mod tests {
/// 转发后 ttl 变化的包必须仍能正常解密。
#[test]
fn test_ttl_change_still_decrypts() {
let crypto = PacketCrypto::new([7u8; 32]);
let crypto = PacketCrypto::new([7u8; 32]).unwrap();
let payload_len = 20;
let mut pkt = build_test_packet(payload_len);
+6 -6
View File
@@ -16,12 +16,12 @@ impl PacketCrypto {
chacha20_poly1305::PacketCrypto::key_sign(s)
}
pub(crate) fn new_from_str(s: Option<&str>) -> Self {
Self {
crypto: s
.map(chacha20_poly1305::PacketCrypto::new_from_str)
.map(Arc::new),
}
pub(crate) fn new_from_str(s: Option<&str>) -> io::Result<Self> {
let crypto = s
.map(chacha20_poly1305::PacketCrypto::new_from_str)
.transpose()?
.map(Arc::new);
Ok(Self { crypto })
}
pub(crate) fn encrypt_reserve(&self) -> usize {
if self.crypto.is_some() { TAG_LEN } else { 0 }
@@ -145,7 +145,7 @@ async fn create_quic_endpoint(
quinn::crypto::rustls::QuicClientConfig::try_from(client_config)
.context("Failed to create QUIC client config")?,
));
client_config.transport_config(build_transport_config());
client_config.transport_config(build_transport_config()?);
let mut endpoint_config = EndpointConfig::default();
endpoint_config.max_udp_payload_size(1300)?;
let mut endpoint = quinn::Endpoint::new_with_abstract_socket(
@@ -159,14 +159,18 @@ async fn create_quic_endpoint(
Ok((inbound, endpoint))
}
fn build_transport_config() -> Arc<TransportConfig> {
fn build_transport_config() -> anyhow::Result<Arc<TransportConfig>> {
let mut transport = TransportConfig::default();
transport.congestion_controller_factory(Arc::new(BbrConfig::default()));
transport.keep_alive_interval(Some(Duration::from_secs(5)));
transport.max_idle_timeout(Some(Duration::from_secs(10).try_into().unwrap()));
transport.max_idle_timeout(Some(
Duration::from_secs(10)
.try_into()
.context("invalid QUIC idle timeout")?,
));
Arc::new(transport)
Ok(Arc::new(transport))
}
async fn ip_stack_recv_task(
+4 -4
View File
@@ -85,9 +85,9 @@ async fn handle_outbound(
} else {
// 创建真实 UDP socket
let bind_addr = if dst.is_ipv4() {
"0.0.0.0:0".parse().expect("valid IPv4 bind address")
SocketAddr::from(([0, 0, 0, 0], 0))
} else {
"[::]:0".parse().expect("valid IPv6 bind address")
SocketAddr::from(([0; 8], 0))
};
let interface = if dst.ip().is_loopback() {
None
@@ -223,9 +223,9 @@ where
.next()
.context("UDP NAT destination resolved to no address")?;
let bind_addr = if destination.is_ipv4() {
"0.0.0.0:0".parse().expect("valid IPv4 bind address")
SocketAddr::from(([0, 0, 0, 0], 0))
} else {
"[::]:0".parse().expect("valid IPv6 bind address")
SocketAddr::from(([0; 8], 0))
};
let interface = if destination.ip().is_loopback() {
None
+37 -32
View File
@@ -18,7 +18,7 @@ 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};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned};
#[derive(Debug, FromBytes, IntoBytes, Unaligned, KnownLayout, Immutable)]
#[repr(C)]
@@ -167,11 +167,6 @@ impl<B: AsRef<[u8]>> NetPacket<B> {
}
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()
}
@@ -182,37 +177,40 @@ impl<B: AsRef<[u8]>> NetPacket<B> {
&self.buffer
}
pub fn msg_type(&self) -> io::Result<MsgType> {
self.header().msg_type().try_into()
(self.buffer.as_ref()[0] & 0x7F).try_into()
}
pub fn max_ttl(&self) -> u8 {
self.header().max_ttl()
self.buffer.as_ref()[1] >> 4
}
pub fn ttl(&self) -> u8 {
self.header().curr_ttl()
self.buffer.as_ref()[1] & 0x0F
}
pub fn seq(&self) -> u32 {
self.header().seq.get()
let buf = self.buffer.as_ref();
u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]])
}
pub fn src_id(&self) -> u32 {
self.header().src_id.get()
let buf = self.buffer.as_ref();
u32::from_be_bytes([buf[8], buf[9], buf[10], buf[11]])
}
pub fn dest_id(&self) -> u32 {
self.header().dest_id.get()
let buf = self.buffer.as_ref();
u32::from_be_bytes([buf[12], buf[13], buf[14], buf[15]])
}
pub fn is_compressed(&self) -> bool {
(self.header().flags_byte & COMPRESSED) != 0
(self.buffer.as_ref()[2] & COMPRESSED) != 0
}
pub fn is_gateway(&self) -> bool {
(self.header().flags_byte & GATEWAY) != 0
(self.buffer.as_ref()[2] & GATEWAY) != 0
}
pub fn is_fec(&self) -> bool {
(self.header().flags_byte & FEC) != 0
(self.buffer.as_ref()[2] & FEC) != 0
}
pub fn is_ethernet(&self) -> bool {
(self.header().flags_byte & ETHERNET) != 0
(self.buffer.as_ref()[2] & ETHERNET) != 0
}
pub fn head(&self) -> &[u8] {
&self.buffer.as_ref()[..HEAD_LENGTH]
@@ -223,47 +221,54 @@ impl<B: AsRef<[u8]>> NetPacket<B> {
}
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());
self.buffer.as_mut()[0] = (u8::from(msg_type) & 0x7F) | 0x80;
}
pub fn decr_ttl(&mut self) {
self.header_mut().decr_ttl()
let ttl_byte = &mut self.buffer.as_mut()[1];
let current = *ttl_byte & 0x0F;
if current != 0 {
*ttl_byte = (*ttl_byte & 0xF0) | (current - 1);
}
}
pub fn set_ttl(&mut self, ttl: u8) {
self.header_mut().set_ttl(ttl, ttl);
self.buffer.as_mut()[1] = (ttl << 4) | (ttl & 0x0F);
}
pub fn set_seq(&mut self, seq: u32) {
self.header_mut().seq.set(seq);
self.buffer.as_mut()[4..8].copy_from_slice(&seq.to_be_bytes());
}
pub fn set_src_id(&mut self, id: u32) {
self.header_mut().src_id.set(id);
self.buffer.as_mut()[8..12].copy_from_slice(&id.to_be_bytes());
}
pub fn set_dest_id(&mut self, id: u32) {
self.header_mut().dest_id.set(id);
self.buffer.as_mut()[12..16].copy_from_slice(&id.to_be_bytes());
}
fn set_flag(&mut self, mask: u8, value: bool) {
let flags = &mut self.buffer.as_mut()[2];
if value {
*flags |= mask;
} else {
*flags &= !mask;
}
}
pub fn set_compressed_flag(&mut self, compressed: bool) {
self.header_mut().set_flag(COMPRESSED, compressed);
self.set_flag(COMPRESSED, compressed);
}
pub fn set_gateway_flag(&mut self, gateway: bool) {
self.header_mut().set_flag(GATEWAY, gateway);
self.set_flag(GATEWAY, gateway);
}
pub fn set_fec_flag(&mut self, fec: bool) {
self.header_mut().set_flag(FEC, fec);
self.set_flag(FEC, fec);
}
pub fn set_ethernet_flag(&mut self, ethernet: bool) {
self.header_mut().set_flag(ETHERNET, ethernet);
self.set_flag(ETHERNET, ethernet);
}
pub fn set_payload(&mut self, data: &[u8]) -> io::Result<()> {
+8 -3
View File
@@ -99,12 +99,17 @@ impl DeviceIOManager {
// 保证失败时调用方状态完整、可以重试
let device_mode = device_config.device_mode;
let device = Arc::new(create_device(device_config)?);
let receiver = receiver.take().unwrap();
let enhanced_outbound = enhanced_outbound.take().unwrap();
let Some(receiver_value) = receiver.take() else {
bail!("device task already started");
};
let Some(enhanced_outbound) = enhanced_outbound.take() else {
*receiver = Some(receiver_value);
bail!("device task already started");
};
let task = create(
&self.task_group,
device,
receiver.receiver,
receiver_value.receiver,
enhanced_outbound,
device_mode,
);
@@ -357,7 +357,9 @@ pub(crate) async fn query_tcp_public_addr_loop(
let addrs: Vec<SocketAddr> = active_connections.keys().cloned().collect();
for addr in addrs {
let (tcp_stream, _) = active_connections.get_mut(&addr).unwrap();
let Some((tcp_stream, _)) = active_connections.get_mut(&addr) else {
continue;
};
let mut buf = [0u8; 1024];
match tcp_stream.try_read(&mut buf) {
@@ -149,8 +149,11 @@ pub async fn ping_all(
ping.set_ttl(1);
ping.set_src_id(src.into());
ping.set_dest_id(id.into());
ping.set_payload(&crate::utils::time::now_ts_ms().to_be_bytes())
.unwrap();
if let Err(error) = ping.set_payload(&crate::utils::time::now_ts_ms().to_be_bytes())
{
log::warn!("failed to build route probe: {error}");
continue;
}
let route_key = route.route_key();
if socket_manager.send_to(ping, &route_key).await.is_ok() {
packet_loss_stats.record_sent(id, route_key);
+4 -1
View File
@@ -151,7 +151,10 @@ impl ServerOutbound {
// 只有一个服务器,直接发送
if map.len() == 1 {
let (server_id, (ips, _)) = map.iter().next().expect("map has exactly one element");
let (server_id, (ips, _)) = map
.iter()
.next()
.context("connected server map unexpectedly became empty")?;
if ips.is_empty() {
return Ok(());
}
+2 -4
View File
@@ -228,11 +228,9 @@ fn bind_udp(
default_interface: &Option<LocalInterface>,
) -> io::Result<UdpSocket> {
let addr: SocketAddr = if name_server.is_ipv4() {
"0.0.0.0:0"
.parse()
.expect("valid IPv4 socket address literal")
SocketAddr::from(([0, 0, 0, 0], 0))
} else {
"[::]:0".parse().expect("valid IPv6 socket address literal")
SocketAddr::from(([0; 8], 0))
};
let socket = rust_p2p_core::socket::bind_udp(addr, default_interface.as_ref())?;
UdpSocket::from_std(socket.into())