#define _POSIX_C_SOURCE 200809L

#include <assert.h>          /* assert */
#include <stdint.h>          /* uint8_t */
#include <sys/socket.h>      /* accept, bind, getsockname, listen, socket */
#include <unistd.h>          /* alarm, close, fork */
#include <netinet/in.h>

/* Naprogramujte proceduru ‹map_ports›, která dostane na vstupu
 * jméno hostitele a rodinu adres, a do výstupní bitmapy nastaví,
 * která čísla portů tohoto hostitele jsou „otevřená“ (přijímají
 * spojení – budeme se zabývat pouze spojovanými protokoly).
 * Ověřujte pouze ty porty, které jsou na začátku v bitmapě přítomné
 * (tzn. jejich hodnota je 1).
 *
 * Návratová hodnota bude -1, nastane-li nějaká fatální chyba, a
 * nula jinak. */

int map_ports( const char *hostname, sa_family_t family,
               uint8_t *bitmap );

/* ┄┄┄┄┄┄┄ %< ┄┄┄┄┄┄┄┄┄┄ následují testy ┄┄┄┄┄┄┄┄┄┄ %< ┄┄┄┄┄┄┄ */

#include <arpa/inet.h>       /* ntohs sockaddr_in, sockaddr_in6 */
#include <err.h>             /* err, warn */
#include <sched.h>           /* sched_yield */
#include <signal.h>          /* kill */
#include <string.h>          /* memset */
#include <sys/wait.h>        /* waitpid */

int get_port( const uint8_t *bitmap, uint16_t port )
{
    return ( bitmap[ port / 8 ] & ( 1 << ( port % 8 ) ) );
}

static void close_or_warn( int fd, const char *name )
{
    if ( close( fd ) == -1 )
        warn( "closing %s", name );
}

struct testserver
{
    pid_t pid;
    struct sockaddr_in6 addr;
};

static int prep_socket( struct sockaddr_in6 *saddr )
{
    int fd = socket( AF_INET6, SOCK_STREAM, 0 );
    if ( fd == -1 )
        err( 2, "socket" );

    const socklen_t sin6len = sizeof( struct sockaddr_in6 );

    memset( saddr, 0, sin6len );
    saddr->sin6_family = AF_INET6;
    saddr->sin6_addr.s6_addr[ 15 ] = 1; /* bind to [::1] */

    if ( bind( fd, ( struct sockaddr * ) saddr, sin6len ) == -1 )
        err( 2, "bind" );

    if ( listen( fd, 1 ) == -1 )
        err( 2, "listen" );

    socklen_t socklen = sin6len;
    if ( getsockname( fd, ( struct sockaddr * ) saddr, &socklen ) == -1 ||
         socklen != sin6len ||
         saddr->sin6_port == 0 )
        err( 2, "getsockname" );

    return fd;
}

static pid_t spawn_server( int sock_fd )
{
    pid_t pid = fork();
    if ( pid == -1 )
        err( 2, "fork" );

    if ( pid > 0 )
        return pid;

    /* server */

    /* Ukončit server nejpozději po dvou sekundách, což je čas více než
     * dostatečný na běh přiložených testů. Zejména ve sdíleném prostředí
     * serveru Aisa je velice důležité servery důsledně ukončovat, aby
     * nedošlo k nedostatku volných TCP portů. */

    alarm( 2 );

    int cfd;
    while ( ( cfd = accept( sock_fd, NULL, NULL ) ) != -1 )
        close_or_warn( cfd, "client fd in server" );

    err( 2, "accept" );
}

static int kill_and_reap( pid_t pid )
{
    if ( kill( pid, SIGTERM ) == -1 )
        err( 2, "kill" );

    int status;

    if ( waitpid( pid, &status, 0 ) == -1 )
        err( 2, "wait" );

    if ( WIFSIGNALED( status ) && WTERMSIG( status ) == SIGTERM )
        return 0;
    if ( WIFEXITED( status ) )
        return WEXITSTATUS( status );
    else
        return -1;
}

void set_port( uint8_t *bitmap, uint16_t port )
{
    bitmap[ port / 8 ] |= 1 << ( port % 8 );
}

void set_ports( uint8_t *bitmap, uint16_t *ports, int n )
{
    memset( bitmap, 0, ( 1 << 16 ) / 8 );
    for ( int i = 0; i < n; ++ i )
        set_port( bitmap, ports[ i ] );
}

int check_ports( uint16_t *ports, int n, uint8_t *bitmap )
{
    int correct = 1 << 16;
    for ( int port = 0; port < ( ( 1 << 16 ) ); ++ port )
        if ( !get_port( bitmap, port ) )
            -- correct;

    for ( int i = 0; i < n; ++ i )
        if ( !get_port( bitmap, ports[ i ] ) )
            -- correct;

    return correct == n;
}

int main( void )
{
    uint8_t bitmap[ ( 1 << 16 ) / 8 ];

    struct testserver server_data = { 0 };
    int serv_fd = prep_socket( &server_data.addr );
    server_data.pid = spawn_server(serv_fd);

    struct testserver server_data2 = { 0 };
    int serv_fd2 = prep_socket( &server_data2.addr );
    server_data2.pid = spawn_server( serv_fd2 );

    sched_yield();

    uint16_t ports_1[ 6 ]     = { 11, 25, 81, 82, 442, 55365 };
    uint16_t ports_out_1[ 6 ] = { 25 };
    uint16_t ports_2[ 6 ]     = { 11, 26, 81, 143, 442, 55365 };
    uint16_t ports_out_2[ 6 ] = { 143 };
    uint16_t ports_3[ 2 ] =
    {
        ntohs( server_data.addr.sin6_port ),
        ntohs( server_data2.addr.sin6_port )
    };

    /* Opened port 25 */
    set_ports( bitmap, ports_1, 6 );
    assert( map_ports( "relay.ip4.fi.muni.cz", AF_INET, bitmap ) == 0 );
    assert( check_ports( ports_out_1, 1, bitmap ) );

    /* Opened port 143 */
    set_ports( bitmap, ports_2, 6 );
    assert( map_ports( "aisa.ip6.fi.muni.cz", AF_INET6, bitmap ) == 0 );
    assert( check_ports( ports_out_2, 1, bitmap ) );

    /* Opened both ports */
    set_ports( bitmap, ports_3, 2 );
    assert( map_ports( "localhost", AF_INET6, bitmap ) == 0 );
    assert( check_ports( ports_3, 2, bitmap ) );

    assert( kill_and_reap( server_data.pid ) == 0 );
    assert( kill_and_reap( server_data2.pid ) == 0 );

    close_or_warn( serv_fd,  "server 1" );
    close_or_warn( serv_fd2, "server 2" );

    return 0;
}