/*
 * File:		tftp.c
 * Purpose:		Trivial File Transfer Protocol driver for reading a file
 * 				from a remote host.
 *
 * Notes:		See RFC 1350
 *
 * Modifications:
 *
 */

#include "src/init/m523xevb.h"
#include "src/init/stdlib.h"
#include "src/ethernet/nif.h"
#include "src/ethernet/tftp/udp.h"
#include "src/ethernet/tftp/ip.h"
#include "src/ethernet/tftp/timer.h"
#include "src/ethernet/tftp/tftp.h"

/********************************************************************/

/* The one and only TFTP connection */
TFTP_Connection tcxn;

/* Progress Indicators */
char hash[] = {'-','\\','|','/'};
int ihash = 0;

/********************************************************************/
static int
tftp_rwrq (void)
{
	NBUF *pNbuf;
	RWRQ *rwrq;
	int i;

	pNbuf = tcxn.nif->tx_alloc();

	/* Make sure I grabbed a valid Tx buffer */
	if (pNbuf == NULL)
	{
		#if defined(DEBUG_PRINT)
			printf("TFTP: tftp_rwrq() couldn't allocate Tx buffer\n");
		#endif
		return 0;
	}

	rwrq = (RWRQ *)&pNbuf->data[TFTP_HDR_OFFSET];

	/* Indicate a R/WRQ */
	rwrq->opcode = tcxn.dir;

	/* Copy in filename */
	strcpy(&rwrq->filename_mode[0], tcxn.file);
	i = strlen(tcxn.file) + 1;

	/* Indicate transfer type */
	strcpy (&rwrq->filename_mode[i], OCTET);

	pNbuf->length = (uint16)(i + strlen(OCTET) + 1 + 2);

	return udp_send(tcxn.nif,
					tcxn.server_ip,
					tcxn.my_port,
					tcxn.server_port,
					pNbuf);
}

/********************************************************************/
static int
tftp_ack (uint16 blocknum)
{
	ACK *ack;
	NBUF *pNbuf;

	pNbuf = tcxn.nif->tx_alloc();

	/* Make sure I grabbed a valid Tx buffer */
	if (pNbuf == NULL)
	{
		#if defined(DEBUG_PRINT)
			printf("TFTP: tftp_ack() couldn't allocate Tx buffer\n");
		#endif
		return 0;
	}

	ack = (ACK *)&pNbuf->data[TFTP_HDR_OFFSET];
	ack->opcode = TFTP_ACK;
	ack->blocknum = blocknum;

	pNbuf->length = 4;

	return udp_send(tcxn.nif,
					tcxn.server_ip,
					tcxn.my_port,
					tcxn.server_port,
					pNbuf);
}

/********************************************************************/
static int
tftp_error (uint16 error_code, uint16 server_port)
{
	ERROR *err;
	NBUF *pNbuf;

	pNbuf = tcxn.nif->tx_alloc();

	/* Make sure I grabbed a valid Tx buffer */
	if (pNbuf == NULL)
	{
		#if defined(DEBUG_PRINT)
			printf("TFTP: tftp_error() couldn't allocate Tx buffer\n");
		#endif
		return 0;
	}

	err = (ERROR *)&pNbuf->data[TFTP_HDR_OFFSET];
	err->opcode = TFTP_ERROR;
	err->code = error_code;
	err->msg[0] = '\0';

	pNbuf->length = 5;

	return udp_send(tcxn.nif,
					tcxn.server_ip,
					tcxn.my_port,
					server_port,
					pNbuf);
}

