add wireguard fragmentation helpers

.gitignore

@@ -1,3 +1,4 @@
 /.idea
+/.devcontainer
 /target
 result*

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"] }

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;

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<Option<Bytes>>,
+}
+
+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())
+            }
+        }
+    }
+}

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<Self> {
+        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<u32>) -> usize {
+    const _: () = assert!(usize::BITS >= 32, "usize smaller than u32");
+    let x_u32: u32 = x.into();
+    x_u32 as usize
+}

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<Bytes>) {
+        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()))
+    }
+}