jkipsec 0.1.0

//! Per-IKE-SA state machine. Internal - drive through [`crate::api`].
//!
//! Owned by `Server` (pub(crate)); one `Session` exists per Initiator SPI.
//! The session holds the keys derived from IKE_SA_INIT and keeps the original
//! request / response bytes around because the IKE_AUTH PSK signature is
//! computed over `RealMessage1 | NonceR | prf(SK_pi, IDi')` (RFC 7296 §2.15).
//!
//! Two suites are supported (selected at runtime against the peer's offer):
//! - AES-GCM-16-256 (AEAD) + PRF SHA-256 + DH 19
//! - AES-CBC-256 + HMAC-SHA-256-128 + PRF SHA-256 + DH 19

// Jackson Coxson

use std::collections::HashMap;
use std::future::Future;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::pin::Pin;

// rand 0.10: `Rng` (re-export of rand_core::Rng) provides `fill_bytes`;
// `RngExt` provides `random()`. `RngCore` was renamed in 0.10.
use rand::{Rng, RngExt};
use sha1::{Digest, Sha1};
use tracing::{debug, info, warn};

use crate::esp::{EspTunnel, InboundFlavor, OutboundFlavor};

use crate::crypto::{self, DhEphemeral, aes_gcm_open, aes_gcm_seal};
use crate::ike::{
    self, ExchangeType, Flags, Message, OutPayload, PayloadKind,
    payloads::{
        AuthMethod, AuthPayload, EncryptedPayload, IdPayload, IdType, KePayload, NoncePayload,
        NotifyPayload, notify_type,
    },
    sa::{AttrValue, Attribute, Proposal, ProtocolId, SaPayload, Transform, TransformType},
    walk_chain,
};

/// Server-wide configuration. The PSK is shared with the iOS client out of band.
/// Handed to the user's auth callback. Carries the peer's identity and all
/// state the callback needs to verify the PSK AUTH payload against a
/// candidate `auth_key`.
///
/// The callback is **not** given a raw PSK because iOS never sends one. Instead,
/// the callback looks up the auth key for this identity (e.g. from a
/// database) and calls [`Self::verify`] to check it against the AUTH payload
/// iOS sent. The callback returns `true` to approve, `false` to reject.
///
/// `auth_key` is the 32-byte output of [`crypto::derive_auth_key`] applied to
/// the raw PSK. Store that in your database; the raw PSK can be discarded.
///
/// RFC 7296 §2.15 defines the verification. We provide [`Self::verify`] so
/// user code stays out of IKEv2 crypto internals.
pub struct AuthChallenge {
    /// Raw IDi body: first byte is id_type, next 3 reserved, rest identity.
    pub identity_raw: Vec<u8>,
    /// Source address of the peer that sent this IKE_AUTH.
    pub peer: SocketAddr,
    // Captured state needed for verification - internal.
    sk_pi: Vec<u8>,
    init_request: Vec<u8>,
    nonce_r: Vec<u8>,
    auth_data: Vec<u8>,
}

impl AuthChallenge {
    /// Identity bytes (after the 4-byte IDi header).
    pub fn identity(&self) -> &[u8] {
        if self.identity_raw.len() > 4 {
            &self.identity_raw[4..]
        } else {
            &self.identity_raw
        }
    }

    /// Lossy UTF-8 view of the identity for display/logging.
    pub fn identity_str(&self) -> String {
        String::from_utf8_lossy(self.identity()).into()
    }

    /// Check whether `auth_key` verifies the client's AUTH payload.
    ///
    /// Call this from the auth callback with a candidate key retrieved from
    /// your store. Returns `true` iff the key matches what iOS authenticated
    /// with. Safe to call multiple times (e.g. to support key rotation).
    pub fn verify(&self, auth_key: &[u8; 32]) -> bool {
        let maced_id = crypto::prf(&self.sk_pi, &self.identity_raw);
        let mut signed =
            Vec::with_capacity(self.init_request.len() + self.nonce_r.len() + maced_id.len());
        signed.extend_from_slice(&self.init_request);
        signed.extend_from_slice(&self.nonce_r);
        signed.extend_from_slice(&maced_id);
        let computed = crypto::prf(auth_key, &signed);
        use subtle::ConstantTimeEq;
        self.auth_data.len() == computed.len() && bool::from(computed.ct_eq(&self.auth_data))
    }
}

/// Decision returned from the auth callback.
///
/// - [`AuthDecision::Approve`] - accept the client. Must carry the 32-byte
///   auth key that verified (the library needs it to compute **our** AUTH
///   payload, which uses the same key). Produced by
///   [`AuthChallenge::approve_with`] or constructed directly after a
///   successful [`AuthChallenge::verify`].
/// - [`AuthDecision::Reject`] - deny the client. The library drops the
///   session and iOS times out.
#[derive(Debug)]
pub enum AuthDecision {
    /// Accept the client. The library uses the supplied `auth_key` to
    /// compute its own AUTH payload (responder side, RFC 7296 §2.15).
    Approve {
        /// 32-byte derived auth key: same value that just verified the
        /// peer's AUTH. Construct via `AuthChallenge::approve_with` to avoid
        /// needing to remember which key matched.
        auth_key: [u8; 32],
    },
    /// Deny the client. The session is dropped and the peer times out.
    Reject,
}

impl AuthChallenge {
    /// Convenience: verify `auth_key` and return an [`AuthDecision`] in one
    /// step. Prefer this when your store has exactly one candidate key per
    /// identity.
    pub fn approve_with(&self, auth_key: &[u8; 32]) -> AuthDecision {
        if self.verify(auth_key) {
            AuthDecision::Approve {
                auth_key: *auth_key,
            }
        } else {
            AuthDecision::Reject
        }
    }
}

/// Async auth callback. Typical implementation:
///
/// ```ignore
/// Box::new(|challenge| Box::pin(async move {
///     match db.get_auth_key(challenge.identity()).await {
///         Some(key) => challenge.approve_with(&key),
///         None => AuthDecision::Reject,
///     }
/// }))
/// ```
///
/// For key-rotation / multi-key scenarios call [`AuthChallenge::verify`]
/// directly and build the [`AuthDecision`] yourself.
///
/// **Note:** the callback runs inline on the single server task. Long-running
/// callbacks delay other IKE traffic; keep it fast (in-memory map or quick
/// DB lookup).
pub type AuthFn =
    Box<dyn Fn(AuthChallenge) -> Pin<Box<dyn Future<Output = AuthDecision> + Send>> + Send + Sync>;

pub(crate) struct Config {
    /// Async callback invoked during IKE_AUTH to validate the connecting
    /// peer's identity and supply the auth key for AUTH verification.
    pub auth: AuthFn,
    /// Address iOS thinks our gateway has. Used in NAT_DETECTION_SOURCE_IP.
    /// Iff this matches what iOS sent for NAT_DETECTION_DESTINATION_IP, iOS
    /// won't think we're behind a NAT.
    pub local_ip: IpAddr,
    /// Port iOS sent the IKE_SA_INIT *to* (typically 500). Independent of the
    /// real bind port - see the README for the port-forwarding setup.
    pub local_port: u16,
    /// Default IDr to announce if the initiator doesn't tell us which identity
    /// it wants. When the initiator *does* send IDr (as iOS does), we echo
    /// that back instead so the profile's Remote ID matches.
    pub our_identity: String,
    /// Virtual IP handed out via CP(CFG_REPLY). Point-to-point /32 is fine.
    pub virtual_ip: Ipv4Addr,
    /// DNS server advertised to the client. Use `0.0.0.0` to skip.
    pub virtual_dns: Ipv4Addr,
    /// Our "gateway" IP inside the tunnel. TSr is narrowed to this single
    /// /32 so iOS only routes traffic destined for it through the VPN -
    /// i.e. a pure split-tunnel where our process is the only reachable
    /// destination. All other iOS traffic goes out its normal default
    /// route.
    pub gateway_ip: Ipv4Addr,
}

// ---------------------------------------------------------------- messages