/********************************************************************/
void
tftp_handler (NIF *nif, NBUF *pNbuf, int hdr_offset)
{
	union TFTPpacket *tftp_pkt;
	udp_frame_hdr *udpframe;
	static int cnt;
	(void) nif;

	/* Accessing shared memory here, disable interrupts */
	asm_set_ipl(6);

	tftp_pkt = (union TFTPpacket *)&pNbuf->data[hdr_offset];
	udpframe = (udp_frame_hdr *)&pNbuf->data[hdr_offset - UDP_HDR_SIZE];
	
	switch (tftp_pkt->generic.opcode)
	{
		case TFTP_DATA:
			if (tftp_pkt->data.blocknum == tcxn.exp_blocknum)
			{	
				/* This is the block expected */
				if (tftp_pkt->data.blocknum == 1)
				{	/* This is the first data block received */

					/* Save the server's transfer ID */
					tcxn.server_port = UDP_SOURCE(udpframe);

					/* Mark the connection as open */
					tcxn.open = TRUE;

					/* Point to the first character in the buffer */
					tcxn.next_char = tftp_pkt->data.data;

					/* Start progress indicator */
					printf("%c",hash[0]);
					cnt = 0;
				}
				else
				{	/* Check the server's transfer ID */
					if (tcxn.server_port != UDP_SOURCE(udpframe))
					{
						#if defined(DEBUG_PRINT)
							printf("TFTP: Invalid server port: %d\n", \
									UDP_SOURCE(udpframe));
						#endif

						/*Send ERROR packet to source */
						tftp_error(TFTP_ERR_TID, UDP_SOURCE(udpframe));
						break;
					}
				}

				/* Add the payload to my data buffer */
				if (tcxn.pkt_data.bytes == 0)
				{
					pNbuf->length -= 4;

					/* Copy data to buffer */
					memcpy (tcxn.pkt_data.data,
							&pNbuf->data[hdr_offset+4],
							pNbuf->length);
					tcxn.pkt_data.bytes = pNbuf->length;

					/* Update number of the next block expected */
					tcxn.exp_blocknum++;

					/* Increment number of bytes received counter */
					tcxn.bytes_recv += pNbuf->length;

					/* Is this the last packet? */
					if (pNbuf->length < TFTP_PKTSIZE)
					{
						/* Mark the connection as closed */
						tcxn.open = FALSE;

						/* ACK the last packet now */
						tftp_ack(tftp_pkt->data.blocknum);
					}

					/* Update progress indicator */
					if (++cnt == 10)
					{
						ihash = (ihash + 1) % 4;
						printf("\b%c",hash[ihash]);
						cnt = 0;
					}
				}
			}
			else
			{
				#if defined(DEBUG_PRINT)
					/* This is NOT the block expected */
					printf("Exp: %d, ", tcxn.exp_blocknum);
					printf("Rcv: %d\n", tftp_pkt->data.blocknum);
				#endif
			}
			break;
		case TFTP_ERROR:
			printf("\nTFTP Error #%d: ",tftp_pkt->error.code);
			printf("%s\n",tftp_pkt->error.msg);
			tcxn.error = TRUE;
			break;
		case TFTP_RRQ:
			printf("\nTFTP Read Request\n");
			break;
		case TFTP_WRQ:
			printf("\nTFTP Write Request\n");
			break;
		case TFTP_ACK:
			if (tftp_pkt->ack.blocknum == tcxn.exp_blocknum)
			{
				if (tftp_pkt->data.blocknum == 0)
				{	/* This is the first ACK received */

					/* Save the server's transfer ID */
					tcxn.server_port = UDP_SOURCE(udpframe);

					/* Mark the connection as open */
					tcxn.open = TRUE;
				}
				else
				{	/* Check the server's transfer ID */
					if (tcxn.server_port != UDP_SOURCE(udpframe))
					{
						#if defined(DEBUG_PRINT)
							printf("TFTP: Invalid server port: %d\n", \
									UDP_SOURCE(udpframe));
						#endif

						/*Send ERROR packet to source */
						tftp_error(TFTP_ERR_TID, UDP_SOURCE(udpframe));
						break;
					}
				}

				tcxn.exp_blocknum++;
			}
			else
			{
				#if defined(DEBUG_PRINT)
					/* This is NOT the block number expected */
					printf("ACK Exp: %d, ", tcxn.exp_blocknum);
					printf("ACK Rcv: %d\n", tftp_pkt->ack.blocknum);
				#endif
			}
			break;
		default:
			break;
	}

	/* Re-enable interrupts now */
	asm_set_ipl(0);
}

/********************************************************************/
void
tftp_end (int success)
{
	if (success && !tcxn.error)
	{
		printf("\bTFTP download successful\n");
		printf("Read %d bytes (%d blocks)\n",	\
			tcxn.bytes_recv, tcxn.exp_blocknum - 1);
	}
	else
	{
		printf("\bErrors in TFTP download.\n");
		printf("Read %d bytes (%d blocks)\n",	\
			tcxn.bytes_recv, tcxn.exp_blocknum - 1);
	}
    tcxn.nif->stop(tcxn.nif);
	udp_free_port(tcxn.my_port);
}

