// Network control handlers — IP allowlist enforcement via seccomp notification.
//
// Intercepts connect/sendto/sendmsg syscalls, extracts the destination IP from
// the child's memory, and checks it against an allowlist of resolved IPs.

use std::collections::HashSet;
use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::os::unix::io::RawFd;
use std::sync::Arc;

use tokio::sync::Mutex;

use std::os::unix::io::AsRawFd;

use crate::seccomp::notif::{read_child_mem, NotifAction, SupervisorState};
use crate::sys::structs::{SeccompNotif, AF_INET, AF_INET6, ECONNREFUSED};

// ============================================================
// parse_ip_from_sockaddr — parse IP from a sockaddr byte buffer
// ============================================================

/// Parse IP address from a sockaddr byte buffer.
/// Returns None for non-IP families (AF_UNIX etc.) — always allowed.
fn parse_ip_from_sockaddr(bytes: &[u8]) -> Option<IpAddr> {
    if bytes.len() < 2 {
        return None;
    }
    let family = u16::from_ne_bytes([bytes[0], bytes[1]]) as u32;
    match family {
        f if f == AF_INET => {
            if bytes.len() < 8 {
                return None;
            }
            Some(IpAddr::V4(Ipv4Addr::new(
                bytes[4], bytes[5], bytes[6], bytes[7],
            )))
        }
        f if f == AF_INET6 => {
            if bytes.len() < 24 {
                return None;
            }
            let mut addr_bytes = [0u8; 16];
            addr_bytes.copy_from_slice(&bytes[8..24]);
            Some(IpAddr::V6(Ipv6Addr::from(addr_bytes)))
        }
        _ => None,
    }
}

// ============================================================
// connect_on_behalf — perform connect() on behalf of the child (TOCTOU-safe)
// ============================================================

/// Perform connect() on behalf of the child process (TOCTOU-safe).
///
/// 1. Copy sockaddr from child memory (our copy — immune to TOCTOU)
/// 2. Check IP against allowlist on our copy
/// 3. Duplicate child's socket fd via pidfd_getfd
/// 4. connect() in supervisor with our validated sockaddr
/// 5. Return result to child
async fn connect_on_behalf(
    notif: &SeccompNotif,
    state: &Arc<Mutex<SupervisorState>>,
    notif_fd: RawFd,
) -> NotifAction {
    let args = &notif.data.args;
    let sockfd = args[0] as i32;
    let addr_ptr = args[1];
    let addr_len = args[2] as u32;

    // 1. Copy sockaddr from child memory
    let addr_bytes =
        match read_child_mem(notif_fd, notif.id, notif.pid, addr_ptr, addr_len as usize) {
            Ok(b) => b,
            Err(_) => return NotifAction::Errno(libc::EIO),
        };

    // 2. Check IP against allowlist
    if let Some(ip) = parse_ip_from_sockaddr(&addr_bytes) {
        let st = state.lock().await;
        if let crate::seccomp::notif::NetworkPolicy::AllowList(ref allowed) =
            st.effective_network_policy(notif.pid)
        {
            if !allowed.contains(&ip) {
                return NotifAction::Errno(ECONNREFUSED);
            }
        }
        let child_pidfd = match st.child_pidfd {
            Some(fd) => fd,
            None => return NotifAction::Errno(libc::ENOSYS),
        };
        drop(st);

        // 3. Duplicate child's socket into supervisor
        let dup_fd = match crate::seccomp::notif::dup_child_fd(child_pidfd, sockfd) {
            Ok(fd) => fd,
            Err(_) => return NotifAction::Errno(libc::ENOSYS),
        };

        // 4. Perform connect in supervisor with our validated sockaddr
        let ret = unsafe {
            libc::connect(
                dup_fd.as_raw_fd(),
                addr_bytes.as_ptr() as *const libc::sockaddr,
                addr_len as libc::socklen_t,
            )
        };

        // 5. Return result
        if ret == 0 {
            NotifAction::ReturnValue(0)
        } else {
            let errno = unsafe { *libc::__errno_location() };
            NotifAction::Errno(errno)
        }
        // dup_fd dropped here, closing supervisor's copy
    } else {
        // Non-IP family (AF_UNIX etc.) — allow through
        NotifAction::Continue
    }
}

