/** @file  wlan_exp_ip_udp_ip_udp.c
 *  @brief Mango wlan_exp IP/UDP Library (IP/UDP/ARP/ICMP Processing)
 *
 *  @copyright Copyright 2014-2019, Mango Communications. All rights reserved.
 *          Distributed under the Mango Reference Design license (https://mangocomm.com/802.11/license)
 */

/***************************** Include Files *********************************/
#include "wlan_mac_high_sw_config.h"

#if WLAN_SW_CONFIG_ENABLE_WLAN_EXP

// Xilinx / Standard library includes
#include <xparameters.h>
#include <xstatus.h>
#include <xil_io.h>
#include <stdlib.h>
#include <stdio.h>

// Mango wlan_exp IP/UDP Library includes
#include "wlan_exp_ip_udp.h"
#include "wlan_exp_ip_udp_socket.h"
#include "wlan_exp_ip_udp_arp.h"

// Mango Framework includes
#include "wlan_high_types.h"
#include "wlan_mac_queue.h"
#include "wlan_mac_common.h"
#include "wlan_exp_transport.h"

/*********************** Global Variable Definitions *************************/

u16 ipv4_id_counter = 0;
static u8 gl_ip_addr[IP_ADDR_LEN];  // Ethernet device IP address
static u8 gl_hw_addr[ETH_ADDR_LEN]; // Ethernet device MAC address

/*************************** Function Prototypes *****************************/

void icmp_echo_reply(eth_rx_queue_buffer_t* eth_rx_queue_buffer, eth_tx_queue_buffer_t* eth_tx_queue_buffer);
int _set_ip_length(u8* pkt, u32 total_length);
int _set_udp_length(u8* pkt, u32 total_length);


/******************************** Functions **********************************/

/*****************************************************************************/
/**
 * Initialize the IP/UDP Library
 *
 * This function will initialize all subsystems within the library:
 *   - Global Ethernet structures
 *   - Socket data structures
 *   - ARP cache
 *   - IP V4 global structures (ie ID counter)
 *
 * @param   None
 *
 * @return  int              - WLAN_SUCCESS (cannot fail)
 *
 ******************************************************************************/
int ip_udp_init(  u8* hw_addr, u8* ip_addr) {

    // Initialize the sockets
    socket_init();

    // Initialize the ARP cache
    arp_init_cache();

    // Initialize the IP v4 global structures
    ipv4_id_counter = 0;

    eth_set_hw_addr(hw_addr);
    eth_set_ip_addr(ip_addr);

    // Return success
    return WLAN_SUCCESS;
}


/*****************************************************************************/
/**
 * Create Ethernet / IP / UDP Headers
 *
 *****************************************************************************/

int ip_udp_template_init(u8* pkt, u32 port, sockaddr_in_t* to) {

    int status;
    sockaddr_in_t from;

    // Get current open socket details
    int socket_index = socket_find_index_by_port(port);

    if (socket_index != SOCKET_INVALID_SOCKET) {
    	from = socket_get(socket_index);
    } else {
    	return WLAN_FAILURE;
    }

    u8 dest_hw_addr[ETH_ADDR_LEN];
    u8 src_hw_addr[ETH_ADDR_LEN];
    u8* src_ip_addr = from.sin_addr.s_addr;
    u32 src_port = from.sin_port;
    u8* dest_ip_addr = to->sin_addr.s_addr;
    u16 dest_port = to->sin_port;

    ethernet_header_t* eth_hdr = (ethernet_header_t*)pkt;
    ipv4_header_t* ip_hdr = (ipv4_header_t*)( pkt + sizeof(ethernet_header_t) );
    udp_header_t* udp_hdr = (udp_header_t*)( pkt + sizeof(ethernet_header_t) + sizeof(ipv4_header_t) );

    eth_get_hw_addr(src_hw_addr);

    // Look up destination HW address in ARP table
    status = arp_get_hw_addr(dest_hw_addr, dest_ip_addr);

    if (status != WLAN_SUCCESS) {
        return WLAN_FAILURE;
    }

    // Update the Ethernet header
    eth_hdr_init(eth_hdr,
    			 dest_hw_addr,
    			 src_hw_addr);

    // Update the IPv4 header
    ip_hdr_init(ip_hdr,
    			src_ip_addr,
    			dest_ip_addr,
    			IP_PROTOCOL_UDP);


    // Update the UDP header
    udp_hdr_init(udp_hdr,
    			 src_port,
    			 dest_port);


    // Send the Ethernet frame using the socket header
    return (sizeof(ethernet_header_t) + sizeof(ipv4_header_t) + sizeof(udp_header_t));
}

