/* vim: set filetype=c:
 *
 **************************************************************************
 *
 * Boot-ROM-Code to load an operating system across a TCP/IP network.
 *
 * Module:  udp.c86
 * Purpose: UDP and IP handling
 * Entries: udp_open, udp_close, udp_read, udp_write, udp_set_mcast
 *
 **************************************************************************
 *
 * Copyright (C) 1995-2007 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: udp.c86,v 1.8 2007/01/06 18:30:47 gkminix Exp $
 */


#include <common.h86>
#include <romconfig.h86>
#include <i386/string.h86>
#include <i386/misc.h86>
#include <i386/time.h86>
#include <i386/conio.h86>
#include <pxe/general.h86>
#include <pxe/udp.h86>
#include <pxe/undi.h86>
#include "headers/net.h86"
#include "headers/arpa.h86"
#include "headers/funcs.h86"
#include "headers/pxe.h86"
#include "ip.h86"
#include "arp.h86"
#include "igmp.h86"
#include "udp.h86"
#include "undi.h86"



/*
 **************************************************************************
 * 
 * Public variables:
 *
 * Note that these are all (except mymtu, readbuf and udpoflag) in network
 * order!!!
 */
unsigned short ippacketid;			/* IP packet id		*/
unsigned short mymtu;				/* my own MTU		*/
unsigned short mytype;				/* my type of network	*/
__u8           myhwaddr[ETH_ALEN];		/* my own hardware addr	*/
t_ipaddr       myipaddr;			/* my own IP address	*/
t_ipaddr       mynetmask;			/* my own netmask	*/
t_ipaddr       mygateway;			/* my gateway to use	*/
bool_t         udpoflag;			/* TRUE if open called	*/



/*
 **************************************************************************
 * 
 * Global variables:
 *
 * These are also all in network order!!!
 */
static unsigned short usource;		/* Source port number		*/
static unsigned short udest;		/* Destination port number	*/
static unsigned short mport;		/* Multicast port number	*/
static t_ipaddr       udaddr;		/* Destination IP address	*/
static t_ipaddr       usaddr;		/* Source IP address		*/
static t_ipaddr       gateway;		/* Gateway IP address		*/
static t_ipaddr       maddr;		/* Multicast IP address		*/



/*
 **************************************************************************
 * 
 * Send buffers: these get shared among various modules to save memory
 */
struct udphdr udpbuf asm("_protobuf") __attribute__((common));
struct iphdr ipbuf asm("_linkbuf") __attribute__((common));
size_t sizes[3] asm("_sendsizes") __attribute__((common));
__fptr bufs[3] asm("_sendbufs") __attribute__((common));



/*
 **************************************************************************
 * 
 * Open a new UDP socket.
 */
short udp_open(t_ipaddr daddr, t_ipaddr gw, short source, short dest)
{
  /* Set global variables */
  usource = source;
  udest = dest;
  mport = 0;
  usaddr = myipaddr;
  udaddr = daddr;
  gateway = gw;
  maddr = htonl(IP_ANY);

  /* Open UNDI interface */
  return(undi_open());
}



/*
 **************************************************************************
 * 
 * Close a UDP socket
 */
short udp_close(void)
{
  /* Leave any multicast group */
  if (maddr != htonl(IP_ANY)) {
	igmp_leave();
	maddr = htonl(IP_ANY);
  }

  /* Close UNDI interface */
  return(undi_close());
}



/*
 **************************************************************************
 * 
 * Receive a UDP packet. If the packet is fragmented, read the next
 * fragment and add it to the end of the buffer. This routine is non-
 * blocking.
 */
