/*
 **************************************************************************
 *
 * Boot-ROM-Code to load an operating system across a TCP/IP network.
 *
 * Module:  tftp.c
 * Purpose: Get a file with TFTP protocol
 * Entries: tftp_open, tftp_get, tftp_close
 *
 **************************************************************************
 *
 * Copyright (C) 1995-2003 Gero Kuhlmann <gero@gkminix.han.de>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * $Id: tftp.c,v 1.7 2003/03/09 00:24:27 gkminix Exp $
 */


#include <general.h>
#include <kernel/net.h>
#include <kernel/arpa.h>
#include <kernel/romlib.h>
#include <pxe/common.h>
#include <pxe/tftp.h>
#include "tftp.h"



/*
 **************************************************************************
 * 
 * Definitions for the open flag:
 */
#define FLAG_UDP_CLOSED		0	/* UDP connection closed	*/
#define FLAG_TFTP_CLOSED	1	/* TFTP connection closed	*/
#define FLAG_TFTP_OPEN		2	/* TFTP connection opened	*/



/*
 **************************************************************************
 * 
 * Global variables:
 */

/* Buffers */
static union {				/* Input buffer			*/
	unsigned char  buf[INPBUFSIZE];
	struct tftphdr hdr;
} inpbuf;

static union {				/* ACK output buffer		*/
	unsigned char  buf[ACKBUFSIZE];
	struct tftphdr hdr;
} ackbuf;


/* Global variables for directly accessing this module */
unsigned char *tftpbuf = &(inpbuf.hdr.th_data);
unsigned int   tftpbufsize;
unsigned int   tftpbuflen;
unsigned long  filesize;
unsigned long  srvrtout;
         int   tftpoflag;


/* Variables local to this module */
static int tftp_sport = 0;		/* TFTP server port number	*/
static int tftp_tout = TFTP_TIMEOUT;	/* TFTP read timeout		*/
static int currblock;			/* Current data block		*/
static int openflag;			/* Open status			*/
static int errval;			/* error status value		*/

/* Definition of option strings */
static unsigned char blksizestr[] = BLKSIZE_STR;
static unsigned char tsizestr[] = TSIZE_STR;
static unsigned char toutstr[] = TOUT_STR;



/*
 **************************************************************************
 * 
 * Convert a string into a long integer value
 */
static unsigned char *getlong(startp, retval)
unsigned char *startp;
unsigned long *retval;
{
  register unsigned char *cp = startp;
  unsigned long val = 0L;

  while (*cp) {
	if (*cp >= '0' && *cp <= '9') {
		val = val * 10;
		val += (*cp - '0');
	}
	cp++;
  }
  *retval = val;
  return(cp);
}



/*
 **************************************************************************
 * 
 * Convert an unsigned integer value into a string
 */
static unsigned char *putnum(startp, val)
unsigned char *startp;
unsigned int val;
{
  register unsigned char *cp;
  unsigned char buf[8];

  /* Convert number into string */
  cp = &buf[7];
  while (val > 0) {
	*cp-- = (val % 10) + '0';
	val /= 10;
  }

  /* Copy string into destination buffer */
  if (cp == &buf[7])
	*startp++ = '0';
  else while (cp < &buf[7])
	*startp++ = *++cp;
  return(startp);
}



/*
 **************************************************************************
 * 
 * Handle an options string
 */
static int handle_options(startp, endp)
unsigned char *startp;
unsigned char *endp;
{
  register unsigned char *cp = startp;
  unsigned long longval;

  while (cp < endp) {
	if (!memcmp(cp, blksizestr, sizeof(BLKSIZE_STR))) {
		/* Handle blocksize option */
		cp += sizeof(BLKSIZE_STR);
		cp = getlong(cp, &longval);
		if (longval < DEFSEGSIZE || longval > MAXSEGSIZE)
			return(FALSE);
		tftpbufsize = (unsigned int)longval;
	} else if (!memcmp(cp, tsizestr, sizeof(TSIZE_STR))) {
		/* Handle filesize option */
		cp += sizeof(TSIZE_STR);
		cp = getlong(cp, &filesize);
	} else if (!memcmp(cp, toutstr, sizeof(TOUT_STR))) {
		/* Handle timeout option */
		cp += sizeof(TOUT_STR);
		cp = getlong(cp, &srvrtout);
	} else
		return(FALSE);
	cp++;
  }
  return(TRUE);
}