// ============================================================
// sendto_on_behalf / sendmsg_on_behalf — on-behalf (TOCTOU-safe)
// ============================================================

/// Perform sendto() on behalf of the child process (TOCTOU-safe).
///
/// 1. Copy sockaddr from child memory (our copy — immune to TOCTOU)
/// 2. Check IP against allowlist on our copy
/// 3. Copy data buffer from child memory
/// 4. Duplicate child's socket fd via pidfd_getfd
/// 5. sendto() in supervisor with validated sockaddr + copied data
/// 6. Return byte count or errno
///
/// Only triggers for unconnected sends (addr_ptr != NULL), which is
/// primarily UDP. Connected sockets (addr_ptr == NULL) use CONTINUE.
async fn sendto_on_behalf(
    notif: &SeccompNotif,
    state: &Arc<Mutex<SupervisorState>>,
    notif_fd: RawFd,
) -> NotifAction {
    let args = &notif.data.args;
    let sockfd = args[0] as i32;
    let buf_ptr = args[1];
    let buf_len = args[2] as usize;
    let flags = args[3] as i32;
    let addr_ptr = args[4];
    let addr_len = args[5] as u32;

    if addr_ptr == 0 {
        return NotifAction::Continue; // connected socket, no addr to check
    }

    // 1. Copy sockaddr from child memory (small: 16-28 bytes)
    let addr_bytes =
        match read_child_mem(notif_fd, notif.id, notif.pid, addr_ptr, addr_len as usize) {
            Ok(b) => b,
            Err(_) => return NotifAction::Errno(libc::EIO),
        };

    // 2. Check IP against allowlist
    if let Some(ip) = parse_ip_from_sockaddr(&addr_bytes) {
        let st = state.lock().await;
        if let crate::seccomp::notif::NetworkPolicy::AllowList(ref allowed) =
            st.effective_network_policy(notif.pid)
        {
            if !allowed.contains(&ip) {
                return NotifAction::Errno(ECONNREFUSED);
            }
        }
        let child_pidfd = match st.child_pidfd {
            Some(fd) => fd,
            None => return NotifAction::Errno(libc::ENOSYS),
        };
        drop(st);

        // 3. Copy data buffer from child memory
        let data = match read_child_mem(notif_fd, notif.id, notif.pid, buf_ptr, buf_len) {
            Ok(b) => b,
            Err(_) => return NotifAction::Errno(libc::EIO),
        };

        // 4. Duplicate child's socket into supervisor
        let dup_fd = match crate::seccomp::notif::dup_child_fd(child_pidfd, sockfd) {
            Ok(fd) => fd,
            Err(_) => return NotifAction::Errno(libc::ENOSYS),
        };

        // 5. Perform sendto in supervisor with validated sockaddr + copied data
        let ret = unsafe {
            libc::sendto(
                dup_fd.as_raw_fd(),
                data.as_ptr() as *const libc::c_void,
                data.len(),
                flags,
                addr_bytes.as_ptr() as *const libc::sockaddr,
                addr_len as libc::socklen_t,
            )
        };

        // 6. Return result
        if ret >= 0 {
            NotifAction::ReturnValue(ret as i64)
        } else {
            let errno = unsafe { *libc::__errno_location() };
            NotifAction::Errno(errno)
        }
    } else {
        // Non-IP family (AF_UNIX etc.) — allow through
        NotifAction::Continue
    }
}