/********************************************************************/
int
tftp_write (NIF *nif, char *fn, IP_ADDR_P server, uint32 begin, uint32 end)
{
	DATA *data;
	NBUF *pNbuf;

	uint32 i, retries, bytes_to_send;
	uint16 blocknum, this_size;
	uint8 success, *current;

	if (fn == 0 || server == 0 || end < begin)
		return 0;

	/* Setup initial connection status */
	tcxn.nif = nif;
	tcxn.file = fn;
	tcxn.server_ip[0] = server[0];
	tcxn.server_ip[1] = server[1];
	tcxn.server_ip[2] = server[2];
	tcxn.server_ip[3] = server[3];
	tcxn.server_port = UDP_PORT_TFTP;
	tcxn.exp_blocknum = 0;
	tcxn.dir = TFTP_WRQ;
	tcxn.open = FALSE;
	tcxn.bytes_sent = 0;
	tcxn.error = FALSE;

	/* Reset the Ethernet Controller */
	tcxn.nif->reset(tcxn.nif);
	tcxn.nif->start(tcxn.nif);

	/* Use Mac address as pseudo-random port */
	udp_prime_port((uint16)((nif->hwa[4] << 8) | nif->hwa[5]));
	tcxn.my_port = udp_obtain_free_port();
	udp_bind_port(tcxn.my_port,&tftp_handler);

	retries = 4;
	success = FALSE;

	while (--retries)
	{
		/* Make the TFTP Read/Write Request */
		if (!tftp_rwrq())
		{
			tcxn.error = TRUE;
			tftp_end(FALSE);
			return FALSE;
		}
	
		/* Enable interrupts so we can get a packet */
		asm_set_ipl(0);

		timer_set_secs(TIMER_NETWORK, TFTP_TIMEOUT);
		while (timer_get_reference(TIMER_NETWORK))
		{
			/* Has the server responded */
			if (tcxn.open)
			{
				success = TRUE;
				break;
			}
		}

		/* Disable interrupts */
		asm_set_ipl(6);

		/* If the connection is open, we are done here */
		if (success || tcxn.error)
			break;
	}
	if (!retries)
	{
		printf("TFTP could not make connection to server.\n");
	    tcxn.nif->stop(tcxn.nif);
		udp_free_port(tcxn.my_port);
		return FALSE;
	}
	else if (tcxn.error)
	{
		printf("\bErrors in TFTP upload.\n");
	    tcxn.nif->stop(tcxn.nif);
		udp_free_port(tcxn.my_port);
		return FALSE;
	}

	bytes_to_send = end - begin;
	current = (uint8 *)begin;
	blocknum = 1;
	retries = 4;
	success = FALSE;

	while (--retries)
	{
		pNbuf = tcxn.nif->tx_alloc();

		/* Make sure I grabbed a valid Tx buffer */
		if (pNbuf == NULL)
		{
			#if defined(DEBUG_PRINT)
				printf("TFTP: tftp_write() couldn't allocate Tx buffer\n");
			#endif
			return FALSE;
		}

		/* Build the packet */
		data = (DATA *)&pNbuf->data[TFTP_HDR_OFFSET];
		data->blocknum = blocknum;
		data->opcode = TFTP_DATA;
	
		this_size = (bytes_to_send > TFTP_PKTSIZE) ? \
					TFTP_PKTSIZE : (uint16)bytes_to_send;

		for (i = 0; i < this_size; i++)
		{
			data->data[i] = current[i];
		}

		/* Set the packet length */
		pNbuf->length = (uint16)(4 + this_size);

		/* Send the packet */
		udp_send(tcxn.nif, tcxn.server_ip, tcxn.my_port, tcxn.server_port, pNbuf);

		/* Enable interrupts so we can get an ACK */
		asm_set_ipl(0);

		timer_set_secs(TIMER_NETWORK, TFTP_TIMEOUT);
		while (timer_get_reference(TIMER_NETWORK))
		{
			/* Has the server responded */
			if ((tcxn.exp_blocknum - 1) == blocknum)
			{
				success = TRUE;
				break;
			}
		}

		/* Disable interrupts */
		asm_set_ipl(6);

		/* TFTP Write Compeleted successfully */
		if (success && (this_size < TFTP_PKTSIZE))
		{
			tcxn.bytes_sent += this_size;
			break;
		}
			
		if (tcxn.error)
			break;

		/* If an ACK was received, keep sending packets */
		if (success)
		{
			tcxn.bytes_sent += TFTP_PKTSIZE;
			bytes_to_send -= TFTP_PKTSIZE;
			current += TFTP_PKTSIZE;
			blocknum++;
			retries = 4;
			success = FALSE;
		}
	}
	if (tcxn.error)
	{
		printf("TFTP lost connection to server.\n");
		printf("Sent %d bytes (%d blocks)\n",	\
			tcxn.bytes_sent, tcxn.exp_blocknum - 1);
		tcxn.nif->stop(tcxn.nif);
		udp_free_port(tcxn.my_port);
		return FALSE;
	}
	else
	{
		printf("\bTFTP upload successful\n");
		printf("Sent %d bytes (%d blocks)\n",	\
			tcxn.bytes_sent, tcxn.exp_blocknum - 1);
		tcxn.nif->stop(tcxn.nif);
		udp_free_port(tcxn.my_port);
		return TRUE;
	}
}

