unix.alu

use io::{Result, Error, FileDescriptor};
use option::{Option};
use io::unix::{errno_try, FdReadWrite};

/// Which end(s) of the socket to shutdown
enum Shutdown {
    /// Shut down the reading end of the socket
    Read = libc::SHUT_RD,
    /// Shut down the writing end of the socket
    Write = libc::SHUT_WR,
    /// Shut down both ends of the socket
    Both = libc::SHUT_RDWR,
}

/// Domain name resolution iterator
///
/// ## Example
/// ```
/// use std::net::NameLookup;
/// let addresses = NameLookup::resolve("www.alumina-lang.net", 443).unwrap();
/// defer addresses.free();
///
/// for addr in addresses {
///     println!("{}", addr);
/// }
/// ```
struct NameLookup {
    original: &mut libc::addrinfo,
    cur: &mut libc::addrinfo,
    port: u16,
}

impl NameLookup {
    /// Resolve a hostname to network addresses
    fn resolve(host: &[u8], port: u16) -> Result<NameLookup> {
        let host = ffi::CString::from_slice(host);
        defer host.free();

        let hints = mem::zeroed::<libc::addrinfo>();
        hints.ai_socktype = libc::SOCK_STREAM;

        let res: &mut libc::addrinfo = null;
        let ret = libc::getaddrinfo(host.ptr, null, &hints, &res);

        if ret == 0 {
            Result::ok(NameLookup {
                original: res,
                cur: res,
                port: port
            })
        } else {
            if ret == libc::EAI_SYSTEM {
                Result::err(Error::from_errno())
            } else {
                Result::err(Error::from_gai(ret))
            }
        }
    }

    /// @ iter::Iterator::next
    fn next(self: &mut NameLookup) -> Option<SocketAddr> {
        if self.cur == null {
            Option::none()
        } else {
            let cur = self.cur;
            self.cur = cur.ai_next;

            let result = internal::sockaddr_to_addr(
                cur.ai_addr as &libc::sockaddr_storage,
                cur.ai_addrlen
            );

            if result.is_some() {
                let result = result.unwrap();
                result.set_port(self.port);

                Option::some(result)
            } else {
                // Ignore addresses that are not IPv6 or IPv4.
                return self.next();
            }
        }
    }

    /// @ mem::Freeable::free
    fn free(self: &mut NameLookup) {
        libc::freeaddrinfo(self.original);
    }

    /// @ mem::Movable::move
    fn move(self: &mut NameLookup) -> NameLookup {
        let res = *self;
        self.original = null;
        res
    }

    mixin std::iter::Iterator<NameLookup, SocketAddr>;
    mixin std::iter::IteratorExt<NameLookup, SocketAddr>;
}

/// A network or local socket
struct Socket {
    fd: FileDescriptor,
}

impl Socket {
    /// Create a new socket with a given family and type
    fn new(family: libc::c_int, type: libc::c_int) -> Result<Socket> {
        let fd: FileDescriptor;

        #[cfg(any(target_os = "android", target_os = "linux"))]
        fd = FileDescriptor::new(errno_try!(libc::socket(family, type | libc::O_CLOEXEC, 0)));

        #[cfg(not(any(target_os = "android", target_os = "linux")))]
        {
            fd = FileDescriptor::new(errno_try!(libc::socket(family, type, 0)));
            fd.set_cloexec()?;
        }

        Result::ok(Socket { fd: fd })
    }

    /// @ io::AsFileDescriptor::as_fd
    fn as_fd(self: &Socket) -> FileDescriptor {
        self.fd
    }

    /// Shuts down the reading or writing end of the socket or both.
    fn shutdown(self: &Socket, how: Shutdown) -> Result<()> {
        errno_try!(libc::shutdown(self.fd.value, how as libc::c_int));
        Result::ok(())
    }

    fn recv_from_with_flags(
        self: &Socket,
        buf: &mut [u8],
        flags: libc::c_int,
    ) -> Result<(usize, SocketAddr)> {
        let storage: libc::sockaddr_storage = mem::zeroed();
        let addrlen = mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;

        let n = errno_try!(
            libc::recvfrom(
                self.fd.value,
                &buf[0] as &mut void,
                buf.len(),
                flags,
                &storage as &mut libc::sockaddr,
                &addrlen,
            )
        );

        Result::ok((n as usize, internal::sockaddr_to_addr(&storage, addrlen).unwrap()))
    }

    fn recv_with_flags(self: &Socket, buf: &mut [u8], flags: libc::c_int) -> Result<usize> {
        let n = errno_try!(
            libc::recv(
                self.fd.value,
                &buf[0] as &mut void,
                buf.len(),
                flags,
            )
        );

        Result::ok(n as usize)
    }