/// Classified datagram from a UDP reader, ready for server dispatch.
pub(crate) enum InboundMsg {
    /// IKE message (non-ESP marker already stripped for 4500 traffic).
    Ike {
        data: Vec<u8>,
        peer: SocketAddr,
        /// Channel to the UDP writer for this socket - the server sends its
        /// reply here, and the writer handles framing (e.g. non-ESP marker).
        reply_tx: crossfire::MTx<ReplyFlavor>,
    },
    /// ESP datagram (starts with 4-byte SPI).
    Esp { data: Vec<u8> },
    /// Disconnect a client identified by its Initiator SPI. Sends IKE DELETE
    /// and cleans up session state.
    Disconnect {
        initiator_spi: u64,
        reply_tx: crossfire::MTx<ReplyFlavor>,
    },
}

/// Outbound datagram sent by the server -> UDP writer.
pub(crate) struct OutboundMsg {
    pub data: Vec<u8>,
    pub peer: SocketAddr,
}

pub(crate) type InboundChanFlavor = crossfire::mpsc::List<InboundMsg>;
pub(crate) type SessionChanFlavor = crossfire::mpsc::List<EstablishedSession>;
/// Type alias for the IKE reply channel flavor used by UDP writers.
pub(crate) type ReplyFlavor = crossfire::mpsc::List<OutboundMsg>;

// ------------------------------------------------------------------ Server

/// Top-level dispatcher: holds one [`Session`] per Initiator SPI. Runs as a
/// **single long-lived task** that owns all session state - no `Arc<Mutex>`
/// needed. Multiple iOS devices each get their own session, CHILD_SA, and
/// [`EstablishedSession`] emitted on the `new_sessions` channel.
pub(crate) struct Server {
    pub config: Config,
    pub sessions: HashMap<u64, Session>,
    inbound_rx: crossfire::AsyncRx<InboundChanFlavor>,
    new_sessions: crossfire::MTx<SessionChanFlavor>,
}

/// Delivered to the caller once `IKE_AUTH` succeeds. Contains everything
/// needed to (a) hand `tunnel` to `jktcp::Adapter::new` and (b) spawn the
/// background task that encrypts outbound IP packets and sends them via UDP.
pub(crate) struct EstablishedSession {
    pub peer: SocketAddr,
    pub initiator_spi: u64,
    /// The identity the peer sent in IKE_AUTH (IDi payload body). For iOS
    /// this is whatever "Local Identifier" is configured in the VPN profile
    /// - e.g. `IdType::Rfc822Addr("alice@fleet")` or `IdType::Fqdn("device-a")`.
    ///   If unconfigured, iOS defaults to its interface IPv6 address.
    pub peer_identity: Vec<u8>,
    pub virtual_ip: Ipv4Addr,
    pub gateway_ip: Ipv4Addr,
    pub tunnel: EspTunnel,
    /// Drain this in a background task to obtain outbound IP packets to
    /// encrypt + `send_to`.
    pub outbound_rx: crossfire::AsyncRx<OutboundFlavor>,
    /// Negotiated suite (so the outbound ESP task knows whether to AEAD-encrypt
    /// or CBC+HMAC-encrypt).
    pub suite: crypto::Suite,
    pub outbound_spi: u32,
    pub outbound_key: Vec<u8>,
    pub outbound_salt: Vec<u8>,
    /// Outbound HMAC integrity key. Empty for AEAD.
    pub outbound_integ: Vec<u8>,
}

impl Server {
    /// Create the server plus its external wiring:
    ///
    /// - `inbound_tx` - clone this into each UDP reader task.
    /// - `new_sessions_rx` - receives one [`EstablishedSession`] per
    ///   successful IKE_AUTH; spawn the per-session outbound ESP task from it.
    pub fn new(
        config: Config,
    ) -> (
        Self,
        crossfire::MTx<InboundChanFlavor>,
        crossfire::AsyncRx<SessionChanFlavor>,
    ) {
        let (in_tx, in_rx) = crossfire::mpsc::unbounded_async::<InboundMsg>();
        let (sess_tx, sess_rx) = crossfire::mpsc::unbounded_async::<EstablishedSession>();
        (
            Self {
                config,
                sessions: HashMap::new(),
                inbound_rx: in_rx,
                new_sessions: sess_tx,
            },
            in_tx,
            sess_rx,
        )
    }

    /// The server task's main loop. Call this from a single `tokio::spawn`.
    /// It owns all session state - no locking needed.
    pub async fn run(&mut self) {
        loop {
            match crossfire::AsyncRxTrait::recv(&self.inbound_rx).await {
                Ok(InboundMsg::Ike {
                    data,
                    peer,
                    reply_tx,
                }) => {
                    if let Some(reply) = self.handle(&data, peer).await {
                        use crossfire::BlockingTxTrait;
                        let _ = reply_tx.send(OutboundMsg { data: reply, peer });
                    }
                }
                Ok(InboundMsg::Esp { data }) => {
                    self.handle_esp(&data);
                }
                Ok(InboundMsg::Disconnect {
                    initiator_spi,
                    reply_tx,
                }) => {
                    self.handle_disconnect(initiator_spi, &reply_tx);
                }
                Err(_) => break, // all senders dropped
            }
        }
    }

    /// Process one inbound datagram. Async because `handle_auth` awaits the
    /// user-supplied auth callback. Returns `Some(reply)` to send back, or
    /// `None` for malformed / unhandled traffic.
    pub async fn handle(&mut self, datagram: &[u8], peer: SocketAddr) -> Option<Vec<u8>> {
        let msg = match Message::parse(datagram) {
            Ok(m) => m,
            Err(e) => {
                warn!(%peer, "parse failed: {e}");
                return None;
            }
        };

        match msg.header.exchange_type {
            ExchangeType::IkeSaInit => self.handle_sa_init(&msg, peer),
            ExchangeType::IkeAuth => self.handle_auth(&msg, peer).await,
            ExchangeType::Informational => self.handle_informational(&msg),
            other => {
                warn!(?other, "unhandled exchange type");
                None
            }
        }
    }

    /// Decrypt an ESP packet and return the inner IP payload plus its Next
    /// Header byte. `datagram` is the full UDP payload (starting with the
    /// 4-byte SPI). The caller dispatches the inner packet - at the moment
    /// just logs it; task #7 wires it into jktcp.
    pub fn handle_esp(&mut self, datagram: &[u8]) -> Option<(u8, Vec<u8>)> {
        if datagram.len() < 4 {
            return None;
        }
        let spi = u32::from_be_bytes(datagram[0..4].try_into().unwrap());
        // Look up the session that owns this inbound SPI. Linear scan is
        // fine - one session per active iOS client, N is tiny.
        let session = self
            .sessions
            .values()
            .find(|s| s.child.as_ref().map(|c| c.inbound_spi) == Some(spi))?;
        let child = session.child.as_ref()?;
        match crate::esp::decrypt(
            child.suite,
            &child.inbound_key,
            &child.inbound_salt,
            &child.inbound_integ,
            datagram,
        ) {
            Ok(d) => {
                debug!(
                    spi = format_args!("{spi:08x}"),
                    seq = d.seq,
                    next_header = d.next_header,
                    len = d.payload.len(),
                    "ESP decrypt ok"
                );
                if let Some(tx) = &session.inbound_tx {
                    use crossfire::BlockingTxTrait;
                    let _ = tx.send(d.payload.clone());
                }
                Some((d.next_header, d.payload))
            }
            Err(e) => {
                warn!(spi = format_args!("{spi:08x}"), "ESP decrypt failed: {e}");
                None
            }
        }
    }

    /// Send an IKE DELETE notification to the peer and remove the session.
    fn handle_disconnect(&mut self, initiator_spi: u64, reply_tx: &crossfire::MTx<ReplyFlavor>) {
        let Some(session) = self.sessions.remove(&initiator_spi) else {
            return;
        };
        if let Some(delete_msg) = build_delete_ike_sa(
            session.initiator_spi,
            session.responder_spi,
            0, // responder-initiated INFORMATIONAL, msg_id=0
            &session.keys,
        ) {
            use crossfire::BlockingTxTrait;
            let _ = reply_tx.send(OutboundMsg {
                data: delete_msg,
                peer: session.peer,
            });
            info!(
                ispi = format_args!("{:016x}", initiator_spi),
                "sent IKE DELETE"
            );
        }
    }

