airdrop/src/connect.rs (a74924e) - apple-crates · Git

connect.rs

//! HTTPS/1.1 to one AirDrop receiver: a keep-alive TLS connection that
//! carries POSTs. Adapted from omdrop-plugin's `omdrop-send` (MIT).
//!
//! Every socket operation is bounded. 021922Z: our third-leg ACK and TLS
//! Client Hello were emitted and lost, the Mac retransmitted SYN-ACK 13 times,
//! and TCP backed off to rto:120000. With no timeout the client sat in that
//! backoff for 4.5 minutes, past the end of the watched hold, instead of
//! failing and reporting. A wedged send must cost seconds, not a hold. An
//! optional absolute deadline caps the whole exchange as well (send-to-peer's
//! name lookups, which Python bounded with a subprocess timeout).

use std::io::{self, Read, Write};
use std::net::{IpAddr, Ipv6Addr, SocketAddr, SocketAddrV6, TcpStream};
use std::sync::Arc;
use std::time::{Duration, Instant};

use rustls::pki_types::ServerName;

fn io_err(
    kind: io::ErrorKind,
    message: impl Into<Box<dyn std::error::Error + Send + Sync>>,
) -> io::Error {
    io::Error::new(kind, message)
}
use rustls::{ClientConfig, ClientConnection, StreamOwned};

/// Responses are small plists (a Mac's /Discover answer is ~4 KB with its
/// record); anything far larger is not an AirDrop receiver talking.
const MAX_RESPONSE_BYTES: usize = 4 << 20;
const MAX_HEAD_BYTES: usize = 64 << 10;
/// Chunk size for a streamed (chunked) request body.
const CHUNK_BYTES: usize = 64 << 10;

/// A receiver's address, and the host text Python used for it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Endpoint {
    pub addr: SocketAddr,
    /// `fe80::1%awdl0`, `::1`, `127.0.0.1`: what the `Host` header carries.
    pub host: String,
}

impl Endpoint {
    /// An IP literal, optionally with a `%zone`, on `port`. A link-local
    /// address without a zone is ambiguous, so it gets `iface`'s (fe80::1 ->
    /// fe80::1%awdl0). Host names are refused, as Python's `ip_address()`
    /// refused them.
    pub fn resolve(host: &str, port: u16, iface: &str) -> io::Result<Endpoint> {
        let (addr_text, zone) = match host.split_once('%') {
            Some((a, z)) => (a, Some(z)),
            None => (host, None),
        };
        let bad = || {
            io_err(
                io::ErrorKind::InvalidInput,
                format!("{host:?} does not appear to be an IPv4 or IPv6 address"),
            )
        };
        let ip: IpAddr = addr_text.parse().map_err(|_| bad())?;
        match ip {
            IpAddr::V4(v4) => {
                if zone.is_some() {
                    return Err(bad());
                }
                Ok(Endpoint {
                    addr: SocketAddr::new(IpAddr::V4(v4), port),
                    host: v4.to_string(),
                })
            }
            IpAddr::V6(v6) => {
                let zone = match zone {
                    Some(z) => Some(z.to_owned()),
                    None if v6.is_unicast_link_local() => Some(iface.to_owned()),
                    None => None,
                };
                Self::v6(v6, port, zone.as_deref())
            }
        }
    }

/// An IPv6 address with an optional zone (interface name or index).
    pub fn v6(addr: Ipv6Addr, port: u16, zone: Option<&str>) -> io::Result<Endpoint> {
        let (scope, host) = match zone {
            Some(z) => {
                let scope = match z.parse::<u32>() {
                    Ok(n) => n,
                    Err(_) => crate::net::scope_id(z)?,
                };
                (scope, format!("{addr}%{z}"))
            }
            None => (0, addr.to_string()),
        };
        Ok(Endpoint {
            addr: SocketAddr::V6(SocketAddrV6::new(addr, port, 0, scope)),
            host,
        })
    }

pub fn port(&self) -> u16 {
        self.addr.port()
    }

/// The `Host` header value, as http.client built it.
    pub fn host_header(&self) -> String {
        let host = if self.host.contains(':') {
            format!("[{}]", self.host)
        } else {
            self.host.clone()
        };
        if self.port() == 443 {
            host
        } else {
            format!("{host}:{}", self.port())
        }
    }
}

