/*
 *   Extremely Deficient TCP Server Connection Emulator (EDTSCE)
 *
 *   This code emulates the server side of a TCP connection, in an
 *    extremely deficient manner. It waits for a SYN on a port,
 *    responds with a SYN+ACK, and then sends several packets back to
 *    the connecting client. It neither stores nor ACKs any client
 *    data; it also does not handle retransmits (if a packet is lost,
 *    it's lost.) 
 *
 *   This code is nothing more than a "Hello world" in a few TCP packets.
 *    It is completely unsuitable as a basis for a TCP implemenation,
 *    which should strictly follow the TCP state machine structure. 
 *    Its inadequacies should be clear from a casual examination---and
 *    you should see where a state machine would look a lot more natural.
 *
 *   The point of this code is only to show that a simple stream of TCP
 *    packets can be made with just a few functions and a bit of logic 
 *    for handling SEQ/ACK numbers for the handshake and tear-down.   
 *
 *   Find bugs, get extra points.
 */ 

#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>

#include <netinet/ip.h>
#include <netinet/tcp.h>

#include "tcb.h"

#define INIT_SEQ 0xAAAABBBB

// global: dummy payload
char payload1[] = "Hello TCP\n";
int payload1_sz = sizeof(payload1);
char payload2[] = "Great to meet you!\n";
int payload2_sz = sizeof(payload2);
char payload3[] = "Nice weather we are having, but they say it will rain on Monday.\n";
int payload3_sz = sizeof(payload3);

// Note the large number of arguments. This is a sign that they should
//   be grouped into meaningful objects, and pointers to these objects should
//   be the arguments (cf. Linux's skbuff, tcp_skb_cb:
//     http://vger.kernel.org/~davem/skb.html
//     http://vger.kernel.org/~davem/tcp_skbcb.html)
void print_packet_info(char *buf, int n, struct iphdr *ip_hdr, struct tcphdr *tcp_hdr);

void make_and_send_synack( int sockfd, char *buf, int n,
			   struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb);
void make_and_send_payload_with_ack( int sockfd, char *buf, int n, char *payload, int payload_sz, 
			   struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb);
void make_and_send_fin( int sockfd, char *buf, int n, 
			struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb );
void make_and_send_finack( int sockfd, char *buf, int n,
			   struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb );

// All of the above make_and_send_* functions end with recomputing
//   checksums and a sendto(); this common part can be factored out.

uint16_t checksum (uint16_t *, int);
uint16_t
tcp4_checksum (char *packet, struct iphdr *ip, struct tcphdr *tcp, int tcp_payload_len);

