/*

    mtr  --  a network diagnostic tool

    Copyright (C) 2016  Matt Kimball

    This program is free software; you can redistribute it and/or modify

    it under the terms of the GNU General Public License version 2 as

    published by the Free Software Foundation.

    This program is distributed in the hope that it will be useful,

    but WITHOUT ANY WARRANTY; without even the implied warranty of

    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License

    along with this program; if not, write to the Free Software

    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

*/

#include "probe.h"

#include <errno.h>

#include <stdio.h>

#include <winternl.h>

#include "protocols.h"

/*  Windows doesn't require any initialization at a privileged level  */

void init_net_state_privileged(

    struct net_state_t *net_state)

{

}

/*  Open the ICMP.DLL interface  */

void init_net_state(

    struct net_state_t *net_state)

{

    memset(net_state, 0, sizeof(struct net_state_t));

    net_state->platform.icmp4 = IcmpCreateFile();

    net_state->platform.icmp6 = Icmp6CreateFile();

    if (net_state->platform.icmp4 == INVALID_HANDLE_VALUE

        && net_state->platform.icmp6 == INVALID_HANDLE_VALUE) {

        fprintf(stderr, "Failure opening ICMP %d\n", GetLastError());

        exit(EXIT_FAILURE);

    }

}

/*

    If we succeeded at opening the ICMP file handle, we can

    assume that IP protocol version is supported.

*/

bool is_ip_version_supported(

    struct net_state_t *net_state,

    int ip_version)

{

    if (ip_version == 4) {

        return (net_state->platform.icmp4 != INVALID_HANDLE_VALUE);

    } else if (ip_version == 6) {

        return (net_state->platform.icmp6 != INVALID_HANDLE_VALUE);

    }

    return false;

}

/*  On Windows, we only support ICMP probes  */

bool is_protocol_supported(

    struct net_state_t * net_state,

    int protocol)

{

    if (protocol == IPPROTO_ICMP) {

        return true;

    }

    return false;

}

/*  Set the back pointer to the net_state when a probe is allocated  */

void platform_alloc_probe(

    struct net_state_t *net_state,

    struct probe_t *probe)

{

    probe->platform.net_state = net_state;

}

/*  Free the reply buffer when the probe is freed  */

void platform_free_probe(

    struct probe_t *probe)

{

    if (probe->platform.reply4) {

        free(probe->platform.reply4);

        probe->platform.reply4 = NULL;

    }

}

/*  Report a windows error code using a platform-independent error string  */

static

void report_win_error(

    int command_token,

    int err)

{

    /*  It could be that we got no reply because of timeout  */

    if (err == IP_REQ_TIMED_OUT || err == IP_SOURCE_QUENCH) {

        printf("%d no-reply\n", command_token);

    } else if (err == IP_DEST_HOST_UNREACHABLE

               || err == IP_DEST_PORT_UNREACHABLE

               || err == IP_DEST_PROT_UNREACHABLE

               || err == IP_DEST_NET_UNREACHABLE

               || err == IP_DEST_UNREACHABLE

               || err == IP_DEST_NO_ROUTE

               || err == IP_BAD_ROUTE || err == IP_BAD_DESTINATION) {

        printf("%d no-route\n", command_token);

    } else if (err == ERROR_INVALID_NETNAME) {

        printf("%d address-not-available\n", command_token);

    } else if (err == ERROR_INVALID_PARAMETER) {

        printf("%d invalid-argument\n", command_token);

    } else {

        printf("%d unexpected-error winerror %d\n", command_token, err);

    }

}

/*

    The overlapped I/O style completion routine to be called by

    Windows during an altertable wait when an ICMP probe has

    completed, either by reply, or by ICMP.DLL timeout.

*/

static

void WINAPI on_icmp_reply(

    PVOID context,

    PIO_STATUS_BLOCK status,

    ULONG reserved)