/*
 **************************************************************************
 * 
 * Send a tftp request packet
 */
static int send_req(fname, fnamlen)
unsigned char *fname;
int            fnamlen;
{
#define OPTSIZE	(sizeof(OCTET_STR) + sizeof(BLKSIZE_STR) + 4 + \
						sizeof(TSIZE_STR) + 2)

  register unsigned char *cp;
  char *octet;

  /* Check that the filename is not too long */
  if (fnamlen > (DEFSEGSIZE - sizeof(inpbuf.hdr.th_op) - OPTSIZE - 2))
	return(PXENV_STATUS_FAILURE);

  /*
   * Setup the output buffer with the request code, the file name,
   * all options and the name of the data format.
   */
  memset(&inpbuf, 0, sizeof(inpbuf));
  inpbuf.hdr.th_op = htons(TFTP_RRQ);
  cp = (unsigned char *)&(inpbuf.hdr.th_block);
  for (; *fname && fnamlen > 0; fnamlen--)
	*cp++ = *fname++;
  cp++;
  memcpy(cp, OCTET_STR, sizeof(OCTET_STR));
  cp += sizeof(OCTET_STR);
  if (tftpbufsize != DEFSEGSIZE) {
	memcpy(cp, blksizestr, sizeof(BLKSIZE_STR));
	cp = putnum(cp + sizeof(BLKSIZE_STR), tftpbufsize) + 1;
  }
  memcpy(cp, tsizestr, sizeof(TSIZE_STR));
  cp += sizeof(TSIZE_STR);
  *cp++ = '0';
  cp++;
  memcpy(cp, toutstr, sizeof(TOUT_STR));
  cp = putnum(cp + sizeof(TOUT_STR), tftp_tout) + 1;

  /* Finally send the request */
  return(udp_write(inpbuf.buf, getds(), (int)(cp - inpbuf.buf)));
}



/*
 **************************************************************************
 * 
 * Send a tftp acknowledge packet
 */
static int send_reply(code, block)
int block;
{
  int len;

  memset(&ackbuf, 0, sizeof(ackbuf));
  ackbuf.hdr.th_op = code;
  ackbuf.hdr.th_block = htons(block);
  len = (int)(ackbuf.hdr.th_data - (unsigned char *)&(ackbuf.hdr));
  if (code == TFTP_ERROR)
	/* Error text is the null string */
	len++;
  return(udp_write(ackbuf.buf, getds(), len));
}

#define send_ack(block)	send_reply(htons(TFTP_ACK), (block))
#define send_nak(err)	send_reply(htons(TFTP_ERROR), (err))



/*
 **************************************************************************
 * 
 * Receive a TFTP data packet. It returns FALSE if the connection should be
 * closed physically.
 */