int main(){
     int sockfd, n;
     socklen_t clilen = sizeof(struct sockaddr_in);    
     const int on = 1;
     struct sockaddr_in cliaddr;
     char buf[IP_MAXPACKET];

     tcb_t tcb;

     // create a raw socket for getting and sending packets
     sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_TCP);
     if (sockfd < 0){
          perror("sock:");
          exit(1);
     }

     // Set flag so socket expects us to provide IPv4 header.
     if (setsockopt (sockfd, IPPROTO_IP, IP_HDRINCL, &on, sizeof (on)) < 0) {
	  perror ("setsockopt() failed to set IP_HDRINCL ");
	  exit (1);
     }
     
     // This should be a state machine implementation, not just a few 
     //   lamely disparate condition checks! 
     while(1){
          printf(" before recvfrom\n");   
          n = recvfrom(sockfd, buf, IP_MAXPACKET, 0, (struct sockaddr *)&cliaddr, &clilen);
          printf(" rec'd %d bytes\n",n);
          
          struct iphdr *ip_hdr = (struct iphdr *) buf;
          struct tcphdr *tcp_hdr = (struct tcphdr *)((char *)ip_hdr + (4 * ip_hdr->ihl));

          // print whatever info helps to debug 
          print_packet_info(buf, n, ip_hdr, tcp_hdr);

	  if( tcp_hdr->syn & !tcp_hdr->ack ){ // respond with SYN+ACK
	       make_and_send_synack( sockfd, buf, n, ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
	  }
	  else if( tcp_hdr-> ack && ! tcp_hdr->syn && ! tcp_hdr->fin ){
	       // Is this my first ack? A proper check (such as a stateful firewall performs)
	       //  would include checking the socket quad (srcIP, dstIP, srcPort, dstPort).
	       if( ntohl(tcp_hdr->ack_seq) == INIT_SEQ+1 ){  // does this ack belong? 
		    printf( "Handshake complete.\n");
		    printf( "Sending payloads:\n");
		    make_and_send_payload_with_ack( sockfd, buf, n, payload1, payload1_sz,
						    ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
		    make_and_send_payload_with_ack( sockfd, buf, n, payload2, payload2_sz,
		    				    ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
		    make_and_send_payload_with_ack( sockfd, buf, n, payload3, payload3_sz,
						    ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
		    //... more payloads to launch...
		    make_and_send_fin( sockfd, buf, n, 
					ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
	       }
	       else{
		    printf( "ACK for %X\n", ntohl(tcp_hdr->ack_seq));
	       }
	  }
	  else if( tcp_hdr->rst ){
	       printf( "Got a RST, exiting\n" );
	       exit(0);
	  }
	  else if( tcp_hdr->fin ){
	       make_and_send_finack( sockfd, buf, n, 
				     ip_hdr, tcp_hdr, (struct sockaddr*)&cliaddr, &tcb);
	  }
          // and then what? A state machine must make sure no cases fall through!
     }
}

void make_and_send_synack( int sockfd, char *buf, int n,
			   struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb )
{
     char res[IP_MAXPACKET]; 
     memcpy(res, buf, n);
	    
     struct iphdr  *r_ip  = (struct iphdr*)  res;
     struct tcphdr *r_tcp = (struct tcphdr*) (res + 4*ip->ihl);
     r_ip->daddr = ip->saddr;
     r_ip->saddr = ip->daddr;
     r_tcp->source = tcp->dest;
     r_tcp->dest   = tcp->source;
     r_tcp->ack_seq = htonl( ntohl(tcp->seq)+1 ); // consume one seq number; all of this to just add 1
     r_tcp->seq = htonl( INIT_SEQ );
     r_tcp->ack = 1;    // flags | 0x02
     r_tcp->check = 0;  // checksum will include this field, so it should be zeroed out

     tcb->seq = INIT_SEQ+1;        // keep this in host order
     tcb->ack = ntohl( r_tcp->ack_seq );  // also host order
     
     // recompute the checksums
     r_ip->check = checksum( (uint16_t *) res, 4 * ip->ihl);
     r_tcp->check = tcp4_checksum (res, r_ip, r_tcp, n - 4*ip->ihl - 4*tcp->doff);
     
     // send packet
     if (sendto (sockfd, res, n, 0,
		    sin, sizeof (struct sockaddr_in)) < 0)  {
	  perror ("sendto() failed ");
	  exit (1);
     }

}

void make_and_send_finack( int sockfd, char *buf, int n,
			   struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb )
{
     char res[IP_MAXPACKET]; 
     memcpy(res, buf, n);
	    
     struct iphdr  *r_ip  = (struct iphdr*)  res;
     struct tcphdr *r_tcp = (struct tcphdr*) (res + 4*ip->ihl);
     r_ip->daddr = ip->saddr;
     r_ip->saddr = ip->daddr;
     r_tcp->source = tcp->dest;
     r_tcp->dest   = tcp->source;
     r_tcp->ack_seq = htonl( ntohl(tcp->seq)+1 ); // consume one seq number; all of this to just add 1
     r_tcp->seq = tcp->ack_seq;
     r_tcp->ack = 1;    // flags | 0x02
     r_tcp->check = 0;

     //recompute the checksums
     r_ip->check = checksum( (uint16_t *) res, 4 * ip->ihl);
     r_tcp->check = tcp4_checksum (res, r_ip, r_tcp, n - 4*ip->ihl - 4*tcp->doff);
     
     //send packet
     if (sendto (sockfd, res, n, 0,
		    sin, sizeof (struct sockaddr_in)) < 0)  {
	  perror ("sendto() failed ");
	  exit (1);
     }

}



void make_and_send_payload_with_ack( int sockfd, char *buf, int n, char *payload, int payload_sz,
				     struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb )
{
     char res[IP_MAXPACKET]; 
     memcpy(res, buf, n);
     memcpy(res + n, payload, payload_sz);
     
     struct iphdr  *r_ip  = (struct iphdr*)  res;
     struct tcphdr *r_tcp = (struct tcphdr*) (res + 4*ip->ihl);
     r_ip->daddr = ip->saddr;
     r_ip->saddr = ip->daddr;

     r_ip->tot_len += payload_sz;
     
     r_tcp->source = tcp->dest;
     r_tcp->dest   = tcp->source;
     r_tcp->ack_seq = tcp->seq;
     r_tcp->seq     = htonl( tcb->seq );
     r_tcp->ack = 1;    // flags | 0x02
     r_tcp->check = 0;

     tcb->seq += payload_sz;
     tcb->ack =  ntohl( tcp->ack_seq );
     
     //recompute the checksums
     r_ip->check = checksum( (uint16_t *) res, 4 * ip->ihl);
     r_tcp->check = tcp4_checksum (res, r_ip, r_tcp, n - 4*ip->ihl - 4*tcp->doff + payload_sz);
     
     //send packet
     if (sendto (sockfd, res, n+payload_sz, 0,
		    sin, sizeof (struct sockaddr_in)) < 0)  {
	  perror ("sendto() failed ");
	  exit (1);
     }

}

void make_and_send_fin( int sockfd, char *buf, int n, 
			struct iphdr *ip, struct tcphdr *tcp, struct sockaddr *sin, tcb_t *tcb )
{
     char res[IP_MAXPACKET]; 
     memcpy(res, buf, n);
     
     struct iphdr  *r_ip  = (struct iphdr*)  res;
     struct tcphdr *r_tcp = (struct tcphdr*) (res + 4*ip->ihl);
     r_ip->daddr = ip->saddr;
     r_ip->saddr = ip->daddr;

     r_tcp->source = tcp->dest;
     r_tcp->dest   = tcp->source;
     r_tcp->ack_seq = tcp->seq;
     r_tcp->seq     = htonl( tcb->seq );
     r_tcp->ack = 0;   // clear ACK, we just want the FIN
     r_tcp->fin = 1;    
     r_tcp->check = 0;

     //recompute the checksums
     r_ip->check = checksum( (uint16_t *) res, 4 * ip->ihl);
     r_tcp->check = tcp4_checksum (res, r_ip, r_tcp, n - 4*ip->ihl - 4*tcp->doff );
     
     //send packet
     if (sendto (sockfd, res, n, 0,
		    sin, sizeof (struct sockaddr_in)) < 0)  {
	  perror ("sendto() failed ");
	  exit (1);
     }
}

// Build IPv4 TCP pseudo-header and call checksum function.
//   This function should really be zero-copy, but it's left as an exercise
//     (or look it up in a TCP implementation).
uint16_t
tcp4_checksum (char *packet, struct iphdr *ip, struct tcphdr *tcp, int tcp_payload_len)
{
     char buf[IP_MAXPACKET];
     char *ptr;
     int chksumlen;

     // ptr points to beginning of buffer buf
     ptr = &buf[0];
     
     struct pseudo_tcp_header
     {
	  in_addr_t saddr, daddr;
	  u_char reserved;
	  u_char protocol;
	  u_short tcp_size;
     } ps_head;

     uint16_t header_len = 4*tcp->doff;
     
     ps_head.saddr    = ip->saddr;
     ps_head.daddr    = ip->daddr;
     ps_head.reserved = 0;
     ps_head.protocol = ip->protocol;
     ps_head.tcp_size = htons(header_len + tcp_payload_len);
     
     memcpy(ptr, &ps_head, sizeof(struct pseudo_tcp_header));
     ptr += sizeof(struct pseudo_tcp_header);
     
     // now copy TCP header
     memcpy (ptr, tcp, header_len);   // this must be result's TCP header!
     ptr += header_len;

     // ... and the TCP payload. 
     memcpy (ptr, packet + ip->ihl*4 + header_len, tcp_payload_len);
     ptr += tcp_payload_len;

     chksumlen = ptr - buf;
     
     return checksum ((uint16_t *) buf, chksumlen);
}

// Computing the internet checksum (RFC 1071).
uint16_t
checksum (uint16_t *addr, int len)
{
     int count = len;
     register uint32_t sum = 0;
     uint16_t answer = 0;

     // Sum up 2-byte values until none or only one byte left.
     while (count > 1) {
	  sum += *(addr++);
	  count -= 2;
     }
     
     // Add left-over byte, if any.
     if (count > 0) {
	  sum += *(uint8_t *) addr;
     }
     
     // Fold 32-bit sum into 16 bits; we lose information by doing this,
     // increasing the chances of a collision.
     // sum = (lower 16 bits) + (upper 16 bits shifted right 16 bits)
     while (sum >> 16) {
	  sum = (sum & 0xffff) + (sum >> 16);
     }
     
     // Checksum is one's compliment of sum.
     answer = ~sum;
     
     return (answer);
}

void print_packet_info(char *buf, int n, struct iphdr *ip, struct tcphdr *tcp)
{
     int i;
     printf("IP header is %d bytes.\n", ip->ihl*4);
          
     for (i = 0; i < n; i++) {
          printf("%02X%s", (uint8_t)buf[i], (i + 1)%16 ? " " : "\n");
     }
     printf("\n");
     
     // Make like a tiny tcpdump. Note that this ignores other flags & combinations!
     printf("TCP sport=%d, dport=%d ", ntohs(tcp->source), ntohs(tcp->dest));
     if (tcp->syn && ! tcp->ack)
          printf("SYN");
     else if( tcp->syn && tcp->ack )
          printf("SYN+ACK");
     else if( tcp->ack && !tcp->fin )
          printf("ACK");
     else if( tcp->ack && tcp->fin )
          printf("FIN+ACK");
     // and then what? No cases should fall through!
     puts("");
}