/// The socket under TLS, with per-operation timeouts capped by a deadline.
#[derive(Debug)]
struct Timed {
    sock: TcpStream,
    op: Option<Duration>,
    deadline: Option<Instant>,
}

impl Timed {
    /// Arm the socket for one operation. Without a deadline the timeouts set
    /// in [`Timed::set_op`] stand.
    fn arm(&self, read: bool) -> io::Result<()> {
        let Some(deadline) = self.deadline else {
            return Ok(());
        };
        let left = deadline.saturating_duration_since(Instant::now());
        if left.is_zero() {
            return Err(io_err(io::ErrorKind::TimedOut, "timed out"));
        }
        let t = Some(self.op.map_or(left, |op| op.min(left)));
        if read {
            self.sock.set_read_timeout(t)
        } else {
            self.sock.set_write_timeout(t)
        }
    }

fn set_op(&mut self, op: Option<Duration>) -> io::Result<()> {
        self.op = op;
        self.sock.set_read_timeout(op)?;
        self.sock.set_write_timeout(op)
    }
}

impl Read for Timed {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.arm(true)?;
        self.sock.read(buf)
    }
}

impl Write for Timed {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.arm(false)?;
        self.sock.write(buf)
    }

fn flush(&mut self) -> io::Result<()> {
        self.sock.flush()
    }
}