static int rcv_packet(block)
int block;
{
#define th_data_ofs ((int)(inpbuf.hdr.th_data - (unsigned char *)&(inpbuf.hdr)))

  int status;
  int isopen = TRUE;

  /* Read packet with timeout */
  errval = PXENV_STATUS_SUCCESS;
  tftpbuflen = sizeof(inpbuf.buf);
  status = udp_read(inpbuf.buf, getds(), &tftpbuflen, tftp_tout, CHR_ESC);
  if (status == PXENV_STATUS_ARP_TIMEOUT) {
	errval = PXENV_STATUS_TFTP_READ_TIMEOUT;
	goto rcv_ack;
  } else if (status != PXENV_STATUS_SUCCESS) {
	errval = PXENV_STATUS_TFTP_READ_CANCELED;
	isopen = FALSE;
	tftpbuflen = 0;
	send_nak(ENOSPACE);
	goto rcv_ret;
  }

  /* Handle option acknowledgement packet */
  if (inpbuf.hdr.th_op == htons(TFTP_OACK)) {
	if (block != 1 ||
	    !handle_options((unsigned char *)&(inpbuf.hdr.th_block),
				(unsigned char *)&(inpbuf.hdr) + tftpbuflen)) {
		errval = PXENV_STATUS_TFTP_ERROR_OP;
		isopen = FALSE;
		tftpbuflen = 0;
		send_nak(EBADOPT);
		goto rcv_ret;
	}
	block = 0;
	tftpbuflen = 0xFFFF;
	goto rcv_ack;
  }

  /* Check that the packet has a correct length */
  if (tftpbuflen < th_data_ofs) {
	errval = PXENV_STATUS_FAILURE;
	goto rcv_ret;
  }
  tftpbuflen -= th_data_ofs;
  if (tftpbuflen > tftpbufsize) {
	errval = PXENV_STATUS_FAILURE;
	goto rcv_ret;
  }

  /* Check if we got an error packet. This will terminate the connection */
  if (inpbuf.hdr.th_op == htons(TFTP_ERROR)) {
	if (inpbuf.hdr.th_error == htons(ENOTFOUND))
		errval = PXENV_STATUS_TFTP_INVALID_FILE;
	else if (inpbuf.hdr.th_error == htons(EACCESS))
		errval = PXENV_STATUS_TFTP_ACCESS;
	else if (inpbuf.hdr.th_error == htons(EBADOP))
		errval = PXENV_STATUS_TFTP_ERROR_OP;
	else
		errval = PXENV_STATUS_FAILURE;
	isopen = FALSE;
	goto rcv_ret;
  }

  /* Check if we got a valid data packet at all. We only need to ack data */
  if (inpbuf.hdr.th_op != htons(TFTP_DATA)) {
	errval = PXENV_STATUS_TFTP_UNKNOWN_OP;
	goto rcv_ack;
  }

  /* Check if the block number of the data packet is correct */
  if (inpbuf.hdr.th_block != htons(block))
	errval = PXENV_STATUS_TFTP_TOO_MANY;

  /*
   * Ack the current packet. This will also resynchronize the sender if
   * the block was wrong.
   */
rcv_ack:
  if (errval != PXENV_STATUS_SUCCESS) {
	block--;
	send_ack(block);
  } else if ((errval = send_ack(block)) != PXENV_STATUS_SUCCESS)
	isopen = FALSE;

rcv_ret:
  return(isopen);

#undef th_data_ofs
}



/*
 **************************************************************************
 * 
 * Open a TFTP connection
 */
int tftp_open(server, gateway, fname, fnamlen)
t_ipaddr        server;
t_ipaddr        gateway;
unsigned char  *fname;
int             fnamlen;
{
  register int retry, tid;

  /* Set default transfer block size */
  if (tftpbufsize == 0)
	tftpbufsize = MAXSEGSIZE;
  else if (tftpbufsize < DEFSEGSIZE)
	tftpbufsize = DEFSEGSIZE;

  /* Set default port */
  if (tftp_sport == 0)
	tftp_sport = htons(TFTP_S_PORT);

  /* Try to open a TFTP connection to the server */
  currblock = 0;
  filesize = -1L;
  srvrtout = -1L;
  errval = PXENV_STATUS_TFTP_CANNOT_OPEN;
  for (retry = 0; openflag == FLAG_UDP_CLOSED && retry < TFTP_RETRY; retry++) {
	/*
	 * First open a new socket. The local port will be the same as the
	 * transaction ID, which should be a random number and different
	 * between each retry.
	 */
	tid = htons(TFTP_C_PORT + random(0x003f) + ((retry * 64) & 0x1c0));
	if (udp_open(server, gateway, tid, tftp_sport) != PXENV_STATUS_SUCCESS)
		break;
	openflag = FLAG_TFTP_CLOSED;

	/*
	 * Send the file request block, and then wait for the first data
	 * block. If there is no response to the query, retry it with
	 * another transaction ID, so that all old packets get discarded
	 * automatically.
	 */
	if ((errval = send_req(fname, fnamlen)) == PXENV_STATUS_SUCCESS) {
		tftpbufsize = DEFSEGSIZE;
		if (rcv_packet(1) && errval == PXENV_STATUS_SUCCESS) {
			openflag = FLAG_TFTP_OPEN;
			return(PXENV_STATUS_SUCCESS);
		}
	}

	/* If an error occurred, retries are useless */
	openflag = FLAG_UDP_CLOSED;
	if (errval != PXENV_STATUS_TFTP_READ_TIMEOUT) {
		udp_close();
		break;
	}
	tftpbufsize = DEFSEGSIZE;
	errval = PXENV_STATUS_TFTP_CANNOT_OPEN;
	if (udp_close() != PXENV_STATUS_SUCCESS)
		break;
  }
  tftpbufsize = 0;
  return(errval);
}