    /// Minimal INFORMATIONAL handler: ack with an empty encrypted response so
    /// the peer doesn't keep retrying. We don't yet act on DELETE notifies, so
    /// a subsequent connect attempt spins up a fresh session anyway.
    fn handle_informational(&mut self, msg: &Message<'_>) -> Option<Vec<u8>> {
        if msg.header.flags.is_response() {
            return None;
        }
        let session = self.sessions.get(&msg.header.initiator_spi)?;
        info!(
            msg_id = msg.header.message_id,
            "INFORMATIONAL from peer - acking with empty response"
        );
        encrypt_empty_informational(
            msg.header.initiator_spi,
            session.responder_spi,
            msg.header.message_id,
            &session.keys,
        )
    }

    async fn handle_auth(&mut self, msg: &Message<'_>, peer: SocketAddr) -> Option<Vec<u8>> {
        if msg.header.flags.is_response() {
            return None;
        }
        let ispi = msg.header.initiator_spi;
        let Some(session) = self.sessions.get_mut(&ispi) else {
            warn!(
                ispi = format_args!("{ispi:016x}"),
                "no session for IKE_AUTH"
            );
            return None;
        };
        if session.responder_spi != msg.header.responder_spi {
            warn!("rSPI mismatch in IKE_AUTH");
            return None;
        }

        // Locate SK payload + decrypt.
        let sk_payload = msg
            .payloads
            .iter()
            .find(|p| p.kind == PayloadKind::Encrypted)?;
        let first_inner_kind = sk_payload.header.next_payload;

        let p = session.keys.suite.params();
        let enc = match EncryptedPayload::split(sk_payload.body, p.encr_iv_bytes, p.encr_icv_bytes)
        {
            Ok(e) => e,
            Err(e) => {
                warn!("SK split failed: {e}");
                return None;
            }
        };
        // For AEAD: AAD = IKE header + SK generic payload header (RFC 5282 §4).
        // For non-AEAD: HMAC over IKE header + SK generic header + IV + ciphertext.
        let sk_body_offset = msg
            .raw
            .len()
            .checked_sub(sk_payload.body.len())
            .expect("SK body lies within raw");
        let aad = &msg.raw[..sk_body_offset];

        let mut ct = enc.ciphertext.to_vec();
        if p.aead {
            if let Err(e) = aes_gcm_open(
                &session.keys.sk_ei,
                &session.keys.salt_ei,
                enc.iv,
                aad,
                &mut ct,
                enc.icv,
            ) {
                warn!("SK decrypt failed: {e}");
                return None;
            }
        } else {
            // RFC 7296 §3.14: HMAC over IKE_hdr || SK_hdr || IV || ciphertext.
            let mut to_mac = Vec::with_capacity(aad.len() + enc.iv.len() + ct.len());
            to_mac.extend_from_slice(aad);
            to_mac.extend_from_slice(enc.iv);
            to_mac.extend_from_slice(&ct);
            let expected_icv = crypto::hmac_sha256_128(&session.keys.sk_ai, &to_mac);
            use subtle::ConstantTimeEq;
            if !bool::from(expected_icv.ct_eq(enc.icv)) {
                warn!("SK HMAC verify failed");
                return None;
            }
            // Decrypt under AES-CBC-256.
            if enc.iv.len() != 16 || session.keys.sk_ei.len() != 32 {
                warn!("SK CBC: unexpected IV/key length");
                return None;
            }
            let mut iv16 = [0u8; 16];
            iv16.copy_from_slice(enc.iv);
            let mut key32 = [0u8; 32];
            key32.copy_from_slice(&session.keys.sk_ei);
            ct = match crypto::aes_cbc_256_decrypt(&key32, &iv16, &ct) {
                Ok(pt) => pt,
                Err(e) => {
                    warn!("SK CBC decrypt failed: {e}");
                    return None;
                }
            };
        }

        // Plaintext = inner payloads | pad | pad_length
        let pad_len = *ct.last()? as usize;
        if 1 + pad_len > ct.len() {
            warn!("bad pad length {pad_len}");
            return None;
        }
        ct.truncate(ct.len() - 1 - pad_len);

        debug!(
            "decrypted SK: pad_len={} actual={}B first_kind={:?}",
            pad_len,
            ct.len(),
            first_inner_kind
        );
        debug!(
            "plaintext hex: {}",
            ct.iter()
                .map(|b| format!("{b:02X}"))
                .collect::<Vec<_>>()
                .join(" ")
        );
        let inner = match walk_chain(first_inner_kind, &ct) {
            Ok(v) => v,
            Err(e) => {
                warn!("inner chain parse failed: {e}");
                return None;
            }
        };

        // Pull the payloads we need.
        let mut idi_body: Option<Vec<u8>> = None;
        let mut idr_body_from_peer: Option<Vec<u8>> = None;
        let mut auth_data: Option<Vec<u8>> = None;
        let mut auth_method: Option<AuthMethod> = None;
        let mut sai2_body: Option<Vec<u8>> = None;
        let mut tsi_body: Option<Vec<u8>> = None;
        let mut tsr_body: Option<Vec<u8>> = None;
        let mut saw_cfg_request = false;
        for p in &inner {
            match p.kind {
                PayloadKind::IdentificationInitiator => idi_body = Some(p.body.to_vec()),
                PayloadKind::IdentificationResponder => idr_body_from_peer = Some(p.body.to_vec()),
                PayloadKind::Authentication => {
                    if let Ok(a) = AuthPayload::parse(p.body) {
                        auth_method = Some(a.method);
                        auth_data = Some(a.data);
                    }
                }
                PayloadKind::SecurityAssociation => sai2_body = Some(p.body.to_vec()),
                PayloadKind::TrafficSelectorInitiator => tsi_body = Some(p.body.to_vec()),
                PayloadKind::TrafficSelectorResponder => tsr_body = Some(p.body.to_vec()),
                PayloadKind::Configuration => {
                    // First byte of body is CFG Type. 1 = CFG_REQUEST.
                    if let Some(&0x01) = p.body.first() {
                        saw_cfg_request = true;
                    }
                }
                PayloadKind::Notify => {
                    if let Ok(n) = NotifyPayload::parse(p.body) {
                        debug!(
                            "inner Notify: type={} ({:#x}) data={}B",
                            n.notify_type,
                            n.notify_type,
                            n.data.len()
                        );
                    }
                }
                _ => debug!("inner payload: {:?} ({} bytes)", p.kind, p.body.len()),
            }
        }
        let (idi_body, auth_data, auth_method, sai2_body, tsi_body, tsr_body) = match (
            idi_body,
            auth_data,
            auth_method,
            sai2_body,
            tsi_body,
            tsr_body,
        ) {
            (Some(a), Some(b), Some(c), Some(d), Some(e), Some(f)) => (a, b, c, d, e, f),
            _ => {
                warn!("IKE_AUTH missing required inner payload");
                return None;
            }
        };
        if auth_method != AuthMethod::SharedKeyMic {
            warn!(
                "auth method {:?} not supported (this MVP is PSK-only)",
                auth_method
            );
            return None;
        }
        debug!(
            "IKE_AUTH inner: IDi={:?} IDr?={} SAi2={}B TSi={}B TSr={}B",
            IdPayload::parse(&idi_body)
                .ok()
                .map(|id| String::from_utf8_lossy(&id.data).into_owned()),
            idr_body_from_peer.is_some(),
            sai2_body.len(),
            tsi_body.len(),
            tsr_body.len(),
        );

        // Build the challenge and hand it to the user callback. The callback
        // looks up a candidate auth_key, calls `challenge.verify(&key)` (or
        // `approve_with`), and returns an `AuthDecision`.
        let challenge = AuthChallenge {
            identity_raw: idi_body.clone(),
            peer,
            sk_pi: session.keys.sk_pi.clone(),
            init_request: session.init_request.clone(),
            nonce_r: session.nonce_r.clone(),
            auth_data: auth_data.clone(),
        };
        let id_display = challenge.identity_str();
        let auth_key = match (self.config.auth)(challenge).await {
            AuthDecision::Approve { auth_key } => auth_key,
            AuthDecision::Reject => {
                warn!(id = %id_display, "auth callback rejected");
                return None;
            }
        };
        info!(
            ispi = format_args!("{ispi:016x}"),
            "AUTH verified - PSK match"
        );

        // Parse SAi2, select an ESP proposal.
        let sai2 = SaPayload::parse(&sai2_body).ok()?;
        let (esp_chosen, peer_esp_spi) = match select_esp_proposal(&sai2, session.keys.suite) {
            Some(v) => v,
            None => {
                warn!("no acceptable ESP proposal in SAi2");
                return None;
            }
        };

        // Generate our ESP SPI.
        let our_esp_spi: u32 = rand::rng().random();

        // Derive CHILD_SA key material. Use the same suite negotiated for IKE
        // - iOS always sends matching ENCR proposals for both legs.
        let child = ChildSa::derive(
            session.keys.suite,
            &session.keys.sk_d,
            &session.nonce_i,
            &session.nonce_r,
            our_esp_spi,
            peer_esp_spi,
        );
        debug!(
            inbound_spi = format_args!("{:08x}", our_esp_spi),
            outbound_spi = format_args!("{peer_esp_spi:08x}"),
            "derived CHILD_SA keys"
        );

        // IDr: if iOS told us what identity it expects (IDr in its request),
        // echo that back. Profiles with a "Remote ID" configured send this and
        // enforce a match. Otherwise fall back to our configured identity.
        let our_idr_body = match idr_body_from_peer {
            Some(b) => b,
            None => {
                let id = IdPayload {
                    id_type: IdType::Rfc822Addr,
                    data: self.config.our_identity.clone().into_bytes(),
                };
                let mut b = Vec::new();
                id.write_body(&mut b);
                b
            }
        };

        // Compute our AUTH using the same auth key.
        let our_auth_mac = compute_psk_auth_responder(session, &auth_key, &our_idr_body);

        // Build inner payloads for response: IDr, AUTH, SAr2, TSi, TSr
        let auth_body = {
            let a = AuthPayload {
                method: AuthMethod::SharedKeyMic,
                data: our_auth_mac.to_vec(),
            };
            let mut b = Vec::new();
            a.write_body(&mut b);
            b
        };
        let sar2_body = build_chosen_esp_sa_body(&esp_chosen, our_esp_spi);

        let mut inner_payloads = vec![
            OutPayload {
                kind: PayloadKind::IdentificationResponder,
                body: our_idr_body,
            },
            OutPayload {
                kind: PayloadKind::Authentication,
                body: auth_body,
            },
        ];
        // CP(CFG_REPLY) - mandatory when the client asked for a virtual IP.
        if saw_cfg_request {
            inner_payloads.push(OutPayload {
                kind: PayloadKind::Configuration,
                body: build_cfg_reply(self.config.virtual_ip, self.config.virtual_dns),
            });
        }
        // Narrow TSi to iOS' virtual IP (the one we assigned via CP) and TSr
        // to our gateway IP. This is what makes iOS install a /32 route for
        // the gateway - it only tunnels traffic for that destination, leaving
        // the rest of its traffic on the default route.
        //
        // NB: RFC 7296 §2.9 requires the responder's narrowed selector to be
        // a subset of what the initiator offered. iOS offered the whole v4
        // address space, so a /32 is trivially a subset.
        let _ = tsi_body;
        let _ = tsr_body;
        let narrowed_tsi = ts_single_ipv4(self.config.virtual_ip);
        let narrowed_tsr = ts_single_ipv4(self.config.gateway_ip);

        inner_payloads.extend([
            OutPayload {
                kind: PayloadKind::SecurityAssociation,
                body: sar2_body,
            },
            OutPayload {
                kind: PayloadKind::TrafficSelectorInitiator,
                body: narrowed_tsi,
            },
            OutPayload {
                kind: PayloadKind::TrafficSelectorResponder,
                body: narrowed_tsr,
            },
        ]);

        let response = encrypt_ike_message(
            ispi,
            session.responder_spi,
            ExchangeType::IkeAuth,
            Flags(Flags::RESPONSE),
            msg.header.message_id,
            &inner_payloads,
            &session.keys,
        )?;

        session.state = State::Authenticated;
        session.peer = peer;

        // Set up the EspTunnel + channels, wire the inbound side into the
        // session, and emit the EstablishedSession for the caller.
        let (tunnel, inbound_tx, outbound_rx) = EspTunnel::channels();
        session.inbound_tx = Some(inbound_tx);
        let established = EstablishedSession {
            peer,
            initiator_spi: ispi,
            peer_identity: idi_body.clone(),
            virtual_ip: self.config.virtual_ip,
            gateway_ip: self.config.gateway_ip,
            tunnel,
            outbound_rx,
            suite: child.suite,
            outbound_spi: child.outbound_spi,
            outbound_key: child.outbound_key.clone(),
            outbound_salt: child.outbound_salt.clone(),
            outbound_integ: child.outbound_integ.clone(),
        };
        session.child = Some(child);

        // Non-fatal if no receiver - just means the caller discarded the
        // channel, but IKE still works.
        {
            use crossfire::BlockingTxTrait;
            if self.new_sessions.send(established).is_err() {
                warn!("new-session channel dropped; tunnel unused");
            }
        }

        info!(%peer, "IKE_AUTH complete - CHILD_SA up");
        Some(response)
    }