/// Perform sendmsg() on behalf of the child process (TOCTOU-safe).
///
/// 1. Copy full msghdr from child memory
/// 2. Copy sockaddr from msg_name (our copy — immune to TOCTOU)
/// 3. Check IP against allowlist on our copy
/// 4. Copy iovec data buffers from child memory
/// 5. Copy control message buffer from child memory
/// 6. Duplicate child's socket fd via pidfd_getfd
/// 7. sendmsg() in supervisor with validated sockaddr + copied data
/// 8. Return byte count or errno
async fn sendmsg_on_behalf(
    notif: &SeccompNotif,
    state: &Arc<Mutex<SupervisorState>>,
    notif_fd: RawFd,
) -> NotifAction {
    let args = &notif.data.args;
    let sockfd = args[0] as i32;
    let msghdr_ptr = args[1];
    let flags = args[2] as i32;

    // 1. Read full msghdr struct (56 bytes on x86_64):
    //   msg_name(8) + msg_namelen(4) + pad(4) + msg_iov(8) + msg_iovlen(8)
    //   + msg_control(8) + msg_controllen(8) + msg_flags(4) + pad(4)
    let msghdr_bytes = match read_child_mem(notif_fd, notif.id, notif.pid, msghdr_ptr, 56) {
        Ok(b) if b.len() >= 56 => b,
        _ => return NotifAction::Continue,
    };

    let msg_name_ptr = u64::from_ne_bytes(msghdr_bytes[0..8].try_into().unwrap());
    let msg_namelen = u32::from_ne_bytes(msghdr_bytes[8..12].try_into().unwrap());
    let msg_iov_ptr = u64::from_ne_bytes(msghdr_bytes[16..24].try_into().unwrap());
    let msg_iovlen = u64::from_ne_bytes(msghdr_bytes[24..32].try_into().unwrap());
    let msg_control_ptr = u64::from_ne_bytes(msghdr_bytes[32..40].try_into().unwrap());
    let msg_controllen = u64::from_ne_bytes(msghdr_bytes[40..48].try_into().unwrap());

    if msg_name_ptr == 0 {
        return NotifAction::Continue; // no address — connected socket
    }

    // 2. Copy sockaddr from msg_name
    let addr_bytes = match read_child_mem(
        notif_fd, notif.id, notif.pid, msg_name_ptr, msg_namelen as usize,
    ) {
        Ok(b) => b,
        Err(_) => return NotifAction::Errno(libc::EIO),
    };

    // 3. Check IP against allowlist
    let ip = match parse_ip_from_sockaddr(&addr_bytes) {
        Some(ip) => ip,
        None => return NotifAction::Continue, // Non-IP family — allow through
    };

    let st = state.lock().await;
    if let crate::seccomp::notif::NetworkPolicy::AllowList(ref allowed) =
        st.effective_network_policy(notif.pid)
    {
        if !allowed.contains(&ip) {
            return NotifAction::Errno(ECONNREFUSED);
        }
    }
    let child_pidfd = match st.child_pidfd {
        Some(fd) => fd,
        None => return NotifAction::Errno(libc::ENOSYS),
    };
    drop(st);

    // 4. Copy iovec entries and their data buffers from child memory
    // Safety: cap iovlen to prevent excessive allocation
    let iovlen = (msg_iovlen as usize).min(1024);
    let iov_size = iovlen * 16; // each iovec is 16 bytes (ptr + len)
    let iov_bytes = match read_child_mem(notif_fd, notif.id, notif.pid, msg_iov_ptr, iov_size) {
        Ok(b) => b,
        Err(_) => return NotifAction::Errno(libc::EIO),
    };

    let mut data_bufs: Vec<Vec<u8>> = Vec::with_capacity(iovlen);
    let mut local_iovs: Vec<libc::iovec> = Vec::with_capacity(iovlen);

    for i in 0..iovlen {
        let off = i * 16;
        if off + 16 > iov_bytes.len() { break; }
        let iov_base = u64::from_ne_bytes(iov_bytes[off..off + 8].try_into().unwrap());
        let iov_len = u64::from_ne_bytes(iov_bytes[off + 8..off + 16].try_into().unwrap()) as usize;

        if iov_base == 0 || iov_len == 0 {
            data_bufs.push(Vec::new());
            continue;
        }

        let buf = match read_child_mem(notif_fd, notif.id, notif.pid, iov_base, iov_len) {
            Ok(b) => b,
            Err(_) => return NotifAction::Errno(libc::EIO),
        };
        data_bufs.push(buf);
    }

    // Build local iovec array pointing to our copied data
    for buf in &data_bufs {
        local_iovs.push(libc::iovec {
            iov_base: buf.as_ptr() as *mut libc::c_void,
            iov_len: buf.len(),
        });
    }

    // 5. Copy control message buffer (ancillary data)
    let control_buf = if msg_control_ptr != 0 && msg_controllen > 0 {
        let len = (msg_controllen as usize).min(4096);
        read_child_mem(notif_fd, notif.id, notif.pid, msg_control_ptr, len).ok()
    } else {
        None
    };

    // 6. Duplicate child's socket into supervisor
    let dup_fd = match crate::seccomp::notif::dup_child_fd(child_pidfd, sockfd) {
        Ok(fd) => fd,
        Err(_) => return NotifAction::Errno(libc::ENOSYS),
    };

    // 7. Build msghdr and perform sendmsg in supervisor
    let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
    msg.msg_name = addr_bytes.as_ptr() as *mut libc::c_void;
    msg.msg_namelen = addr_bytes.len() as u32;
    msg.msg_iov = local_iovs.as_mut_ptr();
    msg.msg_iovlen = local_iovs.len();
    if let Some(ref ctrl) = control_buf {
        msg.msg_control = ctrl.as_ptr() as *mut libc::c_void;
        msg.msg_controllen = ctrl.len();
    }

    let ret = unsafe { libc::sendmsg(dup_fd.as_raw_fd(), &msg, flags) };

    // 8. Return result
    if ret >= 0 {
        NotifAction::ReturnValue(ret as i64)
    } else {
        let errno = unsafe { *libc::__errno_location() };
        NotifAction::Errno(errno)
    }
}