    /// @ io::Readable::read
    fn read(self: &Socket, buf: &mut [u8]) -> Result<usize> {
        self.recv_with_flags(buf, 0)
    }

    /// Returns the address of the remote peer
    fn peer_addr(self: &Socket) -> Result<SocketAddr> {
        let storage = mem::zeroed::<libc::sockaddr_storage>();
        let len: libc::socklen_t = mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;

        errno_try!(libc::getpeername(self.fd.value, &storage as &mut libc::sockaddr, &len));
        Result::ok(internal::sockaddr_to_addr(&storage, len).unwrap())
    }

    /// Returns the address of the remote peer
    fn socket_addr(self: &Socket) -> Result<SocketAddr> {
        let storage = mem::zeroed::<libc::sockaddr_storage>();
        let len: libc::socklen_t = mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;

        errno_try!(libc::getsockname(self.fd.value, &storage as &mut libc::sockaddr, &len));
        Result::ok(internal::sockaddr_to_addr(&storage, len).unwrap())
    }

    /// Closes the socket.
    fn close(self: &mut Socket) -> Result<()> {
        self.fd.close()
    }

    /// @ mem::Movable::move
    fn move(self: &mut Socket) -> Socket {
        Socket { fd: self.fd.move() }
    }

    // Only write, read is overriden
    mixin FdReadWrite<Socket>;
}

/// A TCP connection
///
/// ## Example
/// ```no_run
/// use std::net::{NameLookup, TcpStream, Shutdown};
/// use std::io::{Result, StdioStream, copy};
///
/// fn http_fetch() -> Result<()> {
///     let addrs = NameLookup::resolve("docs.alumina-lang.net", 80)?;
///     defer addrs.free();
///
///     let socket = TcpStream::connect(&addrs.next().unwrap())?;
///     defer socket.close();
///
///     socket.write_all("GET / HTTP/1.1\r\n")?;
///     socket.write_all("Host: docs.alumina-lang.net\r\n")?;
///     socket.write_all("Connection: close\r\n\r\n")?;
///
///     socket.shutdown(Shutdown::Write)?;
///
///     copy(&socket, &StdioStream::stdout())?;
///
///     Result::ok(())
/// }
///
/// http_fetch().unwrap();
/// ```
struct TcpStream {
    socket: Socket,
}

impl TcpStream {
    /// Establishes a TCP connection to a remote host.
    fn connect(addr: &SocketAddr) -> Result<TcpStream> {
        let (family, addr_len) = internal::addr_to_family(addr);

        let sock = Socket::new(family, libc::SOCK_STREAM)?;
        defer sock.close();

        errno_try!(libc::connect(
            sock.fd.value,
            &addr.inner as &libc::sockaddr,
            addr_len as libc::socklen_t
        ));

        Result::ok(from_socket(sock.move()))
    }

    /// Create a new TCP stream from a raw socket
    fn from_socket(socket: Socket) -> TcpStream {
        TcpStream { socket: socket }
    }

    /// @ io::AsFileDescriptor::as_fd
    fn as_fd(self: &TcpStream) -> FileDescriptor {
        self.socket.as_fd()
    }

    /// @ Socket::peer_addr
    fn peer_addr(self: &TcpStream) -> Result<SocketAddr> {
        self.socket.peer_addr()
    }

    /// @ Socket::socket_addr
    fn socket_addr(self: &TcpStream) -> Result<SocketAddr> {
        self.socket.socket_addr()
    }

    /// @ Socket::shutdown
    fn shutdown(self: &TcpStream, how: Shutdown) -> Result<()> {
        self.socket.shutdown(how)
    }

    /// @ io::Readable::read
    fn read(self: &mut TcpStream, buf: &mut [u8]) -> Result<usize> {
        self.socket.read(buf)
    }

    /// @ io::Writable::write
    fn write(self: &mut TcpStream, buf: &[u8]) -> Result<usize> {
        self.socket.write(buf)
    }

    /// @ io::Writable::flush
    fn flush(self: &mut TcpStream) -> Result<()> {
        Result::ok(())
    }

    /// @ Socket::close
    fn close(self: &mut TcpStream) -> Result<()> {
        self.socket.close()
    }

    /// @ mem::Movable::move
    fn move(self: &mut TcpStream) -> TcpStream {
        TcpStream { socket: self.socket.move() }
    }

    mixin io::Readable<TcpStream>;
    mixin io::Writable<TcpStream>;
}