    fn handle_sa_init(&mut self, msg: &Message<'_>, peer: SocketAddr) -> Option<Vec<u8>> {
        if msg.header.flags.is_response() {
            warn!("IKE_SA_INIT marked as response - we're the responder, ignoring");
            return None;
        }
        if msg.header.responder_spi != 0 {
            warn!("IKE_SA_INIT with non-zero responder SPI - ignoring");
            return None;
        }

        // Pull out the three payloads we need, and note which feature notifies
        // iOS advertised so we can echo the ones we support.
        let mut sa_body = None;
        let mut ke = None;
        let mut nonce_i = None;
        let mut peer_wants_fragmentation = false;
        for p in &msg.payloads {
            match p.kind {
                PayloadKind::SecurityAssociation => sa_body = Some(p.body),
                PayloadKind::KeyExchange => {
                    ke = match KePayload::parse(p.body) {
                        Ok(k) => Some(k),
                        Err(e) => {
                            warn!("KE parse failed: {e}");
                            return None;
                        }
                    }
                }
                PayloadKind::Nonce => {
                    nonce_i = match NoncePayload::parse(p.body) {
                        Ok(n) => Some(n.0),
                        Err(e) => {
                            warn!("Ni parse failed: {e}");
                            return None;
                        }
                    }
                }
                PayloadKind::Notify => {
                    if let Ok(n) = NotifyPayload::parse(p.body)
                        && n.notify_type == notify_type::FRAGMENTATION_SUPPORTED
                    {
                        peer_wants_fragmentation = true;
                    }
                }
                _ => {}
            }
        }
        let (sa_body, ke, nonce_i) = match (sa_body, ke, nonce_i) {
            (Some(a), Some(b), Some(c)) => (a, b, c),
            _ => {
                warn!("IKE_SA_INIT missing SA, KE, or Ni");
                return None;
            }
        };

        let sa = match SaPayload::parse(sa_body) {
            Ok(s) => s,
            Err(e) => {
                warn!("SA parse failed: {e}");
                return None;
            }
        };

        // Pick the first proposal we accept (AES-GCM-256 + PRF SHA-256 + DH 19).
        let chosen = match select_proposal(&sa) {
            Some(c) => c,
            None => {
                warn!(
                    "no acceptable proposal - peer offered {} proposal(s):",
                    sa.proposals.len()
                );
                for prop in &sa.proposals {
                    let xforms: Vec<String> = prop
                        .transforms
                        .iter()
                        .map(|t| {
                            format!(
                                "{:?}={}{}",
                                t.transform_type,
                                t.transform_id,
                                t.key_length().map(|k| format!("/{k}")).unwrap_or_default()
                            )
                        })
                        .collect();
                    warn!(
                        "  proposal#{} proto={:?}: {}",
                        prop.proposal_num,
                        prop.protocol,
                        xforms.join(", ")
                    );
                }
                return Some(no_proposal_chosen(msg.header.initiator_spi, 0));
            }
        };
        debug!(?chosen, "selected proposal");

        // RFC 7296 §1.2: the KE payload carries the DH share for ONE group.
        // If the initiator's chosen group doesn't match the proposal we just
        // picked, send INVALID_KE_PAYLOAD with our preferred group. The
        // initiator will retry IKE_SA_INIT with a fresh KE for that group.
        if ke.dh_group != chosen.dh_id {
            info!(
                got = ke.dh_group,
                want = chosen.dh_id,
                "DH group in KE payload doesn't match the chosen proposal - sending INVALID_KE_PAYLOAD"
            );
            return Some(invalid_ke_payload(msg.header.initiator_spi, chosen.dh_id));
        }

        // Generate our half: SPI, DH key pair, nonce.
        let mut rng = rand::rng();
        let responder_spi: u64 = rng.random();
        let dh = DhEphemeral::generate();
        let our_share = dh.public_share();

        let mut nonce_r = vec![0u8; 32];
        rng.fill_bytes(&mut nonce_r);

        // ECDH.
        let g_ir = match dh.diffie_hellman(&ke.key_data) {
            Ok(s) => s,
            Err(e) => {
                warn!("ECDH failed: {e}");
                return None;
            }
        };

        // SKEYSEED = prf(Ni | Nr, g^ir)  (RFC 7296 §2.14)
        let mut prf_key = Vec::with_capacity(nonce_i.len() + nonce_r.len());
        prf_key.extend_from_slice(&nonce_i);
        prf_key.extend_from_slice(&nonce_r);
        let skeyseed = crypto::prf(&prf_key, &g_ir);

        // prf+(SKEYSEED, Ni | Nr | SPIi | SPIr) -> keymat
        let mut seed = Vec::new();
        seed.extend_from_slice(&nonce_i);
        seed.extend_from_slice(&nonce_r);
        seed.extend_from_slice(&msg.header.initiator_spi.to_be_bytes());
        seed.extend_from_slice(&responder_spi.to_be_bytes());
        let suite_params = chosen.suite.params();
        let keymat = crypto::prf_plus(&skeyseed, &seed, suite_params.keymat_len());

        let keys = IkeKeys::split(chosen.suite, &keymat);
        debug!(suite = ?chosen.suite, "derived IKE keys");

        // Build the response payloads.
        let chosen_sa_body = build_chosen_sa_body(&chosen);
        let ke_body = {
            let ke_out = KePayload {
                dh_group: 19,
                key_data: our_share.to_vec(),
            };
            let mut b = Vec::with_capacity(ke_out.encoded_len());
            ke_out.write_body(&mut b);
            b
        };
        let nr_body = nonce_r.clone();

        let nat_src = nat_detection_hash(
            msg.header.initiator_spi,
            responder_spi,
            self.config.local_ip,
            self.config.local_port,
        );
        let nat_dst = nat_detection_hash(
            msg.header.initiator_spi,
            responder_spi,
            peer.ip(),
            peer.port(),
        );

        let mut payloads = vec![
            OutPayload {
                kind: PayloadKind::SecurityAssociation,
                body: chosen_sa_body,
            },
            OutPayload {
                kind: PayloadKind::KeyExchange,
                body: ke_body,
            },
            OutPayload {
                kind: PayloadKind::Nonce,
                body: nr_body,
            },
            OutPayload {
                kind: PayloadKind::Notify,
                body: notify_body(notify_type::NAT_DETECTION_SOURCE_IP, &nat_src),
            },
            OutPayload {
                kind: PayloadKind::Notify,
                body: notify_body(notify_type::NAT_DETECTION_DESTINATION_IP, &nat_dst),
            },
        ];
        if peer_wants_fragmentation {
            payloads.push(OutPayload {
                kind: PayloadKind::Notify,
                body: notify_body(notify_type::FRAGMENTATION_SUPPORTED, &[]),
            });
        }

        let response = ike::build_message(
            msg.header.initiator_spi,
            responder_spi,
            ExchangeType::IkeSaInit,
            Flags(Flags::RESPONSE),
            0,
            &payloads,
        );

        // Persist session state for IKE_AUTH.
        let session = Session {
            initiator_spi: msg.header.initiator_spi,
            responder_spi,
            peer,
            state: State::SaInitDone,
            keys,
            nonce_i,
            nonce_r,
            init_request: msg.raw.to_vec(),
            init_response: response.clone(),
            child: None,
            inbound_tx: None,
        };
        self.sessions.insert(session.initiator_spi, session);

        info!(%peer, ispi = format_args!("{:016x}", msg.header.initiator_spi),
              rspi = format_args!("{:016x}", responder_spi),
              "IKE_SA_INIT complete");
        Some(response)
    }
}