/// A request body.
pub enum Body<'a> {
    /// Sent with a Content-Length.
    Bytes(&'a [u8]),
    /// Sent chunked, as sharingd sends /Upload: the closure writes the body.
    Stream(&'a mut dyn FnMut(&mut dyn Write) -> io::Result<()>),
}

/// One response, read in full.
#[derive(Debug)]
pub struct Response {
    pub status: u16,
    pub reason: String,
    pub body: Vec<u8>,
}

/// A keep-alive HTTPS connection.
pub struct Conn {
    stream: StreamOwned<ClientConnection, Timed>,
    rbuf: Vec<u8>,
    /// The server said `Connection: close`, or the body ran to EOF.
    pub closed: bool,
}

impl Conn {
    /// TCP connect and complete the TLS handshake, each operation bounded by
    /// `op` and everything by `deadline`. A peer's certificate is accepted
    /// whoever signed it (see `crate::tls`).
    pub fn open(
        endpoint: &Endpoint,
        tls: &Arc<ClientConfig>,
        op: Option<Duration>,
        deadline: Option<Instant>,
    ) -> io::Result<Conn> {
        let budget = match (op, deadline) {
            (op, Some(d)) => {
                let left = d.saturating_duration_since(Instant::now());
                if left.is_zero() {
                    return Err(io_err(io::ErrorKind::TimedOut, "timed out"));
                }
                Some(op.map_or(left, |o| o.min(left)))
            }
            (op, None) => op,
        };
        let sock = match budget {
            Some(t) => TcpStream::connect_timeout(&endpoint.addr, t)?,
            None => TcpStream::connect(endpoint.addr)?,
        };
        // Each request is written whole and then answered; Nagle would hold a
        // trailing TLS record for the peer's delayed ACK (measured on
        // loopback: /Discover 0.09 s with Nagle, 0.04 s without).
        sock.set_nodelay(true)?;
        let mut timed = Timed {
            sock,
            op: None,
            deadline,
        };
        timed.set_op(op)?;
        // No SNI: the name is an address, as it was for Python.
        let name = ServerName::IpAddress(endpoint.addr.ip().into());
        let conn = ClientConnection::new(tls.clone(), name).map_err(io::Error::other)?;
        let mut stream = StreamOwned::new(conn, timed);
        while stream.conn.is_handshaking() {
            stream.conn.complete_io(&mut stream.sock)?;
        }
        Ok(Conn {
            stream,
            rbuf: Vec::new(),
            closed: false,
        })
    }

/// The server ended the connection (`Connection: close`, or EOF).
    pub fn is_closed(&self) -> bool {
        self.closed
    }

/// Change the per-operation timeout (the Ask's long one, which then holds
    /// for the Upload on the same connection, as `sock.settimeout` did).
    pub fn set_op_timeout(&mut self, op: Option<Duration>) -> io::Result<()> {
        self.stream.sock.set_op(op)
    }

/// POST and read the whole answer.
    pub fn post(
        &mut self,
        path: &str,
        host_header: &str,
        headers: &[(&str, String)],
        body: Body<'_>,
    ) -> io::Result<Response> {
        // http.client's order: request line, Host, the framing header, then
        // the caller's headers.
        let mut head = format!("POST {path} HTTP/1.1\r\nHost: {host_header}\r\n");
        match &body {
            Body::Bytes(b) => head.push_str(&format!("Content-Length: {}\r\n", b.len())),
            Body::Stream(_) => head.push_str("Transfer-Encoding: chunked\r\n"),
        }
        for (k, v) in headers {
            head.push_str(&format!("{k}: {v}\r\n"));
        }
        head.push_str("\r\n");
        match body {
            Body::Bytes(b) => {
                let mut all = head.into_bytes();
                all.extend_from_slice(b);
                self.stream.write_all(&all)?;
            }
            Body::Stream(produce) => {
                self.stream.write_all(head.as_bytes())?;
                let mut chunked = Chunked {
                    inner: &mut self.stream,
                    buf: Vec::with_capacity(CHUNK_BYTES),
                };
                produce(&mut chunked)?;
                chunked.finish()?;
            }
        }
        self.stream.flush()?;
        self.read_response()
    }

fn fill(&mut self) -> io::Result<usize> {
        let mut chunk = [0u8; 16 << 10];
        let n = self.stream.read(&mut chunk)?;
        self.rbuf
            .extend_from_slice(chunk.get(..n).unwrap_or_default());
        Ok(n)
    }

fn read_response(&mut self) -> io::Result<Response> {
        loop {
            let (status, reason, head_len, framing, close) = self.read_head()?;
            self.rbuf.drain(..head_len);
            // http.client skips interim 100 Continue answers.
            if (100..200).contains(&status) {
                continue;
            }
            let body = match framing {
                Framing::None => Vec::new(),
                Framing::Length(n) => self.read_exact_body(n)?,
                Framing::Chunked => self.read_chunked()?,
                Framing::UntilClose => {
                    self.closed = true;
                    self.read_to_close()?
                }
            };
            if close {
                self.closed = true;
            }
            return Ok(Response {
                status,
                reason,
                body,
            });
        }
    }

fn read_head(&mut self) -> io::Result<(u16, String, usize, Framing, bool)> {
        loop {
            let mut headers = [httparse::EMPTY_HEADER; 64];
            let mut resp = httparse::Response::new(&mut headers);
            match resp.parse(&self.rbuf) {
                Ok(httparse::Status::Complete(n)) => {
                    let status = resp.code.unwrap_or(0);
                    let reason = resp.reason.unwrap_or("").to_owned();
                    let http10 = resp.version == Some(0);
                    let mut length = None;
                    let mut chunked = false;
                    let mut conn_close = http10;
                    for h in resp.headers.iter() {
                        let value = String::from_utf8_lossy(h.value);
                        if h.name.eq_ignore_ascii_case("content-length") {
                            let n: usize = value
                                .trim()
                                .parse()
                                .map_err(|_| bad_status("invalid Content-Length"))?;
                            length = Some(n);
                        } else if h.name.eq_ignore_ascii_case("transfer-encoding") {
                            chunked = value
                                .split(',')
                                .next_back()
                                .is_some_and(|t| t.trim().eq_ignore_ascii_case("chunked"));
                        } else if h.name.eq_ignore_ascii_case("connection") {
                            for token in value.split(',') {
                                let token = token.trim();
                                if token.eq_ignore_ascii_case("close") {
                                    conn_close = true;
                                } else if token.eq_ignore_ascii_case("keep-alive") && http10 {
                                    conn_close = false;
                                }
                            }
                        }
                    }
                    let framing = if status == 204 || status == 304 || (100..200).contains(&status)
                    {
                        Framing::None
                    } else if chunked {
                        Framing::Chunked
                    } else if let Some(n) = length {
                        if n > MAX_RESPONSE_BYTES {
                            return Err(io_err(
                                io::ErrorKind::InvalidData,
                                format!("response of {n} bytes is too large"),
                            ));
                        }
                        Framing::Length(n)
                    } else {
                        Framing::UntilClose
                    };
                    return Ok((status, reason, n, framing, conn_close));
                }
                Ok(httparse::Status::Partial) => {
                    if self.rbuf.len() > MAX_HEAD_BYTES {
                        return Err(io_err(io::ErrorKind::InvalidData, "response head too long"));
                    }
                }
                Err(e) => return Err(bad_status(&e.to_string())),
            }
            let had = self.rbuf.len();
            let n = match self.fill() {
                Ok(n) => n,
                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => 0,
                Err(e) => return Err(e),
            };
            if n == 0 {
                self.closed = true;
                return Err(if had == 0 {
                    io_err(
                        io::ErrorKind::ConnectionAborted,
                        "Remote end closed connection without response",
                    )
                } else {
                    io_err(
                        io::ErrorKind::UnexpectedEof,
                        "connection closed inside the response head",
                    )
                });
            }
        }
    }

fn read_exact_body(&mut self, n: usize) -> io::Result<Vec<u8>> {
        while self.rbuf.len() < n {
            if self.fill_or_eof()? == 0 {
                self.closed = true;
                return Err(io_err(
                    io::ErrorKind::UnexpectedEof,
                    format!(
                        "{} bytes read, {} more expected",
                        self.rbuf.len(),
                        n - self.rbuf.len()
                    ),
                ));
            }
        }
        Ok(self.rbuf.drain(..n).collect())
    }

fn read_line(&mut self) -> io::Result<String> {
        loop {
            if let Some(i) = self.rbuf.windows(2).position(|w| w == b"\r\n") {
                let line: Vec<u8> = self.rbuf.drain(..i + 2).collect();
                return Ok(String::from_utf8_lossy(line.get(..i).unwrap_or_default()).into_owned());
            }
            if self.rbuf.len() > MAX_HEAD_BYTES {
                return Err(io_err(io::ErrorKind::InvalidData, "chunk line too long"));
            }
            if self.fill_or_eof()? == 0 {
                self.closed = true;
                return Err(io_err(
                    io::ErrorKind::UnexpectedEof,
                    "connection closed inside a chunked body",
                ));
            }
        }
    }

fn read_chunked(&mut self) -> io::Result<Vec<u8>> {
        let mut body = Vec::new();
        loop {
            let line = self.read_line()?;
            let size_text = line.split(';').next().unwrap_or("").trim();
            let size = usize::from_str_radix(size_text, 16)
                .map_err(|_| io_err(io::ErrorKind::UnexpectedEof, "invalid chunk size"))?;
            if size == 0 {
                // Trailers, up to the empty line.
                while !self.read_line()?.is_empty() {}
                return Ok(body);
            }
            if body.len().saturating_add(size) > MAX_RESPONSE_BYTES {
                return Err(io_err(
                    io::ErrorKind::InvalidData,
                    "chunked response is too large",
                ));
            }
            let chunk = self.read_exact_body(size)?;
            body.extend_from_slice(&chunk);
            if !self.read_line()?.is_empty() {
                return Err(io_err(
                    io::ErrorKind::UnexpectedEof,
                    "chunk not followed by CRLF",
                ));
            }
        }
    }

fn read_to_close(&mut self) -> io::Result<Vec<u8>> {
        while self.fill_or_eof()? > 0 {
            if self.rbuf.len() > MAX_RESPONSE_BYTES {
                return Err(io_err(io::ErrorKind::InvalidData, "response is too large"));
            }
        }
        Ok(std::mem::take(&mut self.rbuf))
    }

/// A read where a TLS close without close_notify counts as EOF, as
    /// Python's `suppress_ragged_eofs` made it.
    fn fill_or_eof(&mut self) -> io::Result<usize> {
        match self.fill() {
            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => Ok(0),
            other => other,
        }
    }
}

enum Framing {
    None,
    Length(usize),
    Chunked,
    UntilClose,
}

fn bad_status(why: &str) -> io::Error {
    io_err(io::ErrorKind::InvalidData, why)
}

/// HTTP/1.1 chunked transfer coding over a writer.
struct Chunked<'a, W: Write> {
    inner: &'a mut W,
    buf: Vec<u8>,
}