int wlan_exp_ip_udp_set_length(u8* pkt, u32 total_length){
	// total_length refers to the number of bytes that will
	// go over the wire starting with the Ethernet header

	if(pkt==NULL) return WLAN_FAILURE;

	_set_ip_length(pkt, total_length);
	_set_udp_length(pkt, total_length);

    return WLAN_SUCCESS;
}

int _set_ip_length(u8* pkt, u32 total_length){
	// total_length refers to the number of bytes that will
	// go over the wire starting with the Ethernet header

	if(pkt==NULL) return WLAN_FAILURE;

    ipv4_header_t* ip_hdr = (ipv4_header_t*)( pkt + sizeof(ethernet_header_t) );

    // Set IP length
    ip_hdr->total_length = Xil_Htons(total_length - sizeof(ethernet_header_t) );

    // Set IP checksum
    ip_hdr->header_checksum  = Xil_Htons(ipv4_compute_checksum((u8*)ip_hdr, sizeof(ipv4_header_t)));

    return WLAN_SUCCESS;

}
int _set_udp_length(u8* pkt, u32 total_length){
	// total_length refers to the number of bytes that will
	// go over the wire starting with the Ethernet header

	if(pkt==NULL) return WLAN_FAILURE;

    udp_header_t* udp_hdr = (udp_header_t*)( pkt + sizeof(ethernet_header_t) + sizeof(ipv4_header_t) );

    // Set UDP length
    udp_hdr->length = Xil_Htons(total_length - sizeof(ethernet_header_t) - sizeof(ipv4_header_t));

    // If we wanted to use UDP checksums, here is where we would set it

    return WLAN_SUCCESS;
}

/**********************************************************************************************************************/
/**
 * @brief IP Functions 
 *
 **********************************************************************************************************************/


/*****************************************************************************/
/**
 * Process the IP packet
 *
 * @param   eth_rx_queue_buffer - Rx Ethernet packet
 * @param   eth_tx_queue_buffer - Buffer for Tx Ethernet packet to be filled in (if necessary)
 *
 * @return  int - Number of bytes of data in the processed packet; 0 if the packet could not be processed
 *
 * @note    This function assumes that both Ethernet device and buffer are valid.
 *
 *****************************************************************************/