// ------------------------------------------------------------- per-SA state

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum State {
    SaInitDone,
    Authenticated,
}

pub(crate) struct Session {
    pub initiator_spi: u64,
    pub responder_spi: u64,
    pub peer: SocketAddr,
    pub state: State,
    pub keys: IkeKeys,
    pub nonce_i: Vec<u8>,
    pub nonce_r: Vec<u8>,
    pub init_request: Vec<u8>,
    pub init_response: Vec<u8>,
    pub child: Option<ChildSa>,
    /// Delivers decrypted inbound IP packets to the [`EspTunnel`]. `None`
    /// until IKE_AUTH completes.
    pub inbound_tx: Option<crossfire::MTx<InboundFlavor>>,
}

/// ESP CHILD_SA: SPIs for each direction plus per-direction key material
/// (encryption key + optional GCM salt + optional integrity key).
pub(crate) struct ChildSa {
    pub suite: crypto::Suite,
    pub inbound_spi: u32,
    pub outbound_spi: u32,
    /// Inbound AES key (32 bytes for AES-256).
    pub inbound_key: Vec<u8>,
    /// 4-byte GCM salt for AEAD; empty for AES-CBC.
    pub inbound_salt: Vec<u8>,
    /// Inbound HMAC integrity key. Empty for AEAD.
    pub inbound_integ: Vec<u8>,
    pub outbound_key: Vec<u8>,
    pub outbound_salt: Vec<u8>,
    pub outbound_integ: Vec<u8>,
}

impl ChildSa {
    /// Derive `KEYMAT = prf+(SK_d, Ni | Nr)` and split into the per-direction
    /// blobs (RFC 7296 §2.17). Order per direction: encr_key | salt | integ_key
    /// - for AEAD `integ_key` is empty; for AES-CBC `salt` is empty.
    fn derive(
        suite: crypto::Suite,
        sk_d: &[u8],
        ni: &[u8],
        nr: &[u8],
        inbound_spi: u32,
        outbound_spi: u32,
    ) -> Self {
        let p = suite.params();
        let per_dir = p.child_per_dir_len();
        let mut seed = Vec::with_capacity(ni.len() + nr.len());
        seed.extend_from_slice(ni);
        seed.extend_from_slice(nr);
        let km = crypto::prf_plus(sk_d, &seed, per_dir * 2);
        let split_dir = |off: usize| -> (Vec<u8>, Vec<u8>, Vec<u8>) {
            let key = km[off..off + p.encr_key_bytes].to_vec();
            let salt =
                km[off + p.encr_key_bytes..off + p.encr_key_bytes + p.encr_salt_bytes].to_vec();
            let integ = km[off + p.sk_e_len()..off + per_dir].to_vec();
            (key, salt, integ)
        };
        let (inbound_key, inbound_salt, inbound_integ) = split_dir(0);
        let (outbound_key, outbound_salt, outbound_integ) = split_dir(per_dir);
        Self {
            suite,
            inbound_spi,
            outbound_spi,
            inbound_key,
            inbound_salt,
            inbound_integ,
            outbound_key,
            outbound_salt,
            outbound_integ,
        }
    }
}

/// IKE-SA key material derived from `prf+(SKEYSEED, Ni | Nr | SPIi | SPIr)`.
/// Order per RFC 7296 §2.14: SK_d, SK_ai, SK_ar, SK_ei, SK_er, SK_pi, SK_pr.
///
/// For AEAD ciphers (AES-GCM): SK_ai/SK_ar are empty and SK_ei/SK_er carry
/// the AES key followed by the GCM salt.
/// For AES-CBC + HMAC: SK_ai/SK_ar are HMAC keys and SK_ei/SK_er are pure
/// AES keys (no salt).
pub(crate) struct IkeKeys {
    pub suite: crypto::Suite,
    pub sk_d: Vec<u8>,
    pub sk_ai: Vec<u8>,
    pub sk_ar: Vec<u8>,
    pub sk_ei: Vec<u8>,
    pub salt_ei: Vec<u8>,
    pub sk_er: Vec<u8>,
    pub salt_er: Vec<u8>,
    pub sk_pi: Vec<u8>,
    pub sk_pr: Vec<u8>,
}

