/*
 *   Somewhat Deficient TCP Server Connection Emulator (SDTSCE)
 *
 *   This code emulates the server side of a TCP connection, in a
 *    somewhat deficient manner. It waits for a SYN on a port,
 *    responds with a SYN+ACK, and then sends packets representing a
 *    fixed payload response (drawn from a real Apache web server)
 *    back to the connecting client. It neither stores nor ACKs any
 *    client data. 
 *
 *   This version handles retransmissions: for every packet sent, a 
 *    timer is started. If the packet is not ACK-ed within 1 second,
 *    it is retransmitted and the retransmission timer is reset.
 *    Retransmissions will continue until the payload of the packet
 *    is ACK-ed (including cummulatively). 
 *
 *   This code statically pre-allocates the memory for the packets 
 *    in flight and their timers. This is wasteful, and would break
 *    for larger files or smaller MSS of TCP packets.
 *
 *   Find bugs, get extra points.
 */ 

/*
 *  Most code cut, only the time-out logic left.
 */

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

#include <sys/types.h>  // open
#include <sys/stat.h> 
#include <fcntl.h>

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

#include <poll.h>          // poll
#include <sys/timerfd.h>   // timers

#include <sys/socket.h>    // inet_addr
#include <netinet/in.h>
#include <arpa/inet.h>

#include "tcb.h"
#include "pkt.h"

// Keep the initial Seq number the same for easier debugging
#define INIT_SEQ 0xAAAABBBB

// Preallocate space for this many packets in flight and their respective timers
#define MAX_POLLED 200

void set_rto_timer(pkt_t*, int i);
void reset_rto_timer(pkt_t*, int i);
void clear_rto_timer(pkt_t *pkt, int i);

int tcp_raw_socket();
void recv_packet(pkt_t *pkt);

// global timers
struct pollfd polled[MAX_POLLED];
int num_polled = 1; // active polling is from num_acked to num_polled-1
int num_acked  = 0; 

pkt_t pkt_in_flight[MAX_POLLED]; // 0 is not used 

int main(){
     int sockfd;
     int i;
     
     sockfd = tcp_raw_socket(); 

     // Polled will always start with the raw socket, then use 1..MAX_POLLED-1 for timers
     polled[0].fd = sockfd;
     polled[0].events = POLLIN;
     for( i=1; i < MAX_POLLED; i++ ){
	  polled[i].fd     = -1;   // poll ignores negative fds 
	  polled[i].events = POLLIN;
     }
	  
     // State machine implementation. Nothing is allowed to fall through without
     //   at least a message
     while(1){
	  recv_packet(&pkt); // blocking; gets a packet, handles a timer, or exits.
	  if( pkt.invalid ) // timeout in poll   
	       continue;

	  if( ! matches_tcb( &pkt, &tcb) ){
	       // a real TCP stack would send a RST here
	       fprintf(stderr, "Packet doesn't match the TCB, ignoring packet:\n");
	       print_packet_summary(stderr, &pkt);
	       continue;
	  }
	  
	  if( IS_SYN(&pkt) && tcb.state == TCP_LISTEN ){ // send SYN+ACK
	       make_and_send_synack( &pkt, &tcb);
	       tcb.state = TCP_SYN_RECV;
	  }
          // <SKIPPED>
	  }
	  else {  // There must a 
	       fprintf(stderr, "Wrong state %d (%s) for packet, ignoring:\n",
		       tcb.state, tcp_states[tcb.state] );
	       print_packet_info(stderr, &pkt);
	       // A real stack would send a RST here; instead, we just go to LISTEN
	       tcb.state = TCP_LISTEN;
               // clean-up skipped
	  }
     }
}

// When sending a packet, start a timer on it! E.g.:
// <SKIPPED>
     make_and_send_payload_with_ack( pkt, &pkt_in_flight[i], 
   				     payload, payload_sz, tcb );
     set_rto_timer(&pkt_in_flight[i], i);

// <SKIPPED>
// When you receive an ack, find and clear the timer for that packet
     clear_rto_timer(pkt, i); 

// Create a raw socket
int tcp_raw_socket()
{
     const int on = 1;
     int sockfd;
     
     // 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);
     }

     return sockfd;
}