int ipv4_process_packet(eth_rx_queue_buffer_t* eth_rx_queue_buffer, eth_tx_queue_buffer_t* eth_tx_queue_buffer) {

    ipv4_header_t* header;
    ethernet_header_t* eth_header;
    u8 my_ip_addr[IP_ADDR_LEN];

    header = (ipv4_header_t*)((void*)(eth_rx_queue_buffer->pkt) + sizeof(ethernet_header_t));

    u32 addr_check = 0;
    u8* src_ip_addr = header->src_ip_addr;
    u8* dest_ip_addr = header->dest_ip_addr;

    eth_get_ip_addr(my_ip_addr);

    // Get the Ethernet header so we have the source MAC address for the ARP cache
    eth_header = (ethernet_header_t*)eth_rx_queue_buffer->pkt;

    // Check the address of the IP packet
    //     - If the node has not been initialized (eg the node address is 10.0.0.0), then accept broadcast packets from 10.0.X.255
    //     - If the node has been initialized, then accept unicast packets and broadcast packets on the given subnet
    //
    //
    // !!! TBD !!! - Future addition:  The address check should really be more configurable (ie it should be a callback
    //                                 that can be set by the application that uses the library).
    //
    //
    if (my_ip_addr[3] == 0) {

        // Accept broadcast packets from 10.0.X.255
        if ((my_ip_addr[0] == dest_ip_addr[0]) && (my_ip_addr[1] == dest_ip_addr[1]) && (dest_ip_addr[3] == 255)) {
             addr_check = 1;
        }
    } else {
        // Accept unicast packets and broadcast packets on the given subnet
        if ((my_ip_addr[0] == dest_ip_addr[0]) &&
            (my_ip_addr[1] == dest_ip_addr[1]) &&
            (my_ip_addr[2] == dest_ip_addr[2]) &&
            ((my_ip_addr[3] == dest_ip_addr[3]) || (dest_ip_addr[3] == 255))) {
            addr_check = 1;
        }
    }

    //
    // 
    // !!! TBD !!! - Future consideration:  Additional packet checks - Length & Checksum
    //     u16            packet_length  = packet->length;
    //     u16            ip_length      = Xil_Ntohs(header->total_length);
    //
    //
    
    if (addr_check == 1) {
    
        // The Xilinx Ethernet / DMA hardware does not support fragmented Ethernet packets.  However, the 
        // library will still pass the first fragment of a packet up to the higher level transport for
        // processing so that the host that sent the fragmented packet does not have a transport timeout 
        // (This is important when trying to determine the maximum packet size supported by the transport).  
        // If this behavior needs to change the below code will cause the Mango wlan_exp IP/UDP Library to discard
        // packet fragments.
        // 
        // The 'fragment offset field' is 16 bits (see http://en.wikipedia.org/wiki/IPv4#Packet_structure 
        // for more information).  The 'DF' flag can be legitimately set to '1' so we need to mask that 
        // bit before we decide if we can discard the packet.  This means that the frag_off field can 
        // have valid values of either 0x0000 or 0x4000 (big endian).  However, we have to convert to 
        // little endian so the valid values are 0x0000 and 0x0040 (ie byte swapped).
        // 
        // if ((header->fragment_offset & 0xFFBF) != 0x0000) {
        //     xil_printf("ERROR:  Library does not support fragmented packets.\n");
        //     return 0;
        // }

        // Update ARP table (Maps IP address to MAC address)
        arp_update_cache(eth_header->src_mac_addr, src_ip_addr);
        
        // Process the IP packet
        switch (header->protocol) {
            // UDP packets
            case IP_PROTOCOL_UDP:
                return udp_process_packet(eth_rx_queue_buffer);
            break;
            
            // ICMP packets
            case IP_PROTOCOL_ICMP:
                return icmp_process_packet(eth_rx_queue_buffer, eth_tx_queue_buffer);
            break;
            
            // If a packet has made it here, the it is destined for the node but cannot be processed
            // by the library.  Therefore, we need to print an error message.
            default:
                xil_printf("ERROR:  Unknown IP protocol:  %d\n", header->protocol);
            break;
        }
    }

    return 0;
}

void eth_hdr_init(ethernet_header_t* eth_hdr,
			 	 u8* dest_hw_addr,
			 	 u8* src_hw_addr){

	memcpy(eth_hdr->dest_mac_addr, dest_hw_addr, ETH_ADDR_LEN);
	memcpy(eth_hdr->src_mac_addr, src_hw_addr, ETH_ADDR_LEN);
	eth_hdr->ethertype = Xil_Htons(ETHERTYPE_IP_V4);
}

void ip_hdr_init(ipv4_header_t* ip_hdr,
				 u8* src_ip_addr,
				 u8* dest_ip_addr,
				 u8 protocol){

	memcpy(ip_hdr->src_ip_addr, src_ip_addr, IP_ADDR_LEN);
	ip_hdr->version_ihl      = (IP_VERSION_4 << 4) +  IP_HEADER_LEN;
	ip_hdr->dscp_ecn         = (IP_DSCP_CS0 << 2) + IP_ECN_NON_ECT;
	ip_hdr->fragment_offset  = IP_NO_FRAGMENTATION;
	ip_hdr->ttl              = IP_DEFAULT_TTL;
	ip_hdr->total_length     = 0; // Lengths are set after the full packet contents are completed
	ip_hdr->identification   = Xil_Htons(ipv4_id_counter++);
	ip_hdr->protocol         = protocol;
	ip_hdr->header_checksum  = 0; // Checksum is set after the full packet contents are completed
	memcpy(ip_hdr->dest_ip_addr, dest_ip_addr, IP_ADDR_LEN);
}