/// A TCP listener which can accept new connections.
struct TcpListener {
    socket: Socket,
}

impl TcpListener {
    /// Create a new TCP listener from a socket.
    fn from_socket(socket: Socket) -> TcpListener {
        TcpListener { socket: socket }
    }

    /// Creates a new TCP listener bound to the specified address.
    fn bind(addr: &SocketAddr) -> Result<TcpListener> {
        let (family, addr_len) = internal::addr_to_family(addr);

        let sock = Socket::new(family, libc::SOCK_STREAM)?;
        defer sock.close();

        errno_try!(libc::bind(
            sock.fd.value,
            &addr.inner as &libc::sockaddr,
            addr_len
        ));

        errno_try!(libc::listen(
            sock.fd.value,
            128
        ));

        Result::ok(from_socket(sock.move()))
    }

    /// @ io::AsFileDescriptor::as_fd
    fn as_fd(self: &TcpListener) -> FileDescriptor {
        self.socket.as_fd()
    }

    /// Accepts a new incoming connection.
    ///
    /// Returns a new `TcpStream` representing the accepted connection and the address
    /// of the remote peer.
    fn accept(self: &TcpListener) -> Result<(TcpStream, SocketAddr)> {
        let storage = mem::zeroed::<libc::sockaddr_storage>();
        let len: libc::socklen_t = mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;

        let fd: FileDescriptor;

        #[cfg(any(target_os = "android", target_os = "linux"))]
        fd = FileDescriptor::new(errno_try!(libc::accept4(
            self.socket.fd.value,
            &storage as &mut libc::sockaddr,
            &len,
            libc::O_CLOEXEC
        )));


        #[cfg(not(any(target_os = "android", target_os = "linux")))]
        {
            fd = FileDescriptor::new(errno_try!(libc::accept(
                self.socket.fd.value,
                &storage as &mut libc::sockaddr,
                &len
            )));
            fd.set_cloexec()?;
        }

        let socket = Socket { fd: fd };

        Result::ok((
            TcpStream { socket: socket },
            internal::sockaddr_to_addr(&storage, len).unwrap()
        ))
    }

    /// @ Socket::socket_addr
    fn socket_addr(self: &TcpListener) -> Result<SocketAddr> {
        self.socket.socket_addr()
    }

    /// @ Socket::close
    fn close(self: &mut TcpListener) -> Result<()> {
        self.socket.close()
    }

    /// @ mem::Movable::move
    fn move(self: &mut TcpListener) -> TcpStream {
        TcpStream { socket: self.socket.move() }
    }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
use libc::MSG_NOSIGNAL;

#[cfg(not(any(target_os = "linux", target_os = "android")))]
const MSG_NOSIGNAL: libc::c_int = 0;

/// A UDP socket
///
/// ## Example
///
/// This example sends a DNS query to CloudFlare's DNS server and prints the response metadata.
///
/// ```no_run
/// use std::net::{SocketAddr, UdpSocket};
/// use std::io::Result;
///
/// fn dns_query() -> Result<()> {
///     let request = std::concat!(
///         "\x00\x01", // Transaction id
///         "\x01\x20", // Flags (standard query)
///         "\x00\x01", // Question RRs
///         "\x00\x00", // Answer RRs
///         "\x00\x00", // Authority RRs
///         "\x00\x00", // Additional RRs
///         "\x04docs", // Domain name
///         "\x0calumina-lang",
///         "\x03net",
///         "\0",
///         "\x00\x1c", // Type (AAAA)
///         "\x00\x01"  // Class (IN)
///     );
///
///     // Bind on a random local port
///     let local_addr = SocketAddr::parse("[::]:0").unwrap();
///     // Cloudflare DNS resolver
///     let remote_addr = SocketAddr::parse("[2606:4700:4700::1111]:53").unwrap();
///
///     let socket = UdpSocket::bind(&local_addr)?;
///     defer socket.close();
///
///     socket.send_to(request, &remote_addr)?;
///
///     let buf: [u8; 65536];
///     let (size, src) = socket.recv_from(&buf)?;
///
///     println!("Received {} bytes from {}", size, src);
///     println!("  {} questions",   (&buf[4]  as &u16).from_be());
///     println!("  {} answers",     (&buf[6]  as &u16).from_be());
///     println!("  {} authorities", (&buf[8]  as &u16).from_be());
///     println!("  {} additional",  (&buf[10] as &u16).from_be());
///
///     Result::ok(())
/// }
///
/// dns_query().unwrap();
/// ```
struct UdpSocket {
    socket: Socket,
}

impl UdpSocket {
    /// Create a new UDP socket from a raw socket.
    fn from_socket(socket: Socket) -> UdpSocket {
        UdpSocket { socket: socket }
    }