static size_t udp_recv(__fptr buf, size_t bufsize)
{
#define ipp ((struct iphdr *)readbuf.buf)
#define udpp ((struct udphdr *)(readbuf.buf + hdrsize))
#define datap ((unsigned char *)(readbuf.buf + hdrsize + sizeof(struct udphdr)))

  bool_t mcastflag = FALSE;
  bool_t fragflag;
  size_t fragoff;
  size_t readlen;
  size_t size, hdrsize;
  static t_ipaddr dest;
  static struct pshdr rcvpsbuf;
  static struct udphdr rcvudpbuf;
  static size_t totlen;
  static unsigned short pktid = 0xFFFF;

  /* Check if the IGMP timer went off */
  igmp_check();

  /* Loop as long as we have packets to receive */
  while ((readlen = undi_recv()) > sizeof(readbuf.protocol)) {

	/* If this is an ARP packet, do handle it */
	if (readbuf.protocol == UNDI_PROT_ARP) {
		arp_recv(readlen);
		continue;
	}

	/*
	 * Do some checks on the received packet. It must be an IP packet
	 * with a valid version number and header size. Also the checksum
	 * of the header must be correct.
	 */
	hdrsize = (ipp->ip_ihl_version & 0x0F) << 2;
	size = ntohs(ipp->ip_tot_len) - hdrsize;
	if (
	    /* Check for valid version number */
	    (ipp->ip_ihl_version & 0xF0) != (IP_VERSION << 4) ||

	    /* It has to be an IP packet */
	    readbuf.protocol != UNDI_PROT_IP ||

	    /* The header has to have a correct length */
	    hdrsize < IP_MIN_HSIZE ||

	    /* The packet must at least contain an IP header */
	    readlen <= hdrsize ||

	    /* The size value in the packet has to be correct */
	    size > (sizeof(readbuf.buf) - hdrsize) ||

	    /* Check for correct checksum */
	    ipchksum((__u8 *)ipp, hdrsize) != 0)
		continue;

	/* Check if the packet is for UDP */
	if (ipp->ip_protocol == IPPROTO_IGMP)
		igmp_recv(readbuf.buf + hdrsize, size);
	if (ipp->ip_protocol != IPPROTO_UDP)
		continue;
#ifdef UDPDEBUG
	printf("UDP: got packet from %lx.%d for local port %d\n",
			ntohl(ipp->ip_daddr), ntohs(udpp->udp_source),
			ntohs(udpp->udp_dest));
#endif

	/*
	 * Check if this packet is fragmented and there are more frag-
	 * ments following.
	 */
	fragflag = ntohs(ipp->ip_frag_off) & IP_FRAG_MORE;
	fragoff = (ntohs(ipp->ip_frag_off) & IP_FRAG_MASK) << 3;

	/*
	 * If there are no fragments at all or this is the first fragment,
	 * we have to analyze the UDP header and then copy the data contents
	 * into the destination buffer.
	 */
	if (fragoff == 0) {
		/* Check that we have at least a UDP header */
		if (size < sizeof(struct udphdr))
			continue;

		/* Compute total size of data in packet */
		pktid = 0xFFFF;
		totlen = ntohs(udpp->udp_len) - sizeof(struct udphdr);
		if (totlen > bufsize)
			break;

		/* Save UDP header and pseudo header for later */
		rcvpsbuf.srcip = ipp->ip_saddr;
		rcvpsbuf.destip = ipp->ip_daddr;
		rcvpsbuf.pad = 0;
		rcvpsbuf.protocol = ipp->ip_protocol;
		rcvpsbuf.length = udpp->udp_len;
		memcpy((__u8 *)&rcvudpbuf, (__u8 *)udpp, sizeof(rcvudpbuf));

		/*
		 * Check that this packet is for our address and port
		 */
		if (maddr != htonl(IP_ANY) && maddr == ipp->ip_daddr &&
		    mport != 0 && mport == udpp->udp_source)
			mcastflag = TRUE;

		/* Check unicast address */
		if (!mcastflag &&
		    ((usaddr != htonl(IP_ANY) && usaddr != ipp->ip_daddr) ||
		     (usource != 0 && usource != udpp->udp_dest)))
			continue;

		/* Save packet ID and destination address for later */
		pktid = ipp->ip_id;
		dest = ipp->ip_daddr;
	}

	/*
	 * If this is a fragment, we have to copy the fragment data at
	 * the correct place within the destination buffer.
	 */
	if (fragflag || fragoff > 0) {
		/*
		 * Check that this fragment is for the same packet by
		 * using the packet ID and destination address. If this
		 * packet is invalid, discard it.
		 */
		if (ipp->ip_id != pktid ||
		    ipp->ip_daddr != dest)
			continue;

		/* Copy the fragment into the destination buffer */
		if (fragoff == 0) {
			/* First fragment */
			size -= sizeof(struct udphdr);
		} else {
			/* Middle or last fragment */
			fragoff -= sizeof(struct udphdr);
			if ((fragoff + size) > bufsize)
				continue;
		}
		fmemcpy((buf + fragoff), MK_FP(getds(), datap), size);

		/*
		 * Check if this was the last fragment. If not we retry
		 * to read a new packet. Since this routine is non-blocking,
		 * this will likely return us to the caller with zero bytes
		 * read. Therefore, the caller has to repeat this call until
		 * all fragments have been read. In that case, we return the
		 * actual number of bytes read.
		 */
		if (fragflag) {
			/* Wait for next fragment */
			continue;
		} else {
			/* Last fragment, check for checksum */
			goto recv_end;
		}
	}

	/*
	 * If we turn up here the packet is not fragmented. We can simply
	 * copy it into the destination buffer. The packet size can't be
	 * larger than the buffer size as that has already been checked
	 * above.
	 */
	fmemcpy(buf, MK_FP(getds(), datap), totlen);

recv_end:
	/*
	 * Perform a checksum run to see if the packet has been received
	 * correctly. This has only to be done when the checksum in the
	 * UDP packet is not FFFF.
	 */
	if (rcvudpbuf.udp_check && rcvudpbuf.udp_check != 0xFFFF)
		if (udpchksum((__u8 *)&rcvpsbuf, sizeof(rcvpsbuf),
			    (__u8 *)&rcvudpbuf, sizeof(rcvudpbuf),
			    buf, totlen) != 0)
			break;

	/* Setup some global variables. This is required for TFTP and DHCP */
	if (!mcastflag) {
		udest = rcvudpbuf.udp_source;
		udaddr = rcvpsbuf.srcip;
	}

	/* Reset packet ID and return with number of bytes read */
	pktid = 0xFFFF;
	return(totlen);
  }
  return(0);

#undef ipp
#undef udpp
#undef datap
}