// Update the UDP header
void udp_hdr_init(udp_header_t* udp_hdr,
			 	 u16 src_port,
			 	 u16 dest_port){

	udp_hdr->src_port  = Xil_Htons(src_port);
	udp_hdr->dest_port = Xil_Htons(dest_port);
	udp_hdr->length    = 0; // Lengths are set after the full packet contents are completed
	udp_hdr->checksum  = UDP_NO_CHECKSUM;
}


/*****************************************************************************/
/**
 * Compute IP Checksum
 *
 * The ones' complement of the ones' complement sum of the data's 16-bit words
 *
 * @param   data            - Pointer to the data words
 * @param   size            - Size of the data to use
 *
 * @return  u16             - Checksum value
 *
 * @note    IP Checksum Algorithm:  http://en.wikipedia.org/wiki/ipv4_header_t_checksum
 *
 *****************************************************************************/
u16 ipv4_compute_checksum(u8 * data, u32 size) {

    u32   i;
    u32   sum  = 0;
    u16   word = 0;
   
    // Sum all 16-bit words in the header (big-endian)
    for (i = 0; i < size; i = i + 2) {
        word = ((data[i] << 8) & 0xFF00) + (data[i+1] & 0x00FF);
        sum  = sum + ((u32) word);
    }

    // 1's complement 16-bit sum, formed by "end around carry" of 32-bit 2's complement sum
    sum = ((sum & 0xFFFF0000) >> 16) + (sum & 0x0000FFFF);

    // Return the 1's complement of 1's complement 16-bit sum
    return (~sum);
}

 


/**********************************************************************************************************************/
/**
 * @brief UDP Functions 
 *
 **********************************************************************************************************************/


/*****************************************************************************/
/**
 * Process the UDP packet
 *
 * @param   eth_rx_queue_buffer - Ethernet packet
 *
 * @return  int         - Number of bytes of data in the UDP packet; 0 if the packet could not be processed
 *
 * @note    This function assumes that both Ethernet device and buffer are valid.
 *
 *****************************************************************************/
int udp_process_packet(eth_rx_queue_buffer_t* eth_rx_queue_buffer) {

    ipv4_header_t* ip_hdr;
    udp_header_t* udp_hdr;
    int socket_index;
    wlan_exp_eth_rx_queue_buffer_t* wlan_exp_eth_rx_queue_buffer;

    ip_hdr = (ipv4_header_t*)((void*)(eth_rx_queue_buffer->pkt) + sizeof(ethernet_header_t));
    udp_hdr = (udp_header_t*)((void*)ip_hdr + 4*((u8)(ip_hdr->version_ihl) & 0xF));

    
    u32            port_check     = 0;
    u16            dest_port      = Xil_Ntohs(udp_hdr->dest_port);

    // See if there is a socket that corresponds to the UDP packet
    //     - Check all open sockets to see if one matches the port / eth_dev_num
    // 
    socket_index = socket_find_index_by_port(dest_port);
    if (socket_index != SOCKET_INVALID_SOCKET) {
        port_check = 1;

        // Copy socket details into eth_rx_queue_buffer for future wlan_exp processing
        wlan_exp_eth_rx_queue_buffer = (wlan_exp_eth_rx_queue_buffer_t*)eth_rx_queue_buffer;
    	wlan_exp_eth_rx_queue_buffer->rx_from.sin_family = AF_INET;
    	wlan_exp_eth_rx_queue_buffer->rx_from.sin_port = Xil_Ntohs(udp_hdr->src_port);
    	memcpy(wlan_exp_eth_rx_queue_buffer->rx_from.sin_addr.s_addr, ip_hdr->src_ip_addr, IP_ADDR_LEN);
    }
    
    //
    // 
    // !!! TBD !!! - Future consideration:  Additional packet checks - Length & Checksum
    //     u16            packet_length  = packet->length;
    //     u16            udp_length     = Xil_Ntohs(header->length);
    //
    // 

    if (port_check == 1) {
        // Return the length of the UDP data bytes
        return (eth_rx_queue_buffer->length - sizeof(udp_header_t) - sizeof(ipv4_header_t) - sizeof(ethernet_header_t));
    }

    return 0;
}



/*****************************************************************************/
/**
 * Initialize the UDP Header
 *
 * @param   header           - Pointer to the UDP header
 * @param   src_port         - Source port for UDP packet
 *
 * @return  None
 *
 *****************************************************************************/