    /// Creates a new UDP socket bound to the specified address.
    fn bind(addr: &SocketAddr) -> Result<UdpSocket> {
        let (family, addr_len) = internal::addr_to_family(addr);

        let sock = Socket::new(family, libc::SOCK_DGRAM)?;
        defer sock.close();

        errno_try!(libc::bind(
            sock.fd.value,
            &addr.inner as &libc::sockaddr,
            addr_len
        ));

        Result::ok(from_socket(sock.move()))
    }

    /// "Connects" the UDP socket to the specified address.
    ///
    /// As UDP is a connectionless protocol, this does not actually perform any
    /// connection, it just attaches the remote address to the socket.
    ///
    /// Subsequent calls to [recv] and [send] will only receive and send data
    /// from and to this address.
    fn connect(self: &UdpSocket, addr: &SocketAddr) -> Result<()> {
        let (_, addr_len) = internal::addr_to_family(addr);

        errno_try!(libc::connect(
            self.socket.fd.value,
            &addr.inner as &libc::sockaddr,
            addr_len as libc::socklen_t
        ));

        Result::ok(())
    }

    /// Receives data from the socket.
    ///
    /// Returns the number of bytes read and the address from which the data was received.
    fn recv_from(
        self: &UdpSocket,
        buf: &mut [u8],
    ) -> Result<(usize, SocketAddr)> {
        self.socket.recv_from_with_flags(buf, 0)
    }

    /// Receives data from the socket without consuming it.
    ///
    /// Subsequent calls to [recv] will return the same data.
    fn peek_from(
        self: &UdpSocket,
        buf: &mut [u8],
    ) -> Result<(usize, SocketAddr)> {
        self.socket.recv_from_with_flags(buf, libc::MSG_PEEK)
    }

    /// Sends data to a particular address.
    fn send_to(self: &UdpSocket, buf: &[u8], addr: &SocketAddr) -> Result<usize> {
        let (_, addr_len) = internal::addr_to_family(addr);

        let ret = errno_try!(libc::sendto(
            self.socket.fd.value,
            &buf[0] as &void,
            buf.len(),
            MSG_NOSIGNAL,
            &addr.inner as &libc::sockaddr,
            addr_len,
        ));

        Result::ok(ret as usize)
    }

    /// Receives data from the socket.
    fn recv(self: &UdpSocket, buf: &mut [u8]) -> Result<usize> {
        self.socket.recv_with_flags(buf, 0)
    }

    /// Receives data from the socket without consuming it.
    fn peek(self: &UdpSocket, buf: &mut [u8]) -> Result<usize> {
        self.socket.recv_with_flags(buf, libc::MSG_PEEK)
    }

    /// Sends data to the socket.
    fn send(self: &UdpSocket, buf: &[u8]) -> Result<usize> {
        let ret = errno_try!(
            libc::send(self.socket.fd.value, &buf[0] as &void, buf.len(), MSG_NOSIGNAL)
        );

        Result::ok(ret as usize)
    }

    /// @ io::AsFileDescriptor::as_fd
    fn as_fd(self: &UdpSocket) -> FileDescriptor {
        self.socket.as_fd()
    }

    /// @ Socket::peer_addr
    fn peer_addr(self: &UdpSocket) -> Result<SocketAddr> {
        self.socket.peer_addr()
    }

    /// @ Socket::socket_addr
    fn socket_addr(self: &UdpSocket) -> Result<SocketAddr> {
        self.socket.socket_addr()
    }

    /// @ Socket::close
    fn close(self: &mut UdpSocket) -> Result<()> {
        self.socket.close()
    }

    /// @ mem::Movable::move
    fn move(self: &mut UdpSocket) -> UdpSocket {
        UdpSocket { socket: self.socket.move() }
    }
}

#[docs(no_index)]
mod internal {
    fn sockaddr_to_addr(addr: &libc::sockaddr_storage, _len: libc::socklen_t) -> Option<SocketAddr> {
        let result: SocketAddr;
        switch addr.ss_family as libc::c_int {
            libc::AF_INET => {
                result.kind = AddrKind::V4;
                result.inner.v4 = *(addr as &libc::sockaddr_in);
            }
            libc::AF_INET6 => {
                result.kind = AddrKind::V6;
                result.inner.v6 = *(addr as &libc::sockaddr_in6);
            }
            _ => {
                return Option::none();
            }
        };

        Option::some(result)
    }