impl IkeKeys {
    fn split(suite: crypto::Suite, km: &[u8]) -> Self {
        let p = suite.params();
        debug_assert_eq!(km.len(), p.keymat_len());
        let mut cur = 0;
        let mut take = |n: usize| -> Vec<u8> {
            let s = km[cur..cur + n].to_vec();
            cur += n;
            s
        };
        let sk_d = take(p.prf_bytes);
        let sk_ai = take(p.integ_key_bytes);
        let sk_ar = take(p.integ_key_bytes);
        let sk_ei = take(p.encr_key_bytes);
        let salt_ei = take(p.encr_salt_bytes);
        let sk_er = take(p.encr_key_bytes);
        let salt_er = take(p.encr_salt_bytes);
        let sk_pi = take(p.prf_bytes);
        let sk_pr = take(p.prf_bytes);
        debug_assert_eq!(cur, p.keymat_len());
        Self {
            suite,
            sk_d,
            sk_ai,
            sk_ar,
            sk_ei,
            salt_ei,
            sk_er,
            salt_er,
            sk_pi,
            sk_pr,
        }
    }
}

// ------------------------------------------------------------ proposal logic

#[derive(Debug, Clone)]
pub(crate) struct ChosenSuite {
    pub suite: crypto::Suite,
    pub proposal_num: u8,
    pub encr_id: u16,
    pub encr_keylen: u16,
    pub prf_id: u16,
    pub integ_id: Option<u16>,
    pub dh_id: u16,
}

/// Walk peer's IKE proposals and pick the first one matching any of our
/// supported suites (in our preference order: GCM > CBC).
fn select_proposal(sa: &SaPayload) -> Option<ChosenSuite> {
    for our_suite in crypto::Suite::supported() {
        let p = our_suite.params();
        for prop in &sa.proposals {
            if prop.protocol != ProtocolId::Ike {
                continue;
            }
            let encr = prop.transforms.iter().find(|t| {
                t.transform_type == TransformType::Encr
                    && t.transform_id == p.encr_id
                    && t.key_length() == Some(p.encr_keylen_bits)
            });
            let prf = prop
                .transforms
                .iter()
                .find(|t| t.transform_type == TransformType::Prf && t.transform_id == 5);
            let dh = prop
                .transforms
                .iter()
                .find(|t| t.transform_type == TransformType::Dh && t.transform_id == 19);
            // For non-AEAD suites the proposal must also offer the matching
            // INTEG transform. AEAD requires either no INTEG or INTEG=NONE (0).
            let integ_ok = match p.integ_id {
                Some(id) => prop
                    .transforms
                    .iter()
                    .any(|t| t.transform_type == TransformType::Integ && t.transform_id == id),
                None => true, // AEAD - peers always omit INTEG entirely (or NONE)
            };
            if let (Some(e), Some(prf_t), Some(d)) = (encr, prf, dh)
                && integ_ok
            {
                return Some(ChosenSuite {
                    suite: *our_suite,
                    proposal_num: prop.proposal_num,
                    encr_id: e.transform_id,
                    encr_keylen: e.key_length().unwrap_or(0),
                    prf_id: prf_t.transform_id,
                    integ_id: p.integ_id,
                    dh_id: d.transform_id,
                });
            }
        }
    }
    None
}

/// Build the SA payload body containing only the chosen proposal (a single
/// proposal with three transforms: ENCR, PRF, DH).
fn build_chosen_sa_body(c: &ChosenSuite) -> Vec<u8> {
    let mut transforms = vec![
        Transform {
            transform_type: TransformType::Encr,
            transform_id: c.encr_id,
            attributes: vec![Attribute {
                attr_type: Attribute::KEY_LENGTH,
                value: AttrValue::Tv(c.encr_keylen),
            }],
        },
        Transform {
            transform_type: TransformType::Prf,
            transform_id: c.prf_id,
            attributes: Vec::new(),
        },
    ];
    // Non-AEAD suites require an INTEG transform; AEAD suites omit it.
    if let Some(integ_id) = c.integ_id {
        transforms.push(Transform {
            transform_type: TransformType::Integ,
            transform_id: integ_id,
            attributes: Vec::new(),
        });
    }
    transforms.push(Transform {
        transform_type: TransformType::Dh,
        transform_id: c.dh_id,
        attributes: Vec::new(),
    });
    let prop = Proposal {
        proposal_num: c.proposal_num,
        protocol: ProtocolId::Ike,
        spi: Vec::new(),
        transforms,
    };
    let sa = SaPayload {
        proposals: vec![prop],
    };
    let mut buf = Vec::with_capacity(sa.encoded_len());
    sa.write_body(&mut buf);
    buf
}

// ------------------------------------------------------------ helper builders

fn notify_body(notify_type: u16, data: &[u8]) -> Vec<u8> {
    let n = NotifyPayload {
        protocol_id: 0,
        notify_type,
        spi: Vec::new(),
        data: data.to_vec(),
    };
    let mut b = Vec::with_capacity(n.encoded_len());
    n.write_body(&mut b);
    b
}

// ----------------------------------------------------------------- IKE_AUTH

/// Compute our AUTH payload for the responder side of PSK mode.
/// `auth_key` is the same derived key as above.
fn compute_psk_auth_responder(session: &Session, auth_key: &[u8; 32], idr_body: &[u8]) -> [u8; 32] {
    let maced_id = crypto::prf(&session.keys.sk_pr, idr_body);
    let mut signed =
        Vec::with_capacity(session.init_response.len() + session.nonce_i.len() + maced_id.len());
    signed.extend_from_slice(&session.init_response);
    signed.extend_from_slice(&session.nonce_i);
    signed.extend_from_slice(&maced_id);
    crypto::prf(auth_key, &signed)
}

/// Walk peer's ESP proposals (SAi2) and pick the first matching any suite we
/// support. The chosen suite is constrained to match the IKE suite - iOS
/// always offers compatible options for both legs.
fn select_esp_proposal(sa: &SaPayload, ike_suite: crypto::Suite) -> Option<(ChosenEspSuite, u32)> {
    let p = ike_suite.params();
    for prop in &sa.proposals {
        if prop.protocol != ProtocolId::Esp || prop.spi.len() != 4 {
            continue;
        }
        let encr = prop.transforms.iter().find(|t| {
            t.transform_type == TransformType::Encr
                && t.transform_id == p.encr_id
                && t.key_length() == Some(p.encr_keylen_bits)
        });
        let integ_ok = match p.integ_id {
            Some(id) => prop
                .transforms
                .iter()
                .any(|t| t.transform_type == TransformType::Integ && t.transform_id == id),
            None => true,
        };
        // ESN must be present and set to NONE (0).
        let esn = prop
            .transforms
            .iter()
            .find(|t| t.transform_type == TransformType::Esn && t.transform_id == 0);
        if let (Some(e), Some(_esn)) = (encr, esn)
            && integ_ok
        {
            let mut spi_bytes = [0u8; 4];
            spi_bytes.copy_from_slice(&prop.spi);
            return Some((
                ChosenEspSuite {
                    suite: ike_suite,
                    proposal_num: prop.proposal_num,
                    encr_id: e.transform_id,
                    encr_keylen: e.key_length().unwrap_or(0),
                    integ_id: p.integ_id,
                },
                u32::from_be_bytes(spi_bytes),
            ));
        }
    }
    None
}

#[derive(Debug, Clone)]
pub(crate) struct ChosenEspSuite {
    /// Same suite as the IKE-SA - kept for symmetry with `ChosenSuite`.
    #[allow(dead_code)]
    pub suite: crypto::Suite,
    pub proposal_num: u8,
    pub encr_id: u16,
    pub encr_keylen: u16,
    pub integ_id: Option<u16>,
}