void udp_init_header(udp_header_t * header, u16 src_port) {

    // Update the following fields that are static for the socket:
    //   - Source port
    //
    header->src_port  = Xil_Htons(src_port);
}



/*****************************************************************************/
/**
 * Update the UDP Header
 *
 * @param   header           - Pointer to the UDP header
 * @param   dest_port        - Destination port for UDP packet (big endian)
 * @param   udp_length       - Length of the UDP packet (includes UDP header)
 *
 * @return  None
 *
 *****************************************************************************/
void udp_update_header(udp_header_t* header, u16 dest_port, u16 udp_length) {

    // Update the following fields:
    //   - Destination port
    //   - Length
    //
    // NOTE:  We do not need to update the following fields because they are static for the socket:
    //   - Source port
    //
    header->dest_port = dest_port;
    header->length    = Xil_Htons(udp_length);
   
    // Currently, the Mango wlan_exp IP/UDP Library does not use the UDP checksum capabilities.  This is primarily
    // due to the amount of time required to compute the checksum.  Also, given that communication 
    // between hosts and nodes is, in general, fairly localized, there is not as much of a need for 
    // the data integrity check that the UDP checksum provides.
    //
    header->checksum  = UDP_NO_CHECKSUM;
}

/**********************************************************************************************************************/
/**
 * @brief ICMP Functions
 *
 **********************************************************************************************************************/


/*****************************************************************************/
/**
 * Process the ICMP packet
 *
 * @param   eth_rx_queue_buffer - Rx Ethernet packet containing ICMP
 * @param   eth_tx_queue_buffer - Buffer with space to construct Tx Ethernet packet
 *
 * @return  int         - Always returns 0 since we don't want higher level transports
 *                        to process this packet
 *
 * @note    The library only support Echo Request ICMP packets
 * @note    This function assumes that both Ethernet device and buffer are valid.
 *
 *****************************************************************************/
int icmp_process_packet(eth_rx_queue_buffer_t* eth_rx_queue_buffer, eth_tx_queue_buffer_t* eth_tx_queue_buffer) {

    icmp_header_t* icmp;
    ipv4_header_t* ip_hdr;

    ip_hdr = (ipv4_header_t*)(eth_rx_queue_buffer->pkt + sizeof(ethernet_header_t));

    icmp = (icmp_header_t*)((u8*)ip_hdr + sizeof(ipv4_header_t));

    // Check if this is an ICMP Echo Request to the node
    if ((icmp->type == ICMP_ECHO_REQUEST_TYPE) && (icmp->code == ICMP_ECHO_CODE)  ) {
        // Send an ICMP Echo Reply
        icmp_echo_reply(eth_rx_queue_buffer, eth_tx_queue_buffer);
    }

    return 0;    // Upper layer stacks should not process this packet so return zero bytes
}



/*****************************************************************************/
/**
 * Send an ICMP Echo Reply
 *
 * @param   eth_rx_queue_buffer - Rx Ethernet packing containing ICMP Echo Request
 * @param   eth_tx_queue_buffer - Buffer containing space to construct Tx Ethernet packet
 *
 * @return  None
 *
 * @note    This function assumes that both socket and buffer are valid.
 *
 *****************************************************************************/
