feat: First version of broker based WireGuard PSK interface

This allows us to run with minimal priviledges in the Rosenpass process itself

wireguard-broker/src/api/client.rs

@@ -0,0 +1,152 @@
+use std::{borrow::BorrowMut, marker::PhantomData};
+
+use rosenpass_lenses::LenseView;
+
+use crate::{
+    api::msgs::{self, EnvelopeExt, SetPskRequestExt, SetPskResponseExt},
+    WireGuardBroker,
+};
+
+#[derive(thiserror::Error, Debug, Clone, Eq, PartialEq)]
+pub enum BrokerClientPollResponseError<RecvError> {
+    #[error(transparent)]
+    IoError(RecvError),
+    #[error("Invalid message.")]
+    InvalidMessage,
+}
+
+impl<RecvError> From<msgs::InvalidMessageTypeError> for BrokerClientPollResponseError<RecvError> {
+    fn from(value: msgs::InvalidMessageTypeError) -> Self {
+        let msgs::InvalidMessageTypeError = value; // Assert that this is a unit type
+        BrokerClientPollResponseError::<RecvError>::InvalidMessage
+    }
+}
+
+fn io_pollerr<RecvError>(e: RecvError) -> BrokerClientPollResponseError<RecvError> {
+    BrokerClientPollResponseError::<RecvError>::IoError(e)
+}
+
+fn invalid_msg_pollerr<RecvError>() -> BrokerClientPollResponseError<RecvError> {
+    BrokerClientPollResponseError::<RecvError>::InvalidMessage
+}
+
+#[derive(thiserror::Error, Debug, Clone, Eq, PartialEq)]
+pub enum BrokerClientSetPskError<SendError> {
+    #[error(transparent)]
+    IoError(SendError),
+    #[error("Interface name out of bounds")]
+    IfaceOutOfBounds,
+}
+
+pub trait BrokerClientIo {
+    type SendError;
+    type RecvError;
+
+    fn send_msg(&mut self, buf: &[u8]) -> Result<(), Self::SendError>;
+    fn recv_msg(&mut self) -> Result<Option<&[u8]>, Self::RecvError>;
+}
+
+#[derive(Debug)]
+pub struct BrokerClient<'a, Io, IoRef>
+where
+    Io: BrokerClientIo,
+    IoRef: 'a + BorrowMut<Io>,
+{
+    io: IoRef,
+    _phantom_io: PhantomData<&'a mut Io>,
+}
+
+impl<'a, Io, IoRef> BrokerClient<'a, Io, IoRef>
+where
+    Io: BrokerClientIo,
+    IoRef: 'a + BorrowMut<Io>,
+{
+    pub fn new(io: IoRef) -> Self {
+        Self {
+            io,
+            _phantom_io: PhantomData,
+        }
+    }
+
+    pub fn io(&self) -> &IoRef {
+        &self.io
+    }
+
+    pub fn io_mut(&mut self) -> &mut IoRef {
+        &mut self.io
+    }
+
+    pub fn poll_response(
+        &mut self,
+    ) -> Result<Option<msgs::SetPskResult>, BrokerClientPollResponseError<Io::RecvError>> {
+        let res: &[u8] = match self.io.borrow_mut().recv_msg().map_err(io_pollerr)? {
+            Some(r) => r,
+            None => return Ok(None),
+        };
+
+        let typ = res.get(0).ok_or(invalid_msg_pollerr())?;
+        let typ = msgs::MsgType::try_from(*typ)?;
+        let msgs::MsgType::SetPsk = typ; // Assert type
+
+        let res: msgs::Envelope<_, msgs::SetPskResponse<&[u8]>> = res
+            .envelope_truncating()
+            .map_err(|_| invalid_msg_pollerr())?;
+        let res: msgs::SetPskResponse<&[u8]> = res
+            .payload()
+            .set_psk_response()
+            .map_err(|_| invalid_msg_pollerr())?;
+        let res: msgs::SetPskResponseReturnCode = res.return_code()[0]
+            .try_into()
+            .map_err(|_| invalid_msg_pollerr())?;
+        let res: msgs::SetPskResult = res.into();
+
+        Ok(Some(res))
+    }
+}
+
+impl<'a, Io, IoRef> WireGuardBroker for BrokerClient<'a, Io, IoRef>
+where
+    Io: BrokerClientIo,
+    IoRef: 'a + BorrowMut<Io>,
+{
+    type Error = BrokerClientSetPskError<Io::SendError>;
+
+    fn set_psk(
+        &mut self,
+        iface: &str,
+        peer_id: [u8; 32],
+        psk: [u8; 32],
+    ) -> Result<(), Self::Error> {
+        use BrokerClientSetPskError::*;
+        const BUF_SIZE: usize = <msgs::Envelope<(), msgs::SetPskRequest<()>> as LenseView>::LEN;
+
+        // Allocate message
+        let mut req = [0u8; BUF_SIZE];
+
+        // Construct message view
+        let mut req: msgs::Envelope<_, msgs::SetPskRequest<&mut [u8]>> =
+            (&mut req as &mut [u8]).envelope_truncating().unwrap();
+
+        // Populate envelope
+        req.msg_type_mut()
+            .copy_from_slice(&[msgs::MsgType::SetPsk as u8]);
+        {
+            // Derived payload
+            let mut req: msgs::SetPskRequest<&mut [u8]> =
+                req.payload_mut().set_psk_request().unwrap();
+
+            // Populate payload
+            req.peer_id_mut().copy_from_slice(&peer_id);
+            req.psk_mut().copy_from_slice(&psk);
+            req.set_iface(iface).ok_or(IfaceOutOfBounds)?;
+        }
+
+        // Send message
+        self.io
+            .borrow_mut()
+            .send_msg(req.all_bytes())
+            .map_err(IoError)?;
+
+        Ok(())
+    }
+}

