aps/src/protocol.rs (a74924e) - apple-crates · Git

protocol.rs

use crate::PushToken;
use bytes::{Bytes, BytesMut};
use std::io;
use std::ops::Range;
use std::time::Duration;
use thiserror::Error;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

pub(crate) const TIMEOUT: Duration = Duration::from_secs(30);

const MAX_FRAME_SIZE: usize = 1024 * 1024;
// command:u8, then length:u32.
const HEADER_SIZE: usize = 5;
// id:u8, then length:u16.
const FIELD_HEADER_SIZE: usize = 3;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Command {
    Connect,
    Connected,
    Filter,
    Notification,
    Ack,
    Ping,
    Pong,
    SetState,
    Unknown(u8),
}

impl Command {
    pub fn code(self) -> u8 {
        match self {
            Command::Connect => 7,
            Command::Connected => 8,
            Command::Filter => 9,
            Command::Notification => 10,
            Command::Ack => 11,
            Command::Ping => 12,
            Command::Pong => 13,
            Command::SetState => 20,
            Command::Unknown(code) => code,
        }
    }
}

impl From<u8> for Command {
    fn from(code: u8) -> Self {
        match code {
            7 => Command::Connect,
            8 => Command::Connected,
            9 => Command::Filter,
            10 => Command::Notification,
            11 => Command::Ack,
            12 => Command::Ping,
            13 => Command::Pong,
            20 => Command::SetState,
            code => Command::Unknown(code),
        }
    }
}

/// Outgoing messages borrow their fields and encode directly into one wire buffer.
pub(crate) enum Message<'lt> {
    Connect {
        previous: Option<&'lt PushToken>,
        certificate: &'lt [u8],
        nonce: &'lt [u8; 17],
        signature: &'lt [u8],
    },
    Filter {
        token: &'lt PushToken,
        topics: &'lt [[u8; 20]],
    },
    Send {
        token: &'lt PushToken,
        topic: &'lt [u8; 20],
        id: u32,
        payload: &'lt [u8],
    },
    Ack {
        token: &'lt PushToken,
        id: u32,
    },
    Ping,
    Pong,
    SetActive,
}

impl Message<'_> {
    pub async fn write(self, stream: &mut (impl AsyncWrite + Unpin)) -> Result<(), CourierError> {
        let frame = match self {
            Message::Connect {
                previous,
                certificate,
                nonce,
                signature,
            } => {
                let mut frame = Frame::new(Command::Connect);
                if let Some(token) = previous {
                    frame.push(1, token.as_bytes())?;
                }
                frame.push(2, [1])?; // Connect state; Apple's disconnect path sends 2.
                // Presence flags used by rustpush. The Windows client sends 0x4c
                // on connect and 0x48 for presence; individual bits are unverified.
                frame.push(5, 0x41u32.to_be_bytes())?;
                frame.push(12, certificate)?;
                frame.push(13, nonce)?;
                frame.push(14, signature)?;
                // Protocol version used by rustpush. Windows sends 5; macOS 27
                // sends 12 with additional capabilities we do not implement.
                frame.push(16, 9u16.to_be_bytes())?;
                frame
            }
            Message::Filter { token, topics } => {
                let mut frame = Frame::new(Command::Filter);
                frame.push(1, token.as_bytes())?;
                for topic in topics {
                    frame.push(2, topic)?;
                }
                frame
            }
            Message::Send {
                token,
                topic,
                id,
                payload,
            } => {
                let mut frame = Frame::new(Command::Notification);
                frame.push(4, id.to_be_bytes())?;
                // Outgoing messages reverse the incoming token/topic fields.
                frame.push(1, topic)?;
                frame.push(2, token.as_bytes())?;
                frame.push(3, payload)?;
                frame
            }
            Message::Ack { token, id } => {
                let mut frame = Frame::new(Command::Ack);
                frame.push(1, token.as_bytes())?;
                frame.push(4, id.to_be_bytes())?;
                frame.push(8, [0])?; // Success.
                frame
            }
            Message::Ping => Frame::new(Command::Ping),
            Message::Pong => Frame::new(Command::Pong),
            Message::SetActive => {
                let mut frame = Frame::new(Command::SetState);
                frame.push(1, [1])?;
                // rustpush's state interval. INT_MAX is observed; "infinite" is unverified.
                frame.push(2, 0x7fff_ffffu32.to_be_bytes())?;
                frame
            }
        };
        frame.write(stream).await
    }
}