// ============================================================
// handle_net — main handler for connect/sendto/sendmsg
// ============================================================

/// Handle network-related notifications (connect, sendto, sendmsg).
///
/// All three are handled on-behalf (TOCTOU-safe): the supervisor copies data
/// from child memory, validates the destination, duplicates the socket via
/// pidfd_getfd, and performs the syscall itself. The child's memory is never
/// re-read by the kernel after validation.
pub(crate) async fn handle_net(
    notif: &SeccompNotif,
    state: &Arc<Mutex<SupervisorState>>,
    notif_fd: RawFd,
) -> NotifAction {
    let nr = notif.data.nr as i64;

    if nr == libc::SYS_connect {
        connect_on_behalf(notif, state, notif_fd).await
    } else if nr == libc::SYS_sendto {
        sendto_on_behalf(notif, state, notif_fd).await
    } else if nr == libc::SYS_sendmsg {
        sendmsg_on_behalf(notif, state, notif_fd).await
    } else {
        NotifAction::Continue
    }
}

// ============================================================
// resolve_hosts — resolve domain names to IPs
// ============================================================

/// Resolve a list of domain names to IP addresses.
///
/// Always includes loopback addresses (127.0.0.1 and ::1).
/// Uses tokio's async DNS resolver.
pub async fn resolve_hosts(hosts: &[String]) -> io::Result<HashSet<IpAddr>> {
    let mut ips = HashSet::new();

    // Always allow loopback
    ips.insert(IpAddr::V4(Ipv4Addr::LOCALHOST));
    ips.insert(IpAddr::V6(Ipv6Addr::LOCALHOST));

    for host in hosts {
        // Append a dummy port for lookup_host
        let addr = format!("{}:0", host);
        let result = tokio::net::lookup_host(addr.as_str()).await;
        match result {
            Ok(resolved) => {
                for socket_addr in resolved {
                    ips.insert(socket_addr.ip());
                }
            }
            Err(e) => {
                // Return error on DNS failure to avoid silently skipping hosts
                return Err(io::Error::new(
                    e.kind(),
                    format!("failed to resolve host '{}': {}", host, e),
                ));
            }
        }
    }

    Ok(ips)
}

// ============================================================
// Tests
// ============================================================

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

    #[tokio::test]
    async fn test_resolve_hosts_loopback() {
        let ips = resolve_hosts(&[]).await.unwrap();
        assert!(ips.contains(&IpAddr::V4(Ipv4Addr::LOCALHOST)));
        assert!(ips.contains(&IpAddr::V6(Ipv6Addr::LOCALHOST)));
    }

    #[tokio::test]
    async fn test_resolve_hosts_with_domain() {
        let hosts = vec!["localhost".to_string()];
        let ips = resolve_hosts(&hosts).await.unwrap();
        // localhost should resolve to loopback
        assert!(
            ips.contains(&IpAddr::V4(Ipv4Addr::LOCALHOST))
                || ips.contains(&IpAddr::V6(Ipv6Addr::LOCALHOST))
        );
    }
}