void icmp_echo_reply(eth_rx_queue_buffer_t* eth_rx_queue_buffer, eth_tx_queue_buffer_t* eth_tx_queue_buffer) {

    icmp_echo_header_t* recv_icmp_hdr;
    icmp_echo_header_t* send_icmp_hdr;
    ipv4_header_t* recv_ip_hdr;
    ethernet_header_t* recv_eth_hdr;
    u8* recv_icmp_data;
    u32 recv_icmp_data_len;
    ipv4_header_t* ip_hdr;
    
    u8 my_ip_addr[IP_ADDR_LEN];
    eth_get_ip_addr(my_ip_addr);

    u8 eth_hw_addr[ETH_ADDR_LEN];

    eth_get_hw_addr(eth_hw_addr);

    recv_ip_hdr = (ipv4_header_t*)((eth_rx_queue_buffer->pkt) + sizeof(ethernet_header_t));
    recv_eth_hdr = (ethernet_header_t*)(eth_rx_queue_buffer->pkt);

    recv_icmp_hdr = (icmp_echo_header_t*)((void*)recv_ip_hdr + sizeof(ipv4_header_t));
    recv_icmp_data = (u8*)recv_icmp_hdr + sizeof(icmp_echo_header_t);
    recv_icmp_data_len = eth_rx_queue_buffer->length - sizeof(ethernet_header_t) - sizeof(ipv4_header_t) - sizeof(icmp_echo_header_t);

    if (eth_tx_queue_buffer != NULL) {
    
        // Initialize the Ethernet header
        eth_hdr_init((ethernet_header_t*)(eth_tx_queue_buffer->seg0),
        			 recv_eth_hdr->src_mac_addr,
        			 eth_hw_addr);
        
        ip_hdr = (ipv4_header_t*)((eth_tx_queue_buffer->seg0) + sizeof(ethernet_header_t));

        // Initialize the IP header
        ip_hdr_init(ip_hdr,
        			my_ip_addr,
        			recv_ip_hdr->src_ip_addr,
        			IP_PROTOCOL_ICMP);
        
        // Get the pointer to the ICMP reply header
        send_icmp_hdr = (icmp_echo_header_t*)((u8*)ip_hdr + sizeof(ipv4_header_t));

        // Populate the ICMP reply
        send_icmp_hdr->type       = ICMP_ECHO_REPLY_TYPE;
        send_icmp_hdr->code       = ICMP_ECHO_CODE;
        send_icmp_hdr->checksum   = 0;
        send_icmp_hdr->identifier = recv_icmp_hdr->identifier;
        send_icmp_hdr->seq_num    = recv_icmp_hdr->seq_num;
        
        // Copy all the data from the request packet
        memcpy( (u8*)send_icmp_hdr + sizeof(icmp_echo_header_t), recv_icmp_data, recv_icmp_data_len);

        // Calculate the ICMP checksum
        send_icmp_hdr->checksum = Xil_Htons(ipv4_compute_checksum((u8*)send_icmp_hdr,
        														   eth_rx_queue_buffer->length - sizeof(ethernet_header_t) - sizeof(ipv4_header_t)));
        
        
        // By setting the length of segment 0, we inform the calling context we have filled
        //  in an Ethernet frame that needs transmission
        eth_tx_queue_buffer->seg0_len = eth_rx_queue_buffer->length;

        // Set the length in the IP header and update checksum
        _set_ip_length(eth_tx_queue_buffer->seg0, eth_tx_queue_buffer->seg0_len);

    } else {
    	xil_printf("Error: Provided with NULL eth_tx_queue_buffer_t*\n");
    }
}


/*****************************************************************************/
/**
 * Set MAC Address
 *
 * @param   hw_addr          - u8 pointer of MAC address to set
 *
 * @return  None
 *
 ******************************************************************************/
void eth_set_hw_addr(u8* hw_addr) {

	memcpy(gl_hw_addr, hw_addr, ETH_ADDR_LEN);

    return;
}


/*****************************************************************************/
/**
 * Get MAC Address
 *
 * @param   hw_addr          - u8 pointer to where MAC address will be copied to
 *
 * @return  None
 *
 ******************************************************************************/
void eth_get_hw_addr(u8* hw_addr) {

	memcpy(hw_addr, gl_hw_addr, ETH_ADDR_LEN);

    return;
}



/*****************************************************************************/
/**
 * Set IP Address
 *
 * @param   ip_addr          - u8 pointer of IP address to set
 *
 * @return  None
 *
 ******************************************************************************/
void eth_set_ip_addr(u8* ip_addr) {

    memcpy(gl_ip_addr, ip_addr, IP_ADDR_LEN);

    return;
}

/*****************************************************************************/
/**
 * Het IP Address
 *
 * @param   ip_addr          - u8 pointer to where IP address will be copied to
 *
 * @return  None
 *
 ******************************************************************************/
void eth_get_ip_addr(u8* ip_addr) {

	memcpy(ip_addr, gl_ip_addr, IP_ADDR_LEN);

    return;
}

#endif // #if WLAN_SW_CONFIG_ENABLE_WLAN_EXP