{

    struct probe_t *probe = (struct probe_t *) context;

    struct net_state_t *net_state = probe->platform.net_state;

    int icmp_type;

    int round_trip_us = 0;

    int reply_count;

    int reply_status = 0;

    struct sockaddr_storage remote_addr;

    struct sockaddr_in *remote_addr4;

    struct sockaddr_in6 *remote_addr6;

    ICMP_ECHO_REPLY32 *reply4;

    ICMPV6_ECHO_REPLY *reply6;

    if (probe->platform.ip_version == 6) {

        reply6 = probe->platform.reply6;

        reply_count = Icmp6ParseReplies(reply6, sizeof(ICMPV6_ECHO_REPLY));

        if (reply_count > 0) {

            reply_status = reply6->Status;

            /*  Unfortunately, ICMP.DLL only has millisecond precision  */

            round_trip_us = reply6->RoundTripTime * 1000;

            remote_addr6 = (struct sockaddr_in6 *) &remote_addr;

            remote_addr6->sin6_family = AF_INET6;

            remote_addr6->sin6_port = 0;

            remote_addr6->sin6_flowinfo = 0;

            memcpy(&remote_addr6->sin6_addr, reply6->AddressBits,

                   sizeof(struct in6_addr));

            remote_addr6->sin6_scope_id = 0;

        }

    } else {

        reply4 = probe->platform.reply4;

        reply_count = IcmpParseReplies(reply4, sizeof(ICMP_ECHO_REPLY));

        if (reply_count > 0) {

            reply_status = reply4->Status;

            /*  Unfortunately, ICMP.DLL only has millisecond precision  */

            round_trip_us = reply4->RoundTripTime * 1000;

            remote_addr4 = (struct sockaddr_in *) &remote_addr;

            remote_addr4->sin_family = AF_INET;

            remote_addr4->sin_port = 0;

            remote_addr4->sin_addr.s_addr = reply4->Address;

        }

    }

    if (reply_count == 0) {

        reply_status = GetLastError();

    }

    icmp_type = -1;

    if (reply_status == IP_SUCCESS) {

        icmp_type = ICMP_ECHOREPLY;

    } else if (reply_status == IP_TTL_EXPIRED_TRANSIT

               || reply_status == IP_TTL_EXPIRED_REASSEM) {

        icmp_type = ICMP_TIME_EXCEEDED;

    }

    if (icmp_type != -1) {

        /*  Record probe result  */

        respond_to_probe(net_state, probe, icmp_type,

                         &remote_addr, round_trip_us, 0, NULL);

    } else {

        report_win_error(probe->token, reply_status);

        free_probe(net_state, probe);

    }

}

/*  Use ICMP.DLL's send echo support to send a probe  */

static

void icmp_send_probe(

    struct net_state_t *net_state,

    struct probe_t *probe,

    const struct probe_param_t *param,

    struct sockaddr_storage *src_sockaddr,

    struct sockaddr_storage *dest_sockaddr,

    char *payload,

    int payload_size)