#[derive(Debug, Error)]
pub enum CourierError {
    #[error("Courier I/O failed: {0}")]
    Io(#[from] io::Error),
    #[error("APS frame exceeds the size limit")]
    FrameTooLarge,
    #[error("APS field {0} exceeds the wire size limit")]
    FieldTooLarge(u8),
    #[error("Truncated APS field")]
    TruncatedField,
    #[error("Missing APS field {0}")]
    MissingField(u8),
    #[error("Invalid APS field {0}")]
    InvalidField(u8),
    #[error("Unexpected APS command {0}")]
    UnexpectedCommand(u8),
    #[error("The message belongs to another push token")]
    TokenMismatch,
    #[error("Courier operation timed out")]
    Timeout,
    #[error("APS message {0} was not acknowledged before the deadline; reconnect before retrying")]
    AcknowledgementTimeout(u32),
}

/// A header, then (id:u8, length:u16, value) fields, all big endian.
/// Fields borrow the wire buffer instead of allocating a Vec for each TLV.
pub(crate) struct Frame {
    bytes: BytesMut,
}

impl Frame {
    fn new(command: Command) -> Self {
        Frame {
            bytes: BytesMut::from(&[command.code(), 0, 0, 0, 0][..]),
        }
    }

pub fn command(&self) -> Command {
        self.bytes[0].into()
    }

fn push(&mut self, id: u8, value: impl AsRef<[u8]>) -> Result<(), CourierError> {
        let value = value.as_ref();
        let size = u16::try_from(value.len()).map_err(|_| CourierError::FieldTooLarge(id))?;
        let length = self.bytes.len() - HEADER_SIZE + FIELD_HEADER_SIZE + value.len();
        if length > MAX_FRAME_SIZE {
            Err(CourierError::FrameTooLarge)
        } else {
            self.bytes.extend_from_slice(&[id]);
            self.bytes.extend_from_slice(&size.to_be_bytes());
            self.bytes.extend_from_slice(value);
            self.bytes[1..HEADER_SIZE].copy_from_slice(&(length as u32).to_be_bytes());
            Ok(())
        }
    }

fn field_range(&self, id: u8) -> Option<Range<usize>> {
        let mut offset = HEADER_SIZE;
        let mut fields = std::iter::from_fn(|| {
            let header = self.bytes.get(offset..offset + FIELD_HEADER_SIZE)?;
            let size = u16::from_be_bytes([header[1], header[2]]) as usize;
            let value = offset + FIELD_HEADER_SIZE..offset + FIELD_HEADER_SIZE + size;
            offset = value.end;
            Some((header[0], value))
        });
        fields
            .find(|(field_id, _)| *field_id == id)
            .map(|(_, value)| value)
    }

pub fn get(&self, id: u8) -> Option<&[u8]> {
        self.field_range(id).and_then(|value| self.bytes.get(value))
    }

pub fn required(&self, id: u8) -> Result<&[u8], CourierError> {
        self.get(id).ok_or(CourierError::MissingField(id))
    }

/// An optional fixed-size field.
    pub fn get_array<const N: usize>(&self, id: u8) -> Result<Option<[u8; N]>, CourierError> {
        self.get(id)
            .map(|value| value.try_into().map_err(|_| CourierError::InvalidField(id)))
            .transpose()
    }

/// A mandatory fixed-size field.
    pub fn required_array<const N: usize>(&self, id: u8) -> Result<[u8; N], CourierError> {
        self.required(id)?
            .try_into()
            .map_err(|_| CourierError::InvalidField(id))
    }

pub fn into_field(self, id: u8) -> Result<Bytes, CourierError> {
        let value = self.field_range(id).ok_or(CourierError::MissingField(id))?;
        Ok(self.bytes.freeze().slice(value))
    }

async fn write(&self, stream: &mut (impl AsyncWrite + Unpin)) -> Result<(), CourierError> {
        tokio::time::timeout(TIMEOUT, async {
            stream.write_all(&self.bytes).await?;
            stream.flush().await
        })
        .await
        .map_err(|_| CourierError::Timeout)??;
        Ok(())
    }
}

/// Check that the fields exactly fill the frame, so reading them never goes out of bounds.
fn check_fields(mut fields: &[u8]) -> Result<(), CourierError> {
    while let [_, size_high, size_low, rest @ ..] = fields {
        let size = u16::from_be_bytes([*size_high, *size_low]) as usize;
        fields = rest.get(size..).ok_or(CourierError::TruncatedField)?;
    }

if fields.is_empty() {
        Ok(())
    } else {
        Err(CourierError::TruncatedField)
    }
}

#[derive(Default)]
pub(crate) struct FrameReader {
    // Retained when keepalive interrupts a read.
    buffer: BytesMut,
}

impl FrameReader {
    pub async fn read(
        &mut self,
        stream: &mut (impl AsyncRead + Unpin),
    ) -> Result<Frame, CourierError> {
        loop {
            if let Some(frame) = self.split_frame()? {
                return Ok(frame);
            }

self.buffer.reserve(4096);
            if stream.read_buf(&mut self.buffer).await? == 0 {
                return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into());
            }
        }
    }

/// Take the first complete frame out of the buffer.
    fn split_frame(&mut self) -> Result<Option<Frame>, CourierError> {
        let length: Option<[u8; 4]> = self
            .buffer
            .get(1..HEADER_SIZE)
            .and_then(|length| length.try_into().ok());

match length.map(|length| u32::from_be_bytes(length) as usize) {
            None => Ok(None),
            Some(length) if length > MAX_FRAME_SIZE => Err(CourierError::FrameTooLarge),
            Some(length) if self.buffer.len() < HEADER_SIZE + length => Ok(None),
            Some(length) => {
                check_fields(&self.buffer[HEADER_SIZE..HEADER_SIZE + length])?;
                Ok(Some(Frame {
                    bytes: self.buffer.split_to(HEADER_SIZE + length),
                }))
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

#[test]
    fn frames_round_trip_through_the_reader() {
        let mut frame = Frame::new(Command::Ack);
        frame.push(1, [7; 32]).unwrap();
        frame.push(4, 42u32.to_be_bytes()).unwrap();
        frame.push(8, [0]).unwrap();

let mut reader = FrameReader::default();
        reader.buffer.extend_from_slice(&frame.bytes);
        reader
            .buffer
            .extend_from_slice(&[Command::Ping.code(), 0, 0, 0]);

let frame = reader.split_frame().unwrap().unwrap();
        assert_eq!(frame.command(), Command::Ack);
        assert_eq!(frame.get(1), Some(&[7; 32][..]));
        assert_eq!(frame.required_array(4).map(u32::from_be_bytes).unwrap(), 42);
        assert!(matches!(
            frame.get_array::<2>(8),
            Err(CourierError::InvalidField(8))
        ));
        assert!(matches!(
            frame.required(3),
            Err(CourierError::MissingField(3))
        ));
        // The truncated ping stays buffered until the rest of it arrives.
        assert!(reader.split_frame().unwrap().is_none());
        assert_eq!(reader.buffer.len(), 4);
    }

#[test]
    fn truncated_fields_are_rejected() {
        assert!(check_fields(&[]).is_ok());
        assert!(check_fields(&[1, 0, 1, 0xff]).is_ok());
        assert!(matches!(
            check_fields(&[1, 0, 2, 0xff]),
            Err(CourierError::TruncatedField)
        ));
        assert!(matches!(
            check_fields(&[1, 0, 1, 0xff, 2]),
            Err(CourierError::TruncatedField)
        ));
    }
}

apple-crates

Refactor and unify a lot of things for another project and implement airdrop support