wireguard-broker/src/api/mio_client.rs

@@ -0,0 +1,204 @@
+use std::collections::VecDeque;
+use std::io::{ErrorKind, Read, Write};
+
+use anyhow::{bail, ensure};
+
+use crate::WireGuardBroker;
+
+use super::client::{
+    BrokerClient, BrokerClientIo, BrokerClientPollResponseError, BrokerClientSetPskError,
+};
+use super::msgs;
+
+#[derive(Debug)]
+pub struct MioBrokerClient {
+    inner: BrokerClient<'static, MioBrokerClientIo, MioBrokerClientIo>,
+}
+
+#[derive(Debug)]
+struct MioBrokerClientIo {
+    socket: mio::net::UnixStream,
+    send_buf: VecDeque<u8>,
+    receiving_size: bool,
+    recv_buf: Vec<u8>,
+    recv_off: usize,
+}
+
+impl MioBrokerClient {
+    pub fn new(socket: mio::net::UnixStream) -> Self {
+        let io = MioBrokerClientIo {
+            socket,
+            send_buf: VecDeque::new(),
+            receiving_size: false,
+            recv_buf: Vec::new(),
+            recv_off: 0,
+        };
+        let inner = BrokerClient::new(io);
+        Self { inner }
+    }
+
+    pub fn poll(&mut self) -> anyhow::Result<Option<msgs::SetPskResult>> {
+        self.inner.io_mut().flush()?;
+
+        // This sucks
+        match self.inner.poll_response() {
+            Ok(res) => {
+                return Ok(res);
+            }
+            Err(BrokerClientPollResponseError::IoError(e)) => {
+                return Err(e);
+            }
+            Err(BrokerClientPollResponseError::InvalidMessage) => {
+                bail!("Invalid message");
+            }
+        };
+    }
+}
+
+impl WireGuardBroker for MioBrokerClient {
+    type Error = anyhow::Error;
+
+    fn set_psk(&mut self, iface: &str, peer_id: [u8; 32], psk: [u8; 32]) -> anyhow::Result<()> {
+        use BrokerClientSetPskError::*;
+        let e = self.inner.set_psk(iface, peer_id, psk);
+        match e {
+            Ok(()) => Ok(()),
+            Err(IoError(e)) => Err(e),
+            Err(IfaceOutOfBounds) => bail!("Interface name size is out of bounds."),
+        }
+    }
+}
+
+impl BrokerClientIo for MioBrokerClientIo {
+    type SendError = anyhow::Error;
+    type RecvError = anyhow::Error;
+
+    fn send_msg(&mut self, buf: &[u8]) -> Result<(), Self::SendError> {
+        self.flush()?;
+        self.send_or_buffer(&(buf.len() as u64).to_le_bytes())?;
+        self.send_or_buffer(&buf)?;
+        self.flush()?;
+
+        Ok(())
+    }
+
+    fn recv_msg(&mut self) -> Result<Option<&[u8]>, Self::RecvError> {
+        // Stale message in receive buffer. Reset!
+        if self.recv_off == self.recv_buf.len() {
+            self.receiving_size = true;
+            self.recv_off = 0;
+            self.recv_buf.resize(8, 0);
+        }
+
+        // Try filling the receive buffer
+        self.recv_off += raw_recv(&self.socket, &mut self.recv_buf[self.recv_off..])?;
+        if self.recv_off < self.recv_buf.len() {
+            return Ok(None);
+        }
+
+        // Received size, now start receiving
+        if self.receiving_size {
+            // Received the size
+            // Parse the received length
+            let len: &[u8; 8] = self.recv_buf[..].try_into().unwrap();
+            let len: usize = u64::from_le_bytes(*len) as usize;
+
+            ensure!(
+                len <= msgs::RESPONSE_MSG_BUFFER_SIZE,
+                "Oversized buffer ({len}) in psk buffer response."
+            );
+
+            // Prepare the message buffer for receiving an actual message of the given size
+            self.receiving_size = false;
+            self.recv_off = 0;
+            self.recv_buf.resize(len, 0);
+
+            // Try to receive the message
+            return self.recv_msg();
+        }
+
+        // Received an actual message
+        return Ok(Some(&self.recv_buf[..]));
+    }
+}
+
+impl MioBrokerClientIo {
+    fn flush(&mut self) -> anyhow::Result<()> {
+        let (fst, snd) = self.send_buf.as_slices();
+
+        let (written, res) = match raw_send(&self.socket, fst) {
+            Ok(w1) if w1 >= fst.len() => match raw_send(&self.socket, snd) {
+                Ok(w2) => (w1 + w2, Ok(())),
+                Err(e) => (w1, Err(e)),
+            },
+            Ok(w1) => (w1, Ok(())),
+            Err(e) => (0, Err(e)),
+        };
+
+        self.send_buf.drain(..written);
+
+        (&self.socket).try_io(|| (&self.socket).flush())?;
+
+        res
+    }
+
+    fn send_or_buffer(&mut self, buf: &[u8]) -> anyhow::Result<()> {
+        let mut off = 0;
+
+        if self.send_buf.is_empty() {
+            off += raw_send(&self.socket, buf)?;
+        }
+
+        self.send_buf.extend((&buf[off..]).iter());
+
+        Ok(())
+    }
+}
+
+fn raw_send(mut socket: &mio::net::UnixStream, data: &[u8]) -> anyhow::Result<usize> {
+    let mut off = 0;
+
+    socket.try_io(|| {
+        loop {
+            if off == data.len() {
+                return Ok(());
+            }
+            match socket.write(&data[off..]) {
+                Ok(n) => {
+                    off += n;
+                }
+                Err(e) if e.kind() == ErrorKind::Interrupted => {
+                    // pass – retry
+                }
+                Err(e) if off > 0 || e.kind() == ErrorKind::WouldBlock => return Ok(()),
+                Err(e) => return Err(e),
+            }
+        }
+    })?;
+
+    return Ok(off);
+}
+
+fn raw_recv(mut socket: &mio::net::UnixStream, out: &mut [u8]) -> anyhow::Result<usize> {
+    let mut off = 0;
+
+    socket.try_io(|| {
+        loop {
+            if off == out.len() {
+                return Ok(());
+            }
+            match socket.read(&mut out[off..]) {
+                Ok(n) => {
+                    off += n;
+                }
+                Err(e) if e.kind() == ErrorKind::Interrupted => {
+                    // pass – retry
+                }
+                Err(e) if off > 0 || e.kind() == ErrorKind::WouldBlock => return Ok(()),
+                Err(e) => return Err(e),
+            }
+        }
+    })?;
+
+    return Ok(off);
+}