/*
 **************************************************************************
 *
 * Read the next data packet from a TFTP connection
 */
int tftp_get()
{
  register int retry;

  /* Don't do anything if no TFTP connection is active. */
  if (openflag != FLAG_TFTP_OPEN)
	return(PXENV_STATUS_TFTP_CANNOT_READ);

  /*
   * If the block number is 0 then we are still dealing with the first
   * data block after opening a connection. If the data size is smaller
   * than the current block size just close the connection again.
   * However, if the last received packet was an option ack, we still
   * have to read the data packet. In that case, tftpbuflen is negative.
   */
  if (currblock == 0 && tftpbuflen != 0xFFFF) {
	currblock++;
	if (tftpbuflen < tftpbufsize)
		openflag = FLAG_TFTP_CLOSED;
	return(PXENV_STATUS_SUCCESS);
  }

  /*
   * Wait for the next data packet. If no data packet is coming in,
   * resend the ACK for the last packet to restart the sender. Maybe
   * he didn't get our first ACK.
   */
  errval = PXENV_STATUS_TFTP_CANNOT_READ;
  for (retry = 0; retry < TFTP_RETRY; retry++) {
	if (!rcv_packet(currblock + 1))
		break;
	if (errval == PXENV_STATUS_SUCCESS) {
		currblock++;
		if (tftpbuflen < tftpbufsize)
			openflag = FLAG_TFTP_CLOSED;
		return(PXENV_STATUS_SUCCESS);
	}
  }
  openflag = FLAG_TFTP_CLOSED;
  return(errval);
}



/*
 **************************************************************************
 *
 * Close a TFTP connection
 */
int tftp_close()
{
  /* We always succeed in closing an already closed connection */
  if (openflag == FLAG_UDP_CLOSED)
	return(PXENV_STATUS_SUCCESS);

  /*
   * If the connection is still open, read the next packet because the
   * file size can be a multiple of the TFTP packet size, in which case
   * the reading process has already stopped reading while the TFTP server
   * is still waiting to close the connection with a null-sized packet. 
   * If the next packet does not indicate a termination of the TFTP connection,
   * send an error packet.
   */
  errval = PXENV_STATUS_SUCCESS;
  if (openflag == FLAG_TFTP_OPEN &&
      !(currblock == 0 && tftpbuflen < tftpbufsize)) {
	if (!rcv_packet(currblock + 1) || errval != PXENV_STATUS_SUCCESS ||
	    (tftpbuflen >= tftpbufsize && !send_nak(ENOSPACE)))
		errval = PXENV_STATUS_TFTP_CANNOT_READ;
  }

  /* Close UDP connection */
  if (udp_close() != PXENV_STATUS_SUCCESS && errval == PXENV_STATUS_SUCCESS)
	errval = PXENV_STATUS_FAILURE;
  openflag = FLAG_UDP_CLOSED;
  tftpbufsize = 0;
  return(errval);
}



/*
 **************************************************************************
 *
 * Read a file into memory
 */
int tftp_read_file(params)
t_tftp_read_file *params;
{
  register int ret;
  unsigned long loadaddr, topmem;

  /* Set loading address */
  loadaddr = params->buffer;
  topmem = loadaddr + params->bufsize;

  /* Set timeout */
  if (params->opentimeout > 0)
	tftp_tout = params->opentimeout;

  /* Open TFTP connection and read all blocks successively */
  if ((ret = tftp_open(params->server, params->gateway, params->filename,
			sizeof(params->filename))) == PXENV_STATUS_SUCCESS) {
	while (TRUE) {
		if ((ret = tftp_get()) != PXENV_STATUS_SUCCESS)
			break;
		if ((loadaddr + tftpbufsize ) > topmem) {
			ret = PXENV_STATUS_TFTP_TOO_MANY;
			break;
		}
		if (!lmove(loadaddr, tftpbuf, tftpbuflen)) {
			ret = PXENV_STATUS_MCOPY_PROBLEM;
			break;
		}
		loadaddr += tftpbuflen;
		if (tftpbuflen < tftpbufsize) {
			ret = PXENV_STATUS_SUCCESS;
			break;
		}
	}
	params->bufsize = loadaddr - params->buffer;
	tftp_close();
  }
  tftp_tout = TFTP_TIMEOUT;
  return(ret);
}



