From 58aeaf30d7f4fb1a355eb85bcbbd0b077bddd33a Mon Sep 17 00:00:00 2001 From: Florian Uekermann Date: Sat, 21 Feb 2026 17:54:29 +0100 Subject: [PATCH] add wireguard fragmentation helpers --- .gitignore | 1 + Cargo.lock | 1 + Cargo.toml | 1 + src/lib.rs | 1 + src/wg_fragment/merge.rs | 78 ++++++++++++++++++++++++++++++++++++++++ src/wg_fragment/mod.rs | 56 +++++++++++++++++++++++++++++ src/wg_fragment/split.rs | 42 ++++++++++++++++++++++ 7 files changed, 180 insertions(+) create mode 100644 src/wg_fragment/merge.rs create mode 100644 src/wg_fragment/mod.rs create mode 100644 src/wg_fragment/split.rs diff --git a/.gitignore b/.gitignore index 6d69d13..7312923 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /.idea +/.devcontainer /target result* diff --git a/Cargo.lock b/Cargo.lock index 0e4a2ca..d02d7df 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1202,6 +1202,7 @@ version = "0.0.0" dependencies = [ "anyhow", "base64", + "bytes", "clap", "derive_more", "http", diff --git a/Cargo.toml b/Cargo.toml index e7ba1e8..702eb7e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ repository = "https://github.com/Sovereign-Engineering/obscuravpn-api" [dependencies] anyhow = "1.0.98" base64 = "0.22.1" +bytes = "1.10.1" derive_more = { version = "2.0.1", features = ["full"] } http = { version = "1.3.1" } hyper-util = { version = "0.1.19", features = ["client-legacy"] } diff --git a/src/lib.rs b/src/lib.rs index d368080..2c7ecf0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -16,6 +16,7 @@ pub mod relay_protocol; mod resolver_fallback; #[cfg(feature = "client")] mod response; +pub mod wg_fragment; #[cfg(feature = "client")] pub use client::Client; diff --git a/src/wg_fragment/merge.rs b/src/wg_fragment/merge.rs new file mode 100644 index 0000000..eccdce8 --- /dev/null +++ b/src/wg_fragment/merge.rs @@ -0,0 +1,78 @@ +use super::{usize_from, WgFragmentHeaderData, WG_FRAGMENT_MESSAGE_HEADER_SIZE}; +use bytes::Bytes; +use std::num::NonZeroU32; + +pub enum ReassembleResult { + /// Message was not fragmented, returned as-is. + NotFragmented(Bytes), + /// Message was reassembled from two fragments. + Reassembled(Bytes), + /// A single fragment was received with no matching counterpart; the max message size that caused fragmentation. + UnmatchedFragment { max_message_size: u16 }, +} + +pub struct WgFragmentBuffer { + max_fragment_size: u16, + buffer: Vec>, +} + +impl WgFragmentBuffer { + pub fn new(len: NonZeroU32, max_fragment_size: u16) -> Self { + Self { + max_fragment_size, + buffer: vec![None; usize_from(len)], + } + } + + /// Process a potentially fragmented WG message. + /// + /// Returns the message if available, or the max message size from the fragment header if the fragment was buffered. + /// + /// Incorrect reassembly is theoretically possible if two fragmented messages are exactly `u32::MAX` or a multiple thereof apart, all intervening messages sharing the same buffer slot are lost, and only complementary halves of each survive. WireGuard's authenticated encryption will detect this downstream. + pub fn reassemble(&mut self, new_message: Bytes) -> ReassembleResult { + let Some(new_header) = WgFragmentHeaderData::from_message(&new_message) else { + return ReassembleResult::NotFragmented(new_message); + }; + let max_message_size = new_header.mtu; + if new_message.len() > usize_from(self.max_fragment_size) { + tracing::error!( + message_id = "zg4h9td3", + message_len = new_message.len(), + max_fragment_size = self.max_fragment_size, + "ignoring oversized WG message fragment" + ); + return ReassembleResult::UnmatchedFragment { max_message_size }; + } + let buffer_len = self.buffer.len(); + let buffer_element = &mut self.buffer[usize_from(new_header.message_idx) % buffer_len]; + match buffer_element { + None => { + *buffer_element = Some(new_message); + ReassembleResult::UnmatchedFragment { max_message_size } + } + Some(old_message) => { + let old_header = WgFragmentHeaderData::from_message(old_message).unwrap(); + if old_header.message_idx != new_header.message_idx { + *buffer_element = Some(new_message); + return ReassembleResult::UnmatchedFragment { max_message_size }; + } + if old_header.second_fragment == new_header.second_fragment { + *buffer_element = Some(new_message); + return ReassembleResult::UnmatchedFragment { max_message_size }; + } + let (first_msg, second_msg) = if old_header.second_fragment { + (new_message.as_ref(), old_message.as_ref()) + } else { + (old_message.as_ref(), new_message.as_ref()) + }; + let first_data = &first_msg[WG_FRAGMENT_MESSAGE_HEADER_SIZE..]; + let second_data = &second_msg[WG_FRAGMENT_MESSAGE_HEADER_SIZE..]; + let mut reassembled = Vec::with_capacity(first_data.len() + second_data.len()); + reassembled.extend_from_slice(first_data); + reassembled.extend_from_slice(second_data); + *buffer_element = None; + ReassembleResult::Reassembled(reassembled.into()) + } + } + } +} diff --git a/src/wg_fragment/mod.rs b/src/wg_fragment/mod.rs new file mode 100644 index 0000000..e10259e --- /dev/null +++ b/src/wg_fragment/mod.rs @@ -0,0 +1,56 @@ +use bytes::BufMut; + +pub mod merge; +pub mod split; + +// WireGuard uses message types 1-4. We picked 170, mid-range in the unassigned space, to avoid collisions with extensions that claim values near the boundaries. +const WG_FRAGMENT_MESSAGE_TYPE: u8 = 170; + +const WG_FRAGMENT_MESSAGE_HEADER_SIZE: usize = 1 // message type + + 1 // 1st (lowest) bit: fragment index; 2nd-8th bit: reserved + + 2 // max message size (limit that caused fragmentation, little-endian) + + 4; // message index (little-endian) + +pub struct WgFragmentHeaderData { + pub message_idx: u32, + pub second_fragment: bool, + pub mtu: u16, +} + +impl WgFragmentHeaderData { + pub fn header_bytes(&self) -> [u8; WG_FRAGMENT_MESSAGE_HEADER_SIZE] { + let mut header = [0u8; WG_FRAGMENT_MESSAGE_HEADER_SIZE]; + let mut buf = &mut header[..]; + buf.put_u8(WG_FRAGMENT_MESSAGE_TYPE); + buf.put_u8(u8::from(self.second_fragment)); + buf.put_u16_le(self.mtu); + buf.put_u32_le(self.message_idx); + header + } + pub fn from_message(message: &[u8]) -> Option { + let (header, _data) = message.split_at_checked(WG_FRAGMENT_MESSAGE_HEADER_SIZE)?; + let (message_type, header) = header.split_at(1); + if message_type != [WG_FRAGMENT_MESSAGE_TYPE] { + return None; + } + let (flags, header) = header.split_at(1); + let second_fragment = (flags[0] & 1) != 0; + let (mtu, header) = header.split_at(2); + let mtu = u16::from_le_bytes(mtu.try_into().unwrap()); + let (message_idx, header) = header.split_at(4); + let message_idx = u32::from_le_bytes(message_idx.try_into().unwrap()); + _ = header; + + Some(Self { + message_idx, + second_fragment, + mtu, + }) + } +} + +fn usize_from(x: impl Into) -> usize { + const _: () = assert!(usize::BITS >= 32, "usize smaller than u32"); + let x_u32: u32 = x.into(); + x_u32 as usize +} diff --git a/src/wg_fragment/split.rs b/src/wg_fragment/split.rs new file mode 100644 index 0000000..35d040c --- /dev/null +++ b/src/wg_fragment/split.rs @@ -0,0 +1,42 @@ +use super::{WgFragmentHeaderData, WG_FRAGMENT_MESSAGE_HEADER_SIZE}; +use bytes::Bytes; + +#[derive(Debug, Default)] +pub struct WgMessageFragmenter { + next_message_index: u32, +} + +impl WgMessageFragmenter { + pub fn fragment(&mut self, message: Bytes, max_message_size: u16) -> (Bytes, Option) { + if message.len() <= usize::from(max_message_size) { + return (message, None); + } + + let message_idx = self.next_message_index; + self.next_message_index = self.next_message_index.wrapping_add(1); + + let split_point = message.len() / 2; + + let first_header = WgFragmentHeaderData { + message_idx, + second_fragment: false, + mtu: max_message_size, + }; + let second_header = WgFragmentHeaderData { + message_idx, + second_fragment: true, + mtu: max_message_size, + }; + + let mut first_frag = Vec::with_capacity(WG_FRAGMENT_MESSAGE_HEADER_SIZE + split_point); + first_frag.extend_from_slice(&first_header.header_bytes()); + first_frag.extend_from_slice(&message[..split_point]); + + let remainder = message.len() - split_point; + let mut second_frag = Vec::with_capacity(WG_FRAGMENT_MESSAGE_HEADER_SIZE + remainder); + second_frag.extend_from_slice(&second_header.header_bytes()); + second_frag.extend_from_slice(&message[split_point..]); + + (first_frag.into(), Some(second_frag.into())) + } +}