#include "kni_utils.h"

uint16_t kni_ip_checksum(const void *buf, size_t hdr_len){
	unsigned long sum = 0;
	const uint16_t *ip1; 
	ip1 = (const uint16_t *)buf;
	while(hdr_len > 1){
		sum += *ip1++;
		if(sum & 0x80000000){
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		hdr_len -= 2;
	}
	while(sum >> 16){
		sum = (sum & 0xFFFF) + (sum >> 16); 
	}
	return (~sum);
}

uint16_t kni_tcp_checksum(const void *_buf, size_t len, in_addr_t src_addr, in_addr_t dest_addr){
    const uint16_t *buf = (u_int16_t *)_buf;
	uint16_t *ip_src=(uint16_t *)&src_addr, *ip_dst=(uint16_t *)&dest_addr;
	uint32_t sum;
	size_t length=len;
	// Calculate the sum                                            
	sum = 0;
	while(len > 1){
		sum += *buf++;
		if (sum & 0x80000000){
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		len -= 2;
	}
	if(len & 1){
		// Add the padding if the packet lenght is odd          
		sum += *((uint8_t *)buf); 
	}
	// Add the pseudo-header                                        
	sum += *(ip_src++);
	sum += *ip_src;
	sum += *(ip_dst++);
	sum += *ip_dst;
	sum += htons(IPPROTO_TCP);
	sum += htons(length); 
	// Add the carries                                              
	while(sum >> 16){
		sum = (sum & 0xFFFF) + (sum >> 16);
	}
	// Return the one's complement of sum                           
	return ((uint16_t)(~sum));
}

uint16_t kni_udp_checksum(const void *_buf, size_t len, in_addr_t src_addr, in_addr_t dest_addr){
    const uint16_t *buf = (u_int16_t *)_buf;
	uint16_t *ip_src=(u_int16_t *)&src_addr, *ip_dst=(u_int16_t *)&dest_addr;
	uint32_t sum;
	size_t length=len;
	// Calculate the sum                                            
	sum = 0;
	while(len > 1){
		sum += *buf++;
		if (sum & 0x80000000){
			sum = (sum & 0xFFFF) + (sum >> 16);
		}
		len -= 2;
	}
	if(len & 1){
		// Add the padding if the packet lenght is odd          
		sum += *((uint8_t *)buf);
	}

	// Add the pseudo-header                                        
	sum += *(ip_src++);
	sum += *ip_src;
	sum += *(ip_dst++);
	sum += *ip_dst;
	sum += htons(IPPROTO_UDP);
	sum += htons(length);

	// Add the carries                                              
	while(sum >> 16){
		sum = (sum & 0xFFFF) + (sum >> 16);
	}
	// Return the one's complement of sum                           
	return ( (uint16_t)(~sum)  );
 }


struct kni_tcpopt_info* kni_get_tcpopt(struct tcphdr* tcphdr,int tcphdr_len){
    struct kni_tcpopt_info* tcpopt = (struct kni_tcpopt_info*)ALLOC(struct kni_tcpopt_info, 1);
	tcpopt->mss = KNI_DEFAULT_MSS;
	tcpopt->wscale = KNI_DEFAULT_WINSCLE;

	const unsigned char *ptr = ((const unsigned char*)tcphdr + 20);
	int length = tcphdr_len - 20;

	while (length > 0){
		int opcode = *ptr++;
		int opsize;
		switch (opcode){
            case TCPOPT_EOL:
                return tcpopt;
            case TCPOPT_NOP:	/* Ref: RFC 793 section 3.1 */
                length--;
                continue;
            default:
                opsize = *ptr++;
                if (opsize < 2) /* "silly options" */
                    return tcpopt;
                if (opsize > length)
                    return tcpopt;	/* don't parse partial options */
                switch (opcode){
                    case TCPOPT_MAXSEG:
                        if (opsize == TCPOLEN_MAXSEG){
                            uint16_t in_mss = *(uint16_t *)ptr;
                            if(in_mss){ 
                                tcpopt->mss = ntohs(in_mss);
                            }
                        }
                        break;
                        
                    case TCPOPT_WINDOW:
                        if (opsize == TCPOLEN_WINDOW){
                            uint8_t snd_wscale = *(uint8_t *)ptr;
                            // rfc7323 page9: Thus, the shift count MUST be limited to 14 (which allows windows of 2^30 = 1 GiB). 
                            // If a Window Scale option is received with a shift.cnt value larger than 14, 
                            // the TCP SHOULD log the error but MUST use 14 instead of the  specified value.  */
                            tcpopt->wscale = snd_wscale;
                            if(tcpopt->wscale > 14){
                                tcpopt->wscale = 14;
                            }
                            //*wscale_perm=1;
                        }
                        break;
                    case TCPOPT_TIMESTAMP:
                        if ((opsize == TCPOLEN_TIMESTAMP)){
                            tcpopt->ts = 1;
                        }
                        break;
                    case TCPOPT_SACK_PERMITTED:
                        if (opsize == TCPOLEN_SACK_PERMITTED){
                            tcpopt->sack = 1;
                        }
                        break;
                }
                ptr += opsize-2;
                length -= opsize;
        }
	}
	return tcpopt;
}