    fn addr_to_family(addr: &SocketAddr) -> (libc::c_int, libc::socklen_t) {
        let (family, addr_len) = switch addr.kind {
            AddrKind::V4 => (libc::AF_INET, mem::size_of::<libc::sockaddr_in>()),
            AddrKind::V6 => (libc::AF_INET6, mem::size_of::<libc::sockaddr_in6>()),
            _ => unreachable!()
        };

        (family, addr_len as libc::socklen_t)
    }
}

#[cfg(all(test, test_std))]
#[docs(hide)]
mod tests {
    use process::{Forked, Child};
    use string::StringBuf;

    fn fork() -> Option<Child> {
        Forked::new().spawn().unwrap()
    }

    #[test]
    fn test_name_resolution() {
        let addrs = NameLookup::resolve("localhost", 80).unwrap();
        defer addrs.free();

        let v6 = SocketAddr::new(IpAddr::v6(Ipv6Addr::localhost()), 80);
        let v4 = SocketAddr::new(IpAddr::v4(Ipv4Addr::localhost()), 80);

        assert!(addrs.all(|=v6, =v4, a: SocketAddr| -> bool {
            a == v6 || a == v4
        }));
    }

    #[test]
    fn test_basic_tcp() {
        let addr = SocketAddr::parse("[::1]:0").unwrap();
        let listener = TcpListener::bind(&addr).unwrap();
        defer listener.close();

        let local_addr = listener.socket_addr().unwrap();

        if fork().is_none() {
            let socket = TcpStream::connect(&local_addr).unwrap();
            defer socket.close();

            socket.write_all("Hello, world!").unwrap();
            socket.shutdown(Shutdown::Both).unwrap();
        } else {
            let (socket, _) = listener.accept().unwrap();
            defer socket.close();

            let buf = StringBuf::new();
            socket.read_to_end(&buf).unwrap();

            assert_eq!(buf[..] as &[u8], "Hello, world!");
        }
    }

    #[test]
    fn test_basic_udp() {
        let addr = SocketAddr::parse("[::1]:0").unwrap();
        let peer1 = UdpSocket::bind(&addr).unwrap();
        defer peer1.close();

        let peer1_addr = peer1.socket_addr().unwrap();
        if fork().is_none() {
            let addr = SocketAddr::parse("[::1]:0").unwrap();
            let peer2 = UdpSocket::bind(&addr).unwrap();
            defer peer2.close();

            peer2.send_to("Hello, world!", &peer1_addr).unwrap();
        } else {
            let buf: [u8; 128];
            let (n, _) = peer1.recv_from(&buf).unwrap();

            assert_eq!(buf[..n] as &[u8], "Hello, world!");
        }
    }

    #[test]
    fn test_udp_connect() {
        let addr = SocketAddr::parse("[::1]:0").unwrap();
        let peer1 = UdpSocket::bind(&addr).unwrap();
        defer peer1.close();

        let buf: [u8; 128];

        let peer1_addr = peer1.socket_addr().unwrap();
        if fork().is_none() {
            let addr = SocketAddr::parse("[::1]:0").unwrap();
            let peer2 = UdpSocket::bind(&addr).unwrap();
            defer peer2.close();

            peer2.connect(&peer1_addr).unwrap();
            peer2.send("Hello, world!").unwrap();

            let n = peer2.recv(&buf).unwrap();
            peer2.send(buf[..n]).unwrap();

        } else {
            let buf: [u8; 128];
            let (_, peer) = peer1.recv_from(&buf).unwrap();

            peer1.send_to("AAAAAAA", &peer).unwrap();

            let (n, _) = peer1.recv_from(&buf).unwrap();
            assert_eq!(buf[..n] as &[u8], "AAAAAAA");
        }
    }

    #[test]
    fn test_udp_peek() {
        let addr = SocketAddr::parse("[::1]:0").unwrap();
        let peer1 = UdpSocket::bind(&addr).unwrap();
        defer peer1.close();

        let peer1_addr = peer1.socket_addr().unwrap();
        if fork().is_none() {
            let addr = SocketAddr::parse("[::1]:0").unwrap();
            let peer2 = UdpSocket::bind(&addr).unwrap();
            defer peer2.close();

            peer2.send_to("Hello, world!", &peer1_addr).unwrap();
        } else {
            let buf1: [u8; 128];
            let buf2: [u8; 128];
            let (n, _) = peer1.peek_from(&buf1).unwrap();
            assert_eq!(buf1[..n] as &[u8], "Hello, world!");

            let (n, _) = peer1.recv_from(&buf2).unwrap();
            assert_eq!(buf2[..n] as &[u8], "Hello, world!");
        }
    }
}