wireguard-broker/src/api/mod.rs

@@ -0,0 +1,4 @@
+pub mod client;
+pub mod mio_client;
+pub mod msgs;
+pub mod server;

wireguard-broker/src/api/msgs.rs

@@ -0,0 +1,140 @@
+use std::result::Result;
+use std::str::{from_utf8, Utf8Error};
+
+use rosenpass_lenses::{lense, LenseView};
+
+pub const REQUEST_MSG_BUFFER_SIZE: usize = <Envelope<(), SetPskRequest<()>> as LenseView>::LEN;
+pub const RESPONSE_MSG_BUFFER_SIZE: usize = <Envelope<(), SetPskResponse<()>> as LenseView>::LEN;
+
+lense! { Envelope<M> :=
+    /// [MsgType] of this message
+    msg_type: 1,
+    /// Reserved for future use
+    reserved: 3,
+    /// The actual Paylod
+    payload: M::LEN
+}
+
+lense! { SetPskRequest :=
+    peer_id: 32,
+    psk: 32,
+    iface_size: 1, // TODO: We should have variable length strings in lenses
+    iface_buf: 255
+}
+
+impl SetPskRequest<&[u8]> {
+    pub fn iface_bin(&self) -> &[u8] {
+        let len = self.iface_size()[0] as usize;
+        &self.iface_buf()[..len]
+    }
+
+    pub fn iface(&self) -> Result<&str, Utf8Error> {
+        from_utf8(self.iface_bin())
+    }
+}
+
+impl SetPskRequest<&mut [u8]> {
+    pub fn set_iface_bin(&mut self, iface: &[u8]) -> Option<()> {
+        (iface.len() < 256).then_some(())?; // Assert iface.len() < 256
+
+        self.iface_size_mut()[0] = iface.len() as u8;
+
+        self.iface_buf_mut().fill(0);
+        (&mut self.iface_buf_mut()[..iface.len()]).copy_from_slice(iface);
+
+        Some(())
+    }
+
+    pub fn set_iface(&mut self, iface: &str) -> Option<()> {
+        self.set_iface_bin(iface.as_bytes())
+    }
+}
+
+lense! { SetPskResponse :=
+    return_code: 1
+}
+
+#[derive(thiserror::Error, Debug, Clone, Eq, PartialEq)]
+pub enum SetPskError {
+    #[error("The wireguard pre-shared-key assignment broker experienced an internal error.")]
+    InternalError,
+    #[error("The indicated wireguard interface does not exist")]
+    NoSuchInterface,
+    #[error("The indicated peer does not exist on the wireguard interface")]
+    NoSuchPeer,
+}
+
+pub type SetPskResult = Result<(), SetPskError>;
+
+#[repr(u8)]
+#[derive(Hash, PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
+pub enum SetPskResponseReturnCode {
+    Success = 0x00,
+    InternalError = 0x01,
+    NoSuchInterface = 0x02,
+    NoSuchPeer = 0x03,
+}
+
+#[derive(Eq, PartialEq, Debug, Clone)]
+pub struct InvalidSetPskResponseError;
+
+impl TryFrom<u8> for SetPskResponseReturnCode {
+    type Error = InvalidSetPskResponseError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        use SetPskResponseReturnCode::*;
+        match value {
+            0x00 => Ok(Success),
+            0x01 => Ok(InternalError),
+            0x02 => Ok(NoSuchInterface),
+            0x03 => Ok(NoSuchPeer),
+            _ => Err(InvalidSetPskResponseError),
+        }
+    }
+}
+
+impl From<SetPskResponseReturnCode> for SetPskResult {
+    fn from(value: SetPskResponseReturnCode) -> Self {
+        use SetPskError as E;
+        use SetPskResponseReturnCode as C;
+        match value {
+            C::Success => Ok(()),
+            C::InternalError => Err(E::InternalError),
+            C::NoSuchInterface => Err(E::NoSuchInterface),
+            C::NoSuchPeer => Err(E::NoSuchPeer),
+        }
+    }
+}
+
+impl From<SetPskResult> for SetPskResponseReturnCode {
+    fn from(value: SetPskResult) -> Self {
+        use SetPskError as E;
+        use SetPskResponseReturnCode as C;
+        match value {
+            Ok(()) => C::Success,
+            Err(E::InternalError) => C::InternalError,
+            Err(E::NoSuchInterface) => C::NoSuchInterface,
+            Err(E::NoSuchPeer) => C::NoSuchPeer,
+        }
+    }
+}
+
+#[repr(u8)]
+#[derive(Hash, PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy)]
+pub enum MsgType {
+    SetPsk = 0x01,
+}
+
+#[derive(Eq, PartialEq, Debug, Clone)]
+pub struct InvalidMessageTypeError;
+
+impl TryFrom<u8> for MsgType {
+    type Error = InvalidMessageTypeError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0x01 => Ok(MsgType::SetPsk),
+            _ => Err(InvalidMessageTypeError),
+        }
+    }
+}

