/*
 * Copyright 2008-2010, Haiku, Inc. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *              Yin Qiu
 */


#include <errno.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>


#define DGRM_SIZE 1600 // ought to be large enough


struct pseudo_header {
        struct in_addr src;
        struct in_addr dst;     
        u_int8_t zero;
        u_int8_t protocol;
        u_int16_t len;
};


static u_int16_t
cksum(u_int16_t* buf, int nwords)
{
        u_int32_t sum;
        for (sum = 0; nwords > 0; nwords--)
                sum += *(buf++);
        sum = (sum >> 16) + (sum & 0xffff);
        sum += (sum >> 16);
        return ~sum;
}


static struct addrinfo*
host_serv(const char* host, const char* serv, int family, int socktype)
{
        struct addrinfo hints, *res;
        memset(&hints, 0, sizeof(struct addrinfo));
        hints.ai_flags = AI_CANONNAME;
        hints.ai_family = family;
        hints.ai_socktype = socktype;

        int n = getaddrinfo(host, serv, &hints, &res);
        if (n != 0)
                return NULL;

        return res;
}


int
main(int argc, char** argv)
{
        if (argc < 3) {
                printf("Usage: %s <ip-address> <port> [nbytes]\n", argv[0]);
                return 1;
        }

        size_t size;
        if (argc == 3)
                size = DGRM_SIZE;
        else
                size = atoi(argv[3]);

        if (size - sizeof(struct ip) - sizeof(struct udphdr) < 0) {
                fprintf(stderr, "Datagram size is too small\n");
                return 1;
        }

        struct addrinfo* ai = host_serv(argv[1], NULL, 0, 0);
        if (ai == NULL) {
                fprintf(stderr, "Cannot resolve %s\n", argv[1]);
                return 1;
        }

        printf("Using datagram size %zu\n", size);

        char* datagram = (char*)malloc(size);
        if (datagram == NULL) {
                fprintf(stderr, "Out of memory.\n");
                return 1;
        }

        memset(datagram, 0, size);

        struct ip* iph = (struct ip*)datagram;
        iph->ip_hl = 4;
        iph->ip_v = 4;
        iph->ip_tos = 0;
        iph->ip_len = htons(size);
        iph->ip_id = htonl(54321);
        iph->ip_off = 0;
        iph->ip_ttl = 255;
        iph->ip_p = IPPROTO_ICMP;
        iph->ip_sum = 0;
        iph->ip_src.s_addr = INADDR_ANY;
        iph->ip_dst.s_addr = ((struct sockaddr_in*)ai->ai_addr)->sin_addr.s_addr;
        // Calculates IP checksum
        iph->ip_sum = cksum((unsigned short*)datagram, iph->ip_len >> 1);

        struct pseudo_header pHeader; // used to calculate udp checksum

        struct udphdr* udp = (struct udphdr*)(datagram + sizeof(struct ip));
        udp->uh_sport = 0;
        udp->uh_dport = htons(atoi(argv[2]));
        udp->uh_ulen = htons(size - sizeof(struct ip));
        pHeader.src = iph->ip_src;
        pHeader.dst = iph->ip_dst;
        pHeader.zero = 0;
        pHeader.protocol = iph->ip_p;
        pHeader.len = udp->uh_ulen;
        udp->uh_sum = cksum((u_int16_t*)&pHeader, 3);
        
        // Send it via a raw socket
        int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_UDP);
        if (sockfd < 0) {
                fprintf(stderr, "Failed to create raw socket: %s\n", strerror(errno));
                free(datagram);
                return 1;
        }

        int one = 1;
        const int* val = &one;
        if (setsockopt(sockfd, IPPROTO_IP, IP_HDRINCL, val, sizeof(int)) != 0) {
                fprintf(stderr, "Failed to set IP_HDRINCL on socket: %s\n",
                        strerror(errno));
                free(datagram);
                return 1;
        }

        int bytesSent = sendto(sockfd, datagram, size, 0, ai->ai_addr,
                ai->ai_addrlen);
        if (bytesSent < 0) {
                fprintf(stderr, "Failed to send the datagram via the raw socket: %s\n",
                        strerror(errno));
                free(datagram);
                return 1;
        }

        printf("Sent %d bytes.\n", bytesSent);
        return 0;
}