fn build_chosen_esp_sa_body(c: &ChosenEspSuite, our_spi: u32) -> Vec<u8> {
    let mut transforms = vec![Transform {
        transform_type: TransformType::Encr,
        transform_id: c.encr_id,
        attributes: vec![Attribute {
            attr_type: Attribute::KEY_LENGTH,
            value: AttrValue::Tv(c.encr_keylen),
        }],
    }];
    if let Some(integ_id) = c.integ_id {
        transforms.push(Transform {
            transform_type: TransformType::Integ,
            transform_id: integ_id,
            attributes: Vec::new(),
        });
    }
    transforms.push(Transform {
        transform_type: TransformType::Esn,
        transform_id: 0,
        attributes: Vec::new(),
    });
    let prop = Proposal {
        proposal_num: c.proposal_num,
        protocol: ProtocolId::Esp,
        spi: our_spi.to_be_bytes().to_vec(),
        transforms,
    };
    let sa = SaPayload {
        proposals: vec![prop],
    };
    let mut buf = Vec::with_capacity(sa.encoded_len());
    sa.write_body(&mut buf);
    buf
}

/// Encrypt an inner payload chain as an SK payload and wrap it in the outer
/// IKE header. Generic helper used for IKE_AUTH responses, INFORMATIONAL
/// messages, and DELETE notifications.
fn encrypt_ike_message(
    initiator_spi: u64,
    responder_spi: u64,
    exchange_type: ExchangeType,
    flags: Flags,
    message_id: u32,
    inner: &[OutPayload],
    keys: &IkeKeys,
) -> Option<Vec<u8>> {
    let p = keys.suite.params();
    // Build the inner payload chain, then add padding so the trailer (pad +
    // pad_length byte) reaches the cipher's block size: 1 byte for AEAD
    // (no alignment requirement) and 16 bytes for AES-CBC.
    let block = if p.aead { 1 } else { 16 };
    let mut plaintext = Vec::new();
    for (i, op) in inner.iter().enumerate() {
        let next = inner
            .get(i + 1)
            .map(|n| n.kind)
            .unwrap_or(PayloadKind::None);
        let plen = (4 + op.body.len()) as u16;
        plaintext.push(next.as_u8());
        plaintext.push(0);
        plaintext.extend_from_slice(&plen.to_be_bytes());
        plaintext.extend_from_slice(&op.body);
    }
    let unaligned = plaintext.len() + 1; // +pad_length byte
    let pad = (block - (unaligned % block)) % block;
    for i in 1..=pad {
        plaintext.push(i as u8);
    }
    plaintext.push(pad as u8);

    let first_inner_kind = inner.first().map(|p| p.kind).unwrap_or(PayloadKind::None);

    let mut iv = vec![0u8; p.encr_iv_bytes];
    rand::Rng::fill_bytes(&mut rand::rng(), &mut iv);

    // For AEAD, ciphertext length = plaintext length. For AES-CBC same (no
    // chained block expansion).
    let sk_body_len = iv.len() + plaintext.len() + p.encr_icv_bytes;
    let sk_payload_len = 4 + sk_body_len;
    let total_len = ike::HEADER_LEN + sk_payload_len;

    let mut out = Vec::with_capacity(total_len);
    let header = crate::ike::Header {
        initiator_spi,
        responder_spi,
        next_payload: PayloadKind::Encrypted,
        version: ike::IKE_VERSION,
        exchange_type,
        flags,
        message_id,
        length: total_len as u32,
    };
    let mut hbuf = [0u8; ike::HEADER_LEN];
    header.write_into(&mut hbuf);
    out.extend_from_slice(&hbuf);
    out.push(first_inner_kind.as_u8());
    out.push(0);
    out.extend_from_slice(&(sk_payload_len as u16).to_be_bytes());

    if p.aead {
        let aad = out.clone();
        out.extend_from_slice(&iv);
        let mut ct = plaintext;
        let tag = match aes_gcm_seal(&keys.sk_er, &keys.salt_er, &iv, &aad, &mut ct) {
            Ok(t) => t,
            Err(e) => {
                warn!("SK encrypt failed: {e}");
                return None;
            }
        };
        out.extend_from_slice(&ct);
        out.extend_from_slice(&tag);
    } else {
        // AES-CBC + HMAC-SHA-256-128
        let mut iv16 = [0u8; 16];
        iv16.copy_from_slice(&iv);
        let mut key32 = [0u8; 32];
        if keys.sk_er.len() != 32 {
            warn!("SK CBC: bad key length");
            return None;
        }
        key32.copy_from_slice(&keys.sk_er);
        let ct = match crypto::aes_cbc_256_encrypt(&key32, &iv16, &plaintext) {
            Ok(c) => c,
            Err(e) => {
                warn!("SK CBC encrypt failed: {e}");
                return None;
            }
        };
        out.extend_from_slice(&iv);
        out.extend_from_slice(&ct);
        // HMAC over IKE_hdr || SK_hdr || IV || ciphertext.
        let icv = crypto::hmac_sha256_128(&keys.sk_ar, &out);
        out.extend_from_slice(&icv);
    }

    debug_assert_eq!(out.len(), total_len);
    Some(out)
}

/// Build an encrypted INFORMATIONAL **request** carrying a Delete payload
/// that tears down the IKE SA. Per RFC 7296 §1.4, deleting the IKE SA
/// implicitly deletes all CHILD SAs. iOS will ACK and disconnect.
///
/// `message_id` must be the next unused ID for the responder direction. For
/// our MVP this is always 0 (the responder hasn't sent any requests yet).
fn build_delete_ike_sa(
    initiator_spi: u64,
    responder_spi: u64,
    message_id: u32,
    keys: &IkeKeys,
) -> Option<Vec<u8>> {
    // Delete payload for IKE SA: protocol_id=1 (IKE), spi_size=0, num_spis=0.
    let delete_body = {
        let mut b = Vec::new();
        b.push(1); // protocol_id = IKE
        b.push(0); // spi_size = 0 (IKE SA SPIs are in the header)
        b.extend_from_slice(&0u16.to_be_bytes()); // num_spis = 0
        b
    };

    let inner = &[OutPayload {
        kind: PayloadKind::Delete,
        body: delete_body,
    }];

    encrypt_ike_message(
        initiator_spi,
        responder_spi,
        ExchangeType::Informational,
        Flags(0), // request from responder: no INITIATOR, no RESPONSE
        message_id,
        inner,
        keys,
    )
}

/// Encrypt an empty INFORMATIONAL response (no inner payloads). Used to ack
/// DELETE notifies / keep-alives without having to parse them.
fn encrypt_empty_informational(
    initiator_spi: u64,
    responder_spi: u64,
    message_id: u32,
    keys: &IkeKeys,
) -> Option<Vec<u8>> {
    encrypt_ike_message(
        initiator_spi,
        responder_spi,
        ExchangeType::Informational,
        Flags(Flags::RESPONSE),
        message_id,
        &[], // no inner payloads
        keys,
    )
}

/// Build a Traffic Selector payload body covering exactly one IPv4 address
/// (a /32 single-host range), any port, any protocol. This is what iOS
/// interprets as "only tunnel traffic to this destination".
///
/// Layout:
/// ```text
/// | num(1) | RESERVED(3) | TS { ts_type=7 | proto=0 | length=16
///                                | start_port=0 | end_port=0xFFFF
///                                | start_addr | end_addr } |
/// ```
fn ts_single_ipv4(addr: Ipv4Addr) -> Vec<u8> {
    let mut sel = Vec::with_capacity(16);
    sel.push(7); // TS_IPV4_ADDR_RANGE
    sel.push(0); // IP protocol = ANY
    sel.extend_from_slice(&16u16.to_be_bytes()); // length
    sel.extend_from_slice(&0u16.to_be_bytes()); // start_port
    sel.extend_from_slice(&0xFFFFu16.to_be_bytes()); // end_port
    sel.extend_from_slice(&addr.octets()); // start addr
    sel.extend_from_slice(&addr.octets()); // end addr (same = /32)

    let mut body = Vec::with_capacity(4 + sel.len());
    body.push(1); // num selectors
    body.extend_from_slice(&[0; 3]); // reserved
    body.extend_from_slice(&sel);
    body
}