impl<W: Write> Chunked<'_, W> {
    fn emit(&mut self) -> io::Result<()> {
        if self.buf.is_empty() {
            return Ok(());
        }
        let mut frame = format!("{:x}\r\n", self.buf.len()).into_bytes();
        frame.extend_from_slice(&self.buf);
        frame.extend_from_slice(b"\r\n");
        self.inner.write_all(&frame)?;
        self.buf.clear();
        Ok(())
    }

fn finish(mut self) -> io::Result<()> {
        self.emit()?;
        self.inner.write_all(b"0\r\n\r\n")
    }
}

impl<W: Write> Write for Chunked<'_, W> {
    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
        let room = CHUNK_BYTES - self.buf.len();
        let n = room.min(data.len());
        self.buf
            .extend_from_slice(data.get(..n).unwrap_or_default());
        if self.buf.len() >= CHUNK_BYTES {
            self.emit()?;
        }
        Ok(n)
    }

fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

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

#[test]
    fn endpoints_carry_the_zone_python_gave_them() {
        let lo = Endpoint::resolve("::1", 8771, "awdl0").unwrap();
        assert_eq!(lo.host_header(), "[::1]:8771");
        let v4 = Endpoint::resolve("127.0.0.1", 443, "awdl0").unwrap();
        assert_eq!(v4.host_header(), "127.0.0.1");
        // A numeric zone needs no interface lookup.
        let ll = Endpoint::resolve("fe80::e09d:3ff:fe6c:5823%7", 8770, "awdl0").unwrap();
        assert_eq!(ll.host_header(), "[fe80::e09d:3ff:fe6c:5823%7]:8770");
        match ll.addr {
            SocketAddr::V6(a) => assert_eq!(a.scope_id(), 7),
            SocketAddr::V4(_) => panic!("v4"),
        }
        let lo_iface = Endpoint::resolve("fe80::1", 1, "lo").unwrap();
        assert_eq!(lo_iface.host, "fe80::1%lo");
        let err = Endpoint::resolve("localhost", 1, "lo").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
        assert!(Endpoint::resolve("127.0.0.1%lo", 1, "lo").is_err());
    }

#[test]
    fn chunked_framing() {
        let mut out = Vec::new();
        let mut c = Chunked {
            inner: &mut out,
            buf: Vec::new(),
        };
        c.write_all(b"hello").unwrap();
        c.finish().unwrap();
        assert_eq!(out, b"5\r\nhello\r\n0\r\n\r\n");
    }
}

apple-crates

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