/*
 **************************************************************************
 *
 * TFTP open PXE function
 */
static int pxe_tftp_open(params)
t_tftp_open *params;
{
  int status;

  /* Check that we are allowed to open a TFTP connection */
  if (udpoflag)
	return(PXENV_STATUS_UDP_OPEN);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Open TFTP connection */
  tftp_sport = params->port;
  tftpbufsize = params->pktsize;
  if (tftpbufsize < DEFSEGSIZE)
	tftpbufsize = DEFSEGSIZE;
  else if (tftpbufsize > MAXSEGSIZE)
	tftpbufsize = MAXSEGSIZE;
  if ((status = tftp_open(params->server, params->gateway,
       params->filename, sizeof(params->filename))) == PXENV_STATUS_SUCCESS)
	tftpoflag = TRUE;
  return(status);

}



/*
 **************************************************************************
 *
 * TFTP close PXE function
 */
static int pxe_tftp_close(params)
t_tftp_close *params;
{
  int status;

  /* Check that we are allowed to close the TFTP connection */
  if (!tftpoflag)
	return(PXENV_STATUS_TFTP_CLOSED);

  /* Close TFTP connection */
  tftp_sport = 0;
  tftpbufsize = 0;
  if ((status = tftp_close()) == PXENV_STATUS_SUCCESS)
	tftpoflag = FALSE;
  return(status);
}



/*
 **************************************************************************
 *
 * TFTP read data PXE function
 */
static int pxe_tftp_read_data(params)
t_tftp_read_data *params;
{
  int ret = PXENV_STATUS_TFTP_TOO_MANY;

  /* Check that we are allowed to read data from TFTP connection */
  if (!tftpoflag)
	return(PXENV_STATUS_TFTP_CLOSED);

  /* Read data from open TFTP connection */
  if (params->bufsize >= tftpbufsize &&
      (ret = tftp_get()) == PXENV_STATUS_SUCCESS) {
	fmemcpy(params->bufofs, params->bufseg, tftpbuf, tftpbuflen);
	params->bufsize = tftpbuflen;
	params->packetno = currblock;
  }
  return(ret);
}



/*
 **************************************************************************
 *
 * TFTP read file PXE function
 */
int pxe_tftp_read_file(params)
t_tftp_read_file *params;
{
  /* Check if we are allowed to read a file */
  if (udpoflag)
	return(PXENV_STATUS_UDP_OPEN);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Simply call our local function to read the file */
  return(tftp_read_file(params));
}



/*
 **************************************************************************
 *
 * Get file size.
 */
static int pxe_get_file_size(params)
t_tftp_get_file_size *params;
{
  int ret = PXENV_STATUS_FAILURE;

  /* No TFTP connection is allowed to be open */
  if (udpoflag)
	return(PXENV_STATUS_UDP_OPEN);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Open and close TFTP connection. This will get us the file size */
  tftp_sport = 0;
  tftpbufsize = DEFSEGSIZE;
  if (tftp_open(params->server, params->gateway,
	params->filename, sizeof(params->filename)) == PXENV_STATUS_SUCCESS &&
      tftp_close() == PXENV_STATUS_SUCCESS && filesize != -1L) {
	params->filesize = filesize;
	ret = PXENV_STATUS_SUCCESS;
  }
  tftpbufsize = 0;
  return(ret);
}



/*
 **************************************************************************
 *
 * Function table for TFTP PXE functions
 */
struct functab tftp_func_tab = {
	TFTP_MINNUM, TFTP_MAXNUM,
	{{ pxe_tftp_open,	sizeof(t_tftp_open)},
	 { pxe_tftp_close,	sizeof(t_tftp_close)},
	 { pxe_tftp_read_data,	sizeof(t_tftp_read_data)},
	 { pxe_tftp_read_file,	sizeof(t_tftp_read_file)},
	 { pxe_get_file_size,	sizeof(t_tftp_get_file_size)}
	}
};



syntax highlighted by Code2HTML, v. 0.9.1