/// Minimal CP(CFG_REPLY) body: assign an IPv4 address, /32 netmask, and
/// optionally a DNS server. RFC 7296 §3.15.
///
/// Body layout:
/// ```text
/// | CFG Type (1) | RESERVED (3) | Configuration Attributes... |
/// ```
/// Each attribute: `| R | 15-bit type | 16-bit length | value |`.
fn build_cfg_reply(ip: Ipv4Addr, dns: Ipv4Addr) -> Vec<u8> {
    const CFG_REPLY: u8 = 2;
    const INTERNAL_IP4_ADDRESS: u16 = 1;
    const INTERNAL_IP4_NETMASK: u16 = 2;
    const INTERNAL_IP4_DNS: u16 = 3;

    let mut body = Vec::new();
    body.push(CFG_REPLY);
    body.extend_from_slice(&[0, 0, 0]); // reserved

    // Helper to append a TLV attribute.
    let mut push_attr = |ty: u16, value: &[u8]| {
        body.extend_from_slice(&ty.to_be_bytes());
        body.extend_from_slice(&(value.len() as u16).to_be_bytes());
        body.extend_from_slice(value);
    };

    push_attr(INTERNAL_IP4_ADDRESS, &ip.octets());
    // /32 - point-to-point link, routing is up to the client.
    push_attr(INTERNAL_IP4_NETMASK, &[0xff, 0xff, 0xff, 0xff]);
    if dns.octets() != [0, 0, 0, 0] {
        push_attr(INTERNAL_IP4_DNS, &dns.octets());
    }
    body
}

/// SHA-1(SPI_i || SPI_r || IP || port) per RFC 3947 §3.2.
fn nat_detection_hash(spi_i: u64, spi_r: u64, ip: IpAddr, port: u16) -> [u8; 20] {
    let mut h = Sha1::new();
    h.update(spi_i.to_be_bytes());
    h.update(spi_r.to_be_bytes());
    match ip {
        IpAddr::V4(v4) => h.update(v4.octets()),
        IpAddr::V6(v6) => h.update(v6.octets()),
    }
    h.update(port.to_be_bytes());
    h.finalize().into()
}

/// Build a NO_PROPOSAL_CHOSEN response (header + Notify payload, unencrypted
/// because no SA exists yet).
fn no_proposal_chosen(initiator_spi: u64, message_id: u32) -> Vec<u8> {
    let body = notify_body(notify_type::NO_PROPOSAL_CHOSEN, &[]);
    ike::build_message(
        initiator_spi,
        0,
        ExchangeType::IkeSaInit,
        Flags(Flags::RESPONSE),
        message_id,
        &[OutPayload {
            kind: PayloadKind::Notify,
            body,
        }],
    )
}

/// Build an INVALID_KE_PAYLOAD response (RFC 7296 §1.2 / §3.10.1). The Notify
/// data is the 2-byte big-endian DH group ID we want the initiator to use.
fn invalid_ke_payload(initiator_spi: u64, want_dh_group: u16) -> Vec<u8> {
    let body = notify_body(
        notify_type::INVALID_KE_PAYLOAD,
        &want_dh_group.to_be_bytes(),
    );
    ike::build_message(
        initiator_spi,
        0,
        ExchangeType::IkeSaInit,
        Flags(Flags::RESPONSE),
        0,
        &[OutPayload {
            kind: PayloadKind::Notify,
            body,
        }],
    )
}

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

    /// End-to-end: feed the captured iOS IKE_SA_INIT into the server, verify
    /// we produce a parseable response with the expected payload chain.
    #[tokio::test]
    async fn ios_init_round_trip() {
        let config = Config {
            auth: Box::new(|challenge| {
                Box::pin(async move {
                    challenge.approve_with(&crate::crypto::derive_auth_key(b"secret"))
                })
            }),
            local_ip: "192.0.2.1".parse().unwrap(),
            local_port: 500,
            our_identity: "vpn@local".to_string(),
            virtual_ip: "10.8.0.2".parse().unwrap(),
            virtual_dns: "1.1.1.1".parse().unwrap(),
            gateway_ip: "10.8.0.1".parse().unwrap(),
        };
        let (mut server, _inbound_tx, _new_sessions) = Server::new(config);

        let datagram = ios_init_bytes();
        let peer: SocketAddr = "203.0.113.5:500".parse().unwrap();
        let reply = server.handle(&datagram, peer).await.expect("reply");

        // Parse the reply and inspect.
        let parsed = Message::parse(&reply).expect("parse");
        assert_eq!(parsed.header.exchange_type, ExchangeType::IkeSaInit);
        assert!(parsed.header.flags.is_response());
        assert_eq!(parsed.header.message_id, 0);
        assert_eq!(parsed.header.initiator_spi, 0x36195e76eb4b9a11);
        assert_ne!(parsed.header.responder_spi, 0);
        let kinds: Vec<_> = parsed.payloads.iter().map(|p| p.kind).collect();
        assert_eq!(
            kinds,
            vec![
                PayloadKind::SecurityAssociation,
                PayloadKind::KeyExchange,
                PayloadKind::Nonce,
                PayloadKind::Notify, // NAT_DETECTION_SOURCE_IP
                PayloadKind::Notify, // NAT_DETECTION_DESTINATION_IP
                PayloadKind::Notify, // FRAGMENTATION_SUPPORTED (echoed)
            ]
        );

        // Session must be saved.
        assert!(server.sessions.contains_key(&0x36195e76eb4b9a11));
    }

    fn ios_init_bytes() -> Vec<u8> {
        const HEX: &str = "
            36 19 5E 76 EB 4B 9A 11 00 00 00 00 00 00 00 00
            21 20 22 08 00 00 00 00 00 00 02 3A 22 00 01 64
            02 00 00 2C 01 01 00 04 03 00 00 0C 01 00 00 14
            80 0E 01 00 03 00 00 08 02 00 00 05 03 00 00 08
            06 00 00 24 00 00 00 08 04 00 00 13 02 00 00 2C
            02 01 00 04 03 00 00 0C 01 00 00 14 80 0E 01 00
            03 00 00 08 02 00 00 05 03 00 00 08 06 00 00 24
            00 00 00 08 04 00 00 0E 02 00 00 34 03 01 00 05
            03 00 00 0C 01 00 00 0C 80 0E 01 00 03 00 00 08
            03 00 00 0C 03 00 00 08 02 00 00 05 03 00 00 08
            06 00 00 24 00 00 00 08 04 00 00 13 02 00 00 34
            04 01 00 05 03 00 00 0C 01 00 00 0C 80 0E 01 00
            03 00 00 08 03 00 00 0C 03 00 00 08 02 00 00 05
            03 00 00 08 06 00 00 24 00 00 00 08 04 00 00 0E
            02 00 00 24 05 01 00 03 03 00 00 0C 01 00 00 14
            80 0E 01 00 03 00 00 08 02 00 00 05 00 00 00 08
            04 00 00 13 02 00 00 24 06 01 00 03 03 00 00 0C
            01 00 00 14 80 0E 01 00 03 00 00 08 02 00 00 05
            00 00 00 08 04 00 00 0E 02 00 00 2C 07 01 00 04
            03 00 00 0C 01 00 00 0C 80 0E 01 00 03 00 00 08
            03 00 00 0C 03 00 00 08 02 00 00 05 00 00 00 08
            04 00 00 13 00 00 00 2C 08 01 00 04 03 00 00 0C
            01 00 00 0C 80 0E 01 00 03 00 00 08 03 00 00 0C
            03 00 00 08 02 00 00 05 00 00 00 08 04 00 00 0E
            28 00 00 48 00 13 00 00 76 4C 55 51 F0 73 E8 8E
            AE 62 8C 82 10 32 D3 DB 10 7A 24 27 9C B4 A0 F0
            A0 31 C3 FF 8E D5 10 7A 01 43 CC 3F 51 0A 66 4A
            15 7D EF 81 D0 55 FD 58 60 BC 71 9A C7 FA 2C 13
            EB 8A DD CA 3E 71 53 2D 29 00 00 14 08 CE DD 13
            DE 02 33 7B 8A 72 3F 21 7B 07 2C C2 29 00 00 08
            00 00 40 16 29 00 00 1C 00 00 40 04 FE 75 8E 84
            ED 11 8A 37 8D FB F9 05 0B CB 40 D8 9C 88 24 25
            29 00 00 1C 00 00 40 05 D2 DC 28 01 6E 4C 2F 7B
            2B 61 97 BD 71 00 13 35 8C 6C 4D B1 29 00 00 08
            00 00 40 2E 29 00 00 0E 00 00 40 2F 00 04 00 03
            00 02 00 00 00 08 00 00 40 36
        ";
        HEX.split_ascii_whitespace()
            .map(|h| u8::from_str_radix(h, 16).unwrap())
            .collect()
    }
}