{

    IP_OPTION_INFORMATION option;

    DWORD timeout;

    DWORD send_result;

    int reply_size;

    int err;

    struct sockaddr_in *dest_sockaddr4;

    struct sockaddr_in6 *src_sockaddr6;

    struct sockaddr_in6 *dest_sockaddr6;

    if (param->timeout > 0) {

        timeout = 1000 * param->timeout;

    } else {

        /*

           IcmpSendEcho2 will return invalid argument on a timeout of 

           zero.  Our Unix implementation allows it.  Bump up the timeout

           to 1 millisecond.

         */

        timeout = 1;

    }

    memset(&option, 0, sizeof(IP_OPTION_INFORMATION32));

    option.Ttl = param->ttl;

    if (param->ip_version == 6) {

        reply_size = sizeof(ICMPV6_ECHO_REPLY) + payload_size;

    } else {

        reply_size = sizeof(ICMP_ECHO_REPLY32) + payload_size;

    }

    probe->platform.reply4 = malloc(reply_size);

    if (probe->platform.reply4 == NULL) {

        perror("failure to allocate reply buffer");

        exit(EXIT_FAILURE);

    }

    if (param->ip_version == 6) {

        src_sockaddr6 = (struct sockaddr_in6 *) src_sockaddr;

        dest_sockaddr6 = (struct sockaddr_in6 *) dest_sockaddr;

        send_result = Icmp6SendEcho2(net_state->platform.icmp6, NULL,

                                     (FARPROC) on_icmp_reply, probe,

                                     src_sockaddr6, dest_sockaddr6,

                                     payload, payload_size, &option,

                                     probe->platform.reply6, reply_size,

                                     timeout);

    } else {

        dest_sockaddr4 = (struct sockaddr_in *) dest_sockaddr;

        send_result = IcmpSendEcho2(net_state->platform.icmp4, NULL,

                                    (FARPROC) on_icmp_reply, probe,

                                    dest_sockaddr4->sin_addr.s_addr,

                                    payload, payload_size, &option,

                                    probe->platform.reply4, reply_size,

                                    timeout);

    }

    if (send_result == 0) {

        err = GetLastError();

        /*

           ERROR_IO_PENDING is expected for asynchronous probes,

           but any other error is unexpected.

         */

        if (err != ERROR_IO_PENDING) {

            report_win_error(probe->token, err);

            free_probe(net_state, probe);

        }

    }

}

/*  Fill the payload of the packet as specified by the probe parameters  */

static

int fill_payload(

    const struct probe_param_t *param,

    char *payload,

    int payload_buffer_size)

{

    int ip_icmp_size;

    int payload_size;

    if (param->ip_version == 6) {

        ip_icmp_size =

            sizeof(struct IP6Header) + sizeof(struct ICMPHeader);

    } else if (param->ip_version == 4) {

        ip_icmp_size = sizeof(struct IPHeader) + sizeof(struct ICMPHeader);

    } else {

        errno = EINVAL;

        return -1;

    }

    payload_size = param->packet_size - ip_icmp_size;

    if (payload_size < 0) {

        payload_size = 0;

    }

    if (payload_size > payload_buffer_size) {

        errno = EINVAL;

        return -1;

    }

    memset(payload, param->bit_pattern, payload_size);

    return payload_size;

}

/*  Decode the probe parameters and send a probe  */

void send_probe(

    struct net_state_t *net_state,

    const struct probe_param_t *param)

{

    struct probe_t *probe;

    struct sockaddr_storage dest_sockaddr;

    struct sockaddr_storage src_sockaddr;

    char payload[PACKET_BUFFER_SIZE];

    int payload_size;

    if (resolve_probe_addresses(param, &dest_sockaddr, &src_sockaddr)) {

        printf("%d invalid-argument\n", param->command_token);

        return;

    }

    probe = alloc_probe(net_state, param->command_token);

    if (probe == NULL) {

        printf("%d probes-exhausted\n", param->command_token);

        return;

    }

    probe->platform.ip_version = param->ip_version;

    payload_size = fill_payload(param, payload, PACKET_BUFFER_SIZE);

    if (payload_size < 0) {

        perror("Error construction packet");

        exit(EXIT_FAILURE);

    }

    icmp_send_probe(net_state, probe, param,

                    &src_sockaddr, &dest_sockaddr, payload, payload_size);

}

/*

    On Windows, an implementation of receive_replies is unnecessary, because,

    unlike Unix, replies are completed using Overlapped I/O during an

    alertable wait, and don't require explicit reads.

*/

void receive_replies(

    struct net_state_t *net_state)

{

}

/*

    On Windows, an implementation of check_probe_timeout is unnecesary because

    timeouts are managed by ICMP.DLL, including a call to the I/O completion

    routine when the time fully expires.

*/

void check_probe_timeouts(

    struct net_state_t *net_state)

{

}

/*

    As in the case of check_probe_timeout, getting the next probe timeout is

    unnecessary under Windows, as ICMP.DLL manages timeouts for us.

*/

bool get_next_probe_timeout(

    const struct net_state_t *net_state,

    struct timeval *timeout)

{

    return false;

}