// Listens on global polled[], takes action on expired timers 
void recv_packet(pkt_t *pkt)
{
     int n, rv, i;
     socklen_t clilen = sizeof(struct sockaddr_in);    

     int sockfd = polled[0].fd;
     char *buf = pkt->buf;

     pkt->invalid = 0;
     
     printf(" at poll\n");
     rv = poll( polled, num_polled, 10000 );
     if (rv == -1) {
	  perror("poll:"); // error occurred in poll()
     } else if (rv == 0) {
	  fprintf(stderr, "Timeout occurred!  No data after 10 seconds.\n");
	  pkt->invalid = 1;
	  return;
     } else {
	  // We got a packet
	  if (polled[0].revents & POLLIN) {
	       n = recvfrom(sockfd, buf, IP_MAXPACKET, 0,
			    (struct sockaddr *)&(pkt->sin), &clilen);
	       
	       if( n < 0 ){
		    perror("recv: ");
		    exit(1);
	       }
	       else if( n == 0 ){
		    fprintf( stderr, "Raw socket recv returned 0; this should not happen\n");
		    exit(1);
	       }
	       else
		    printf(" rec'd %d bytes\n", n);
	       
               // Build packet pkt
               // <SKIPPED>
	       
	       // print whatever info helps to debug 
	       print_packet_info(stdout, pkt);
	       
	  } // done with a packet, now handle timers

	  for( i=1; i < num_polled; i++ ){  // we can actually start with num_acked
	       if(polled[i].revents & POLLIN) {
		    uint64_t exp; // must read the expired timer
		    int s = read(polled[i].fd, &exp, sizeof(uint64_t));
		    if (s != sizeof(uint64_t)){
			 perror("timer read: ");
			 exit(1);
		    }
		    fprintf(stderr, "Timeout of TRO %d: %lu\n", i, exp);
		    // restransmit timed-out packet
		    retransmit_on_timeout(&pkt_in_flight[i], i);
	       }
	  }  // done with timers
     } // end poll
}

// retransmit packet and reset its timer
void retransmit_on_timeout( pkt_t *pkt, int i )
{
     // resend packet with sendto()
     int sockfd = pkt->sockfd;
     struct sockaddr *sin = (struct sockaddr*)&(pkt->sin);

     fprintf(stderr, "Retransmitting pkt %d ack %X  (sock %d)\n", i, pkt->ack, sockfd );
     
     if (sendto (sockfd, pkt->buf, pkt->len, 0,
		 sin, sizeof (struct sockaddr_in)) < 0)  {
	  perror ("sendto() failed on retransmission");
	  exit (1);
     }

     reset_rto_timer( pkt, i );
}

// set as timer i in polled
void set_rto_timer(pkt_t *pkt, int i)
{
     struct itimerspec itspec;
     int tfd = timerfd_create(CLOCK_MONOTONIC, 0);
     if (tfd == -1){
	  perror("timerfd_create: ");
	  exit(1);
     }

     itspec.it_value.tv_sec = 1;          // expire in 1 sec
     itspec.it_value.tv_nsec = 0;
     itspec.it_interval.tv_sec = 0;       // ...just once
     itspec.it_interval.tv_nsec = 0;

     if (timerfd_settime(tfd, 0, &itspec, NULL) == -1){
	  perror("timerfd_settime: ");
	  exit(1);
     }

     // now activate in polled
     polled[i].fd = tfd;
     polled[i].events = POLLIN;
     num_polled = i+1;  // now i is polled

     pkt->rto_fd = tfd;
}

// reset as timer i in polled
void reset_rto_timer(pkt_t *pkt, int i)
{
     struct itimerspec itspec;
     int tfd = pkt->rto_fd;

     itspec.it_value.tv_sec = 1;          // expire in 1 sec
     itspec.it_value.tv_nsec = 0;
     itspec.it_interval.tv_sec = 0;       // ...just once
     itspec.it_interval.tv_nsec = 0;

     if (timerfd_settime(tfd, 0, &itspec, NULL) == -1){
	  perror("reset RTO timerfd_settime: ");
	  exit(1);
     }
}

// clear as timer i in polled
void clear_rto_timer(pkt_t *pkt, int i)
{
     int tfd = pkt->rto_fd;

     // de-activate in polled[]. poll() ignores negative fds
     polled[i].fd = -tfd; // or -1. I left -tfd for debugging. 
     close(tfd);  // this will dellocated the timer, otherwise you'll run out of FDs
}