wireguard-broker/src/api/server.rs

@@ -0,0 +1,99 @@
+use std::borrow::BorrowMut;
+use std::marker::PhantomData;
+use std::result::Result;
+
+use rosenpass_lenses::LenseError;
+
+use crate::api::msgs::{self, EnvelopeExt, SetPskRequestExt, SetPskResponseExt};
+use crate::WireGuardBroker;
+
+#[derive(thiserror::Error, Debug, Clone, Eq, PartialEq)]
+pub enum BrokerServerError {
+    #[error("No such request type: {}", .0)]
+    NoSuchRequestType(u8),
+    #[error("Invalid message received.")]
+    InvalidMessage,
+}
+
+impl From<LenseError> for BrokerServerError {
+    fn from(value: LenseError) -> Self {
+        use BrokerServerError as Be;
+        use LenseError as Le;
+        match value {
+            Le::BufferSizeMismatch => Be::InvalidMessage,
+        }
+    }
+}
+
+impl From<msgs::InvalidMessageTypeError> for BrokerServerError {
+    fn from(value: msgs::InvalidMessageTypeError) -> Self {
+        let msgs::InvalidMessageTypeError = value; // Assert that this is a unit type
+        BrokerServerError::InvalidMessage
+    }
+}
+
+pub struct BrokerServer<'a, Err, Inner, Ref>
+where
+    msgs::SetPskError: From<Err>,
+    Inner: WireGuardBroker<Error = Err>,
+    Ref: BorrowMut<Inner> + 'a,
+{
+    inner: Ref,
+    _phantom: PhantomData<&'a mut Inner>,
+}
+
+impl<'a, Err, Inner, Ref> BrokerServer<'a, Err, Inner, Ref>
+where
+    msgs::SetPskError: From<Err>,
+    Inner: WireGuardBroker<Error = Err>,
+    Ref: 'a + BorrowMut<Inner>,
+{
+    pub fn new(inner: Ref) -> Self {
+        Self {
+            inner,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub fn handle_message(
+        &mut self,
+        req: &[u8],
+        res: &mut [u8; msgs::RESPONSE_MSG_BUFFER_SIZE],
+    ) -> Result<usize, BrokerServerError> {
+        use BrokerServerError::*;
+
+        let typ = req.get(0).ok_or(InvalidMessage)?;
+        let typ = msgs::MsgType::try_from(*typ)?;
+        let msgs::MsgType::SetPsk = typ; // Assert type
+
+        let req: msgs::Envelope<_, msgs::SetPskRequest<&[u8]>> = req.envelope_truncating()?;
+        let mut res: msgs::Envelope<_, msgs::SetPskResponse<&mut [u8]>> =
+            (res as &mut [u8]).envelope_truncating()?;
+        (&mut res).msg_type_mut()[0] = msgs::MsgType::SetPsk as u8;
+        self.handle_set_psk(
+            req.payload().set_psk_request()?,
+            res.payload_mut().set_psk_response()?,
+        )?;
+        Ok(res.all_bytes().len())
+    }
+
+    fn handle_set_psk(
+        &mut self,
+        req: msgs::SetPskRequest<&[u8]>,
+        mut res: msgs::SetPskResponse<&mut [u8]>,
+    ) -> Result<(), BrokerServerError> {
+        // Using unwrap here since lenses can not return fixed-size arrays
+        // TODO: Slices should give access to fixed size arrays
+        let r: Result<(), Err> = self.inner.borrow_mut().set_psk(
+            req.iface()
+                .map_err(|_e| BrokerServerError::InvalidMessage)?,
+            req.peer_id().try_into().unwrap(),
+            req.psk().try_into().unwrap(),
+        );
+        let r: msgs::SetPskResult = r.map_err(|e| e.into());
+        let r: msgs::SetPskResponseReturnCode = r.into();
+        res.return_code_mut()[0] = r as u8;
+
+        Ok(())
+    }
+}