/*
 **************************************************************************
 * 
 * Display a spinner on the console
 */
static void spinner(void)
{
  static char spinchars[] = "/-\\|";
  static unsigned long lasttick = 0;
  register unsigned long curtick;

  curtick = get_ticks();
  if (curtick != lasttick) {
	lasttick = curtick;
	charout(spinchars[(curtick >> 1) & 3]);
	charout('\b');
  }
}



/*
 **************************************************************************
 * 
 * Read one packet from a UDP socket. This routine is blocking.
 */
short udp_read(__fptr buf, size_t *bufsize, unsigned short timeout,
					char abortch, bool_t dospin)
{
  size_t buflen;

  /* Wait until we receive a correct packet */
  set_timeout(timeout);
  while (!chk_timeout()) {
	if (statuscallout == FNULL) {
		if (abortch && (chkkey() == abortch))
			return(PXENV_STATUS_ARP_CANCELED);
		if (dospin)
			spinner();
	} else if (statuscallout != (__fptr)(-1)) {
		if (callout(0) != 0)
			return(PXENV_STATUS_ARP_CANCELED);
	}
	if ((buflen = udp_recv(buf, *bufsize)) > 0) {
		*bufsize = buflen;
		return(PXENV_STATUS_SUCCESS);
	}
  }
  return(PXENV_STATUS_ARP_TIMEOUT);
}



/*
 **************************************************************************
 * 
 * Write one packet to a UDP socket
 */
short udp_write(__fptr buf, size_t len)
{
  register short status;
  __u8 udhw[ETH_ALEN];
  struct pshdr psbuf;

  /* Determine the hardware address of the recipient */
  if ((status = ip_resolve(usaddr, udaddr, gateway, udhw)) !=
						PXENV_STATUS_SUCCESS)
	goto write_ret;

#ifdef UDPDEBUG
  printf("UDP: sending %d bytes from %ld.%d to %ld.%d (%x:%x:%x:%x:%x:%x)\n",
			len,
			ntohl(usaddr), ntohs(usource),
			ntohl(udaddr), ntohs(udest),
			udhw[0], udhw[1], udhw[2], udhw[3], udhw[4], udhw[5]);
#endif

  /* Prepare UDP datagram */
  udpbuf.udp_source = usource;
  udpbuf.udp_dest = udest;
  udpbuf.udp_len = htons(len + sizeof(udpbuf));
  udpbuf.udp_check = 0;

  /* Prepare UDP pseudo header for calculating the checksum */
  psbuf.srcip = usaddr;
  psbuf.destip = udaddr;
  psbuf.pad = 0;
  psbuf.protocol = IPPROTO_UDP;
  psbuf.length = udpbuf.udp_len;

  /* Compute UDP checksum */
  udpbuf.udp_check = udpchksum((__u8 *)&psbuf, sizeof(psbuf),
				(__u8 *)&udpbuf, sizeof(udpbuf),
				buf, len);

  /* Prepare IP header */
  ipbuf.ip_ihl_version = ((IP_VERSION << 4) |
					((sizeof(ipbuf)/4) & 0x0f)) & 0xff;
  ipbuf.ip_tos         = IP_STD_TOS;
  ipbuf.ip_tot_len     = htons(sizeof(ipbuf) + sizeof(udpbuf) + len);
  ipbuf.ip_id          = htons(++ippacketid);
  ipbuf.ip_frag_off    = htons(IP_FRAG_NONE);
  ipbuf.ip_check       = 0;
  ipbuf.ip_ttl         = IP_TTL;
  ipbuf.ip_protocol    = IPPROTO_UDP;
  ipbuf.ip_daddr       = udaddr;
  ipbuf.ip_saddr       = usaddr;
  ipbuf.ip_check       = ipchksum((__u8 *)&ipbuf, sizeof(ipbuf));

  /* Finally send out the packet */
  sizes[0] = sizeof(ipbuf);
  sizes[1] = sizeof(udpbuf);
  bufs[0] = MK_FP(getds(), &ipbuf);
  bufs[1] = MK_FP(getds(), &udpbuf);
  if (len > (mymtu - sizeof(struct iphdr) - sizeof(struct udphdr))) {
	/* We need to fragment the outgoing packet */
	size_t fraglen;
	size_t fragoff;

	/* Broadcast packets can't get fragmented */
	if (udaddr == htonl(IP_BROADCAST) ||
	    !((udaddr ^ htonl(IP_BROADCAST)) & ~mynetmask)) {
		status = PXENV_STATUS_FAILURE;
		goto write_ret;
	}

	/* Compute length of each fragment. Has to be on 8 byte boundary */
	fraglen = (mymtu - sizeof(ipbuf)) & ~0x0007;

	/* Setup first packet */
	ipbuf.ip_frag_off = htons(IP_FRAG_MORE);
	ipbuf.ip_tot_len = htons(sizeof(ipbuf) + fraglen);
	ipbuf.ip_check = 0;
	ipbuf.ip_check = ipchksum((__u8 *)&ipbuf, sizeof(ipbuf));

	/* Write first packet */
	sizes[2] = fraglen - sizeof(udpbuf);
	bufs[2] = buf;
	if ((status = undi_send(udhw, UNDI_PROT_IP, 3, sizes, bufs)) !=
							PXENV_STATUS_SUCCESS)
		goto write_ret;
	
	/* Now write all full fragments */
	fragoff = fraglen;
	buf += fraglen - sizeof(udpbuf);
	len -= fraglen - sizeof(udpbuf);
	while (len > fraglen) {
		/* Setup packet for next fragment */
		ipbuf.ip_frag_off = htons(IP_FRAG_MORE | (fragoff >> 3));
		ipbuf.ip_tot_len = htons(sizeof(ipbuf) + fraglen);
		ipbuf.ip_check = 0;
		ipbuf.ip_check = ipchksum((__u8 *)&ipbuf, sizeof(ipbuf));

		/* Write next packet */
		sizes[1] = fraglen;
		bufs[1] = buf;
		if ((status = undi_send(udhw, UNDI_PROT_IP, 2, sizes, bufs)) !=
							PXENV_STATUS_SUCCESS)
			goto write_ret;

		/* Proceed with next packet */
		fragoff += fraglen;
		buf += fraglen;
		len -= fraglen;
	}	

	/* Setup packet for last fragment */
	ipbuf.ip_frag_off = htons(fragoff >> 3);
	ipbuf.ip_tot_len = htons(sizeof(ipbuf) + len);
	ipbuf.ip_check = 0;
	ipbuf.ip_check = ipchksum((__u8 *)&ipbuf, sizeof(ipbuf));

	/* Write last packet */
	sizes[1] = len;
	bufs[1] = buf;
	status = undi_send(udhw, UNDI_PROT_IP, 2, sizes, bufs);
  } else {
	/* No fragmentation necessary */
	sizes[2] = len;
	bufs[2] = buf;
	status = undi_send(udhw, UNDI_PROT_IP, 3, sizes, bufs);
  }

write_ret:
#ifdef UDPDEBUG
  printf("UDP: send status %d\n", status);
#endif
  return(status);
}



/*
 **************************************************************************
 * 
 * Set multicast receive address and server port
 */
short udp_set_mcast(t_ipaddr newmaddr, short newport)
{
  register short ret;
  short count;
  t_ipaddr addrs[2];

  maddr = newmaddr;
  mport = newport;
  count = 0;
  if (newmaddr != htonl(IP_ANY)) {
	addrs[0] = newmaddr;
	addrs[1] = htonl(IP_ALLHOSTS);
	count = 2;
  }
  if ((ret = undi_set_mcast(addrs, count)) == PXENV_STATUS_SUCCESS) {
	if (newmaddr != htonl(IP_ANY))
		igmp_join(maddr);
	else
		igmp_leave();
  }
  return(ret);
}



/*
 **************************************************************************
 * 
 * UDP open PXE API function
 */
static short pxe_udp_open(t_udp_open *params)
{
  register short status;

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

  /* Set my own IP address */
  if (params->src_ip == htonl(IP_ANY))
	params->src_ip = myipaddr;
  usaddr = params->src_ip;

  /* Finally open the network interface */
  if ((status = undi_open()) == PXENV_STATUS_SUCCESS)
	udpoflag = TRUE;
  return(status);
}



/*
 **************************************************************************
 * 
 * UDP close PXE API function
 */
static short pxe_udp_close(t_udp_close *params)
{
  register short status;

  /* Check that we are allowed to close the UDP channel */
  if (!udpoflag)
	return(PXENV_STATUS_UDP_CLOSED);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Close the UDP channel */
  if ((status = udp_close()) == PXENV_STATUS_SUCCESS)
	udpoflag = FALSE;
  return(status);
}



/*
 **************************************************************************
 * 
 * UDP read PXE API function
 */
static short pxe_udp_read(t_udp_read *params)
{
  register short status;
  size_t buflen = 0;
  t_ipaddr oldsaddr = usaddr;

  /* Check that we are allowed to rerad from the UDP channel */
  if (!udpoflag)
	return(PXENV_STATUS_UDP_CLOSED);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Set global variables for this call only */
  usource = params->d_port;
  udest = htons(-1);
  usaddr = params->dest_ip;
  udaddr = htonl(IP_ANY);

  /* Read a packet using UNDI */
  buflen = udp_recv(params->buffer, params->bufsize);
  if (buflen == 0)
	status = PXENV_STATUS_FAILURE;
  else {
	params->src_ip = udaddr;
	params->s_port = udest;
	params->bufsize = buflen;
	status = PXENV_STATUS_SUCCESS;
  }
  usaddr = oldsaddr;
  return(status);
}



/*
 **************************************************************************
 * 
 * UDP write PXE API function
 */
static short pxe_udp_write(t_udp_write *params)
{
  /* Check that we are allowed to write to the UDP channel */
  if (!udpoflag)
	return(PXENV_STATUS_UDP_CLOSED);
  if (tftpoflag)
	return(PXENV_STATUS_TFTP_OPEN);

  /* Set global variables and write the packet */
  if (params->src_port == 0)
	usource = ntohs(PXE_UDP_DEF_PORT);
  else
	usource = params->src_port;
  udest = params->dest_port;
  udaddr = params->ip;
  gateway = params->gw;
  return(udp_write(params->buffer, params->bufsize));
}



/*
 **************************************************************************
 * 
 * Function table for UDP PXE functions
 */
struct functab udp_func_tab = {
	UDP_MINNUM, UDP_MAXNUM,
	{{ pxe_udp_open,	sizeof(t_udp_open)},
	 { pxe_udp_close, 	sizeof(t_udp_close)},
	 { pxe_udp_read, 	sizeof(t_udp_read)},
	 { pxe_udp_write, 	sizeof(t_udp_write)}
	}
};