/********************************************************************/
int
tftp_read (NIF *nif, char *fn, IP_ADDR_P server)
{
	uint32 retries;
	uint8 success;

	/* Disable interrupts */
	asm_set_ipl(6);

	if (fn == 0 || server == 0)
		return 0;

	/* Setup initial connection status */
	tcxn.nif = nif;
	tcxn.file = fn;
	tcxn.server_ip[0] = server[0];
	tcxn.server_ip[1] = server[1];
	tcxn.server_ip[2] = server[2];
	tcxn.server_ip[3] = server[3];
	tcxn.server_port = UDP_PORT_TFTP;
	tcxn.exp_blocknum = 1;
	tcxn.dir = TFTP_RRQ;
	tcxn.open = FALSE;
	tcxn.bytes_recv = 0;
	tcxn.pkt_data.bytes = 0;
	tcxn.error = FALSE;

	/* Reset the Ethernet Controller */
	tcxn.nif->reset(tcxn.nif);
	tcxn.nif->start(tcxn.nif);

	/* Use Mac address as pseudo-random port */
	udp_prime_port((uint16)((nif->hwa[4] << 8) | nif->hwa[5]));
	tcxn.my_port = udp_obtain_free_port();
	udp_bind_port(tcxn.my_port,&tftp_handler);

	retries = 4;
	success = FALSE;

	while (--retries)
	{
		/* Make the TFTP Read/Write Request */
		if (!tftp_rwrq())
		{
			tcxn.error = TRUE;
			tftp_end(FALSE);
			return FALSE;
		}
	
		/* Enable interrupts so we can get a packet */
		asm_set_ipl(0);

		timer_set_secs(TIMER_NETWORK, TFTP_TIMEOUT);
		while (timer_get_reference(TIMER_NETWORK))
		{
			/* Has the server responded */
			if (tcxn.bytes_recv)
			{
				success = TRUE;
				break;
			}
		}

		/* Disable interrupts */
		asm_set_ipl(6);

		/* If the connection is open, we are done here */
		if (success || tcxn.error)
			break;
	}
	if (!retries)
	{
		printf("TFTP could not make connection to server.\n");
		tftp_end(FALSE);
		return FALSE;
	}
	else if (tcxn.error)
	{
		tftp_end(FALSE);
		return FALSE;
	}
	else
	{
		/* Enable interrupts so we can receive more packets */
		asm_set_ipl(0);
		return TRUE;
	}
}

/********************************************************************/
int
tftp_in_char (void)
{
	int retries, success;
	int ch;
	
	/* Accessing shared memory here, disable interrupts */
	asm_set_ipl(6);

	if (tcxn.pkt_data.bytes)
	{	/* The current buffer isn't exhausted yet */
		tcxn.pkt_data.bytes--;
		ch = *tcxn.next_char++;
	}
	else
	{
		/* Make sure the connection is still open */
		if (!tcxn.open)
			return -1;
	
		retries = 6;
		success = FALSE;

		while (--retries)
		{
			/* ACK the last packet received */
			tftp_ack((uint16)(tcxn.exp_blocknum - 1));

			/* Enable interrupts so we can get a packet */
			asm_set_ipl(0);

			timer_set_secs(TIMER_NETWORK, TFTP_TIMEOUT);
			while (timer_get_reference(TIMER_NETWORK))
			{
				/* Has the server sent another DATA packet? */
				if (tcxn.pkt_data.bytes)
				{
					success = TRUE;
					break;
				}
			}

			/* Disable interrupts */
			asm_set_ipl(6);

			if (success)
				break;
		}
		if (!retries)
		{	
			/* Error: Lost connection to server */
			printf("TFTP lost connection to server.\n");
			/* Send error packet to stifle the server */
			tftp_error(TFTP_ERR_ILL, tcxn.server_port);
			tcxn.open = FALSE;
			tcxn.error = TRUE;
			return -1;
		}
		else
		{	/* Server sent more data */
			tcxn.pkt_data.bytes--;
			tcxn.next_char = tcxn.pkt_data.data;
			ch = *tcxn.next_char++;
		}
	}
	/* Enable interrupts so we can receive more packets */
	asm_set_ipl(0);

	return ch;
}

/********************************************************************/