/*
 *  VAS_EBOOT  --  GRand Unified Bootloader
 *  Copyright (C) 2010,2011  Free Software Foundation, Inc.
 *
 *  VAS_EBOOT 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 3 of the License, or
 *  (at your option) any later version.
 *
 *  VAS_EBOOT 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 VAS_EBOOT.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/misc.h>
#include <VasEBoot/net/udp.h>
#include <VasEBoot/net/ip.h>
#include <VasEBoot/net/ethernet.h>
#include <VasEBoot/net/netbuff.h>
#include <VasEBoot/net.h>
#include <VasEBoot/mm.h>
#include <VasEBoot/dl.h>
#include <VasEBoot/file.h>
#include <VasEBoot/i18n.h>

VAS_EBOOT_MOD_LICENSE ("GPLv3+");

/* IP port for the MTFTP server used for Intel's PXE */
enum
  {
    MTFTP_SERVER_PORT = 75,
    MTFTP_CLIENT_PORT = 76,
    /* IP port for the TFTP server */
    TFTP_SERVER_PORT = 69
  };

enum
  {
    TFTP_DEFAULTSIZE_PACKET = 512,
  };

enum
  {
    TFTP_CODE_EOF = 1,
    TFTP_CODE_MORE = 2,
    TFTP_CODE_ERROR = 3,
    TFTP_CODE_BOOT = 4,
    TFTP_CODE_CFG = 5
  };

enum
  {
    TFTP_RRQ = 1,
    TFTP_WRQ = 2,
    TFTP_DATA = 3,
    TFTP_ACK = 4,
    TFTP_ERROR = 5,
    TFTP_OACK = 6
  };

enum
  {
    TFTP_EUNDEF = 0,                   /* not defined */
    TFTP_ENOTFOUND = 1,                /* file not found */
    TFTP_EACCESS = 2,                  /* access violation */
    TFTP_ENOSPACE = 3,                 /* disk full or allocation exceeded */
    TFTP_EBADOP = 4,                   /* illegal TFTP operation */
    TFTP_EBADID = 5,                   /* unknown transfer ID */
    TFTP_EEXISTS = 6,                  /* file already exists */
    TFTP_ENOUSER = 7                  /* no such user */
  };

struct tftphdr {
  VasEBoot_uint16_t opcode;
  union {
    VasEBoot_int8_t rrq[TFTP_DEFAULTSIZE_PACKET];
    struct {
      VasEBoot_uint16_t block;
      VasEBoot_int8_t download[0];
    } data;
    struct {
      VasEBoot_uint16_t block;
    } ack;
    struct {
      VasEBoot_uint16_t errcode;
      VasEBoot_int8_t errmsg[TFTP_DEFAULTSIZE_PACKET];
    } err;
    struct {
      VasEBoot_int8_t data[TFTP_DEFAULTSIZE_PACKET+2];
    } oack;
  } u;
} VAS_EBOOT_PACKED ;


typedef struct tftp_data
{
  VasEBoot_uint64_t file_size;
  VasEBoot_uint64_t block;
  VasEBoot_uint32_t block_size;
  VasEBoot_uint64_t ack_sent;
  int have_oack;
  struct VasEBoot_error_saved save_err;
  VasEBoot_net_udp_socket_t sock;
} *tftp_data_t;

static VasEBoot_err_t
ack (tftp_data_t data, VasEBoot_uint64_t block)
{
  struct tftphdr *tftph_ack;
  VasEBoot_uint8_t nbdata[512];
  struct VasEBoot_net_buff nb_ack;
  VasEBoot_err_t err;

  nb_ack.head = nbdata;
  nb_ack.end = nbdata + sizeof (nbdata);
  VasEBoot_netbuff_clear (&nb_ack);
  VasEBoot_netbuff_reserve (&nb_ack, 512);
  err = VasEBoot_netbuff_push (&nb_ack, sizeof (tftph_ack->opcode)
			   + sizeof (tftph_ack->u.ack.block));
  if (err)
    return err;

  tftph_ack = (struct tftphdr *) nb_ack.data;
  tftph_ack->opcode = VasEBoot_cpu_to_be16_compile_time (TFTP_ACK);
  tftph_ack->u.ack.block = VasEBoot_cpu_to_be16 (block);

  err = VasEBoot_net_send_udp_packet (data->sock, &nb_ack);
  if (err)
    return err;
  data->ack_sent = block;
  return VAS_EBOOT_ERR_NONE;
}

static VasEBoot_err_t
tftp_receive (VasEBoot_net_udp_socket_t sock __attribute__ ((unused)),
	      struct VasEBoot_net_buff *nb,
	      void *f)
{
  VasEBoot_file_t file = f;
  struct tftphdr *tftph = (void *) nb->data;
  tftp_data_t data = file->data;
  VasEBoot_err_t err;
  VasEBoot_uint8_t *ptr;

  if (nb->tail - nb->data < (VasEBoot_ssize_t) sizeof (tftph->opcode))
    {
      VasEBoot_dprintf ("tftp", "TFTP packet too small\n");
      return VAS_EBOOT_ERR_NONE;
    }

  tftph = (struct tftphdr *) nb->data;
  switch (VasEBoot_be_to_cpu16 (tftph->opcode))
    {
    case TFTP_OACK:
      data->block_size = TFTP_DEFAULTSIZE_PACKET;
      data->have_oack = 1;
      for (ptr = nb->data + sizeof (tftph->opcode); ptr < nb->tail;)
	{
	  if (VasEBoot_memcmp (ptr, "tsize\0", sizeof ("tsize\0") - 1) == 0)
	    data->file_size = VasEBoot_strtoul ((char *) ptr + sizeof ("tsize\0")
					    - 1, 0, 0);
	  if (VasEBoot_memcmp (ptr, "blksize\0", sizeof ("blksize\0") - 1) == 0)
	    data->block_size = VasEBoot_strtoul ((char *) ptr + sizeof ("blksize\0")
					     - 1, 0, 0);
	  while (ptr < nb->tail && *ptr)
	    ptr++;
	  ptr++;
	}
      data->block = 0;
      VasEBoot_netbuff_free (nb);
      err = ack (data, 0);
      VasEBoot_error_save (&data->save_err);
      return VAS_EBOOT_ERR_NONE;
    case TFTP_DATA:
      if (nb->tail - nb->data < (VasEBoot_ssize_t) (sizeof (tftph->opcode)
						+ sizeof (tftph->u.data.block)))
	{
	  VasEBoot_dprintf ("tftp", "TFTP packet too small\n");
	  return VAS_EBOOT_ERR_NONE;
	}

      /*
       * Ack old/retransmitted block.
       *
       * The block number is a 16-bit counter, thus the maximum file size that
       * could be transfered is 65535 * block size. Most TFTP hosts support to
       * roll-over the block counter to allow unlimited transfer file size.
       *
       * This behavior is not defined in the RFC 1350 [0] but is implemented by
       * most TFTP clients and hosts.
       *
       * [0]: https://tools.ietf.org/html/rfc1350
       */
      if (VasEBoot_be_to_cpu16 (tftph->u.data.block) < ((VasEBoot_uint16_t) (data->block + 1)))
	ack (data, VasEBoot_be_to_cpu16 (tftph->u.data.block));
      /* Ignore unexpected block. */
      else if (VasEBoot_be_to_cpu16 (tftph->u.data.block) > ((VasEBoot_uint16_t) (data->block + 1)))
	VasEBoot_dprintf ("tftp", "TFTP unexpected block # %d\n", tftph->u.data.block);
      else
	{
	  unsigned size;

	  if (file->device->net->packs.count < 50)
	    {
	      err = ack (data, data->block + 1);
	      if (err)
		return err;
	    }
	  else
	    file->device->net->stall = 1;

	  err = VasEBoot_netbuff_pull (nb, sizeof (tftph->opcode) +
				   sizeof (tftph->u.data.block));
	  if (err)
	    return err;
	  size = nb->tail - nb->data;

	  data->block++;
	  if (size < data->block_size)
	    {
	      if (data->ack_sent < data->block)
		ack (data, data->block);
	      file->device->net->eof = 1;
	      file->device->net->stall = 1;
	      VasEBoot_net_udp_close (data->sock);
	      data->sock = NULL;
	    }
	  /*
	   * Prevent garbage in broken cards. Is it still necessary
	   * given that IP implementation has been fixed?
	   */
	  if (size > data->block_size)
	    {
	      err = VasEBoot_netbuff_unput (nb, size - data->block_size);
	      if (err)
		return err;
	    }
	  /* If there is data, puts packet in socket list. */
	  if ((nb->tail - nb->data) > 0)
	    {
	      VasEBoot_net_put_packet (&file->device->net->packs, nb);
	      /* Do not free nb. */
	      return VAS_EBOOT_ERR_NONE;
	    }
	}
      VasEBoot_netbuff_free (nb);
      return VAS_EBOOT_ERR_NONE;
    case TFTP_ERROR:
      data->have_oack = 1;
      VasEBoot_error (VAS_EBOOT_ERR_IO, "%s", tftph->u.err.errmsg);
      VasEBoot_error_save (&data->save_err);
      VasEBoot_netbuff_free (nb);
      return VAS_EBOOT_ERR_NONE;
    default:
      VasEBoot_netbuff_free (nb);
      return VAS_EBOOT_ERR_NONE;
    }
}

/*
 * Create a normalized copy of the filename. Compress any string of consecutive
 * forward slashes to a single forward slash.
 */
static void
VasEBoot_normalize_filename (char *normalized, const char *filename, int c)
{
  char *dest = normalized;
  const char *src = filename;

  while (*src != '\0' && c > 0)
    {
      if (src[0] == '/' && src[1] == '/')
        src++;
      else {
        c--;
        *dest++ = *src++;
      }
    }
  *dest = '\0';
}

static VasEBoot_err_t
tftp_open (struct VasEBoot_file *file, const char *filename)
{
  struct tftphdr *tftph;
  char *rrq;
  int i;
  int rrqlen, rrqsize;
  int hdrlen;
  VasEBoot_uint8_t open_data[1500];
  struct VasEBoot_net_buff nb;
  tftp_data_t data;
  VasEBoot_err_t err;
  VasEBoot_uint8_t *nbd;
  VasEBoot_net_network_level_address_t addr;
  int port = file->device->net->port;

  data = VasEBoot_zalloc (sizeof (*data));
  if (!data)
    return VasEBoot_errno;

  nb.head = open_data;
  nb.end = open_data + sizeof (open_data);
  VasEBoot_netbuff_clear (&nb);

  VasEBoot_netbuff_reserve (&nb, 1500);
  err = VasEBoot_netbuff_push (&nb, sizeof (*tftph));
  if (err)
    {
      VasEBoot_free (data);
      return err;
    }

  tftph = (struct tftphdr *) nb.data;

  tftph->opcode = VasEBoot_cpu_to_be16_compile_time (TFTP_RRQ);

  rrq = (char *) tftph->u.rrq;
  rrqsize = sizeof (tftph->u.rrq);

  /*
   * Copy and normalize the filename to work-around issues on some TFTP
   * servers when file names are being matched for remapping.
   */
  VasEBoot_normalize_filename (rrq, filename, rrqsize);

  rrqlen = VasEBoot_strlen (rrq) + 1;
  rrq += VasEBoot_strlen (rrq) + 1;

  /* Verify there is enough space for the remaining components. */
  rrqlen += VasEBoot_strlen ("octet") + 1;
  rrqlen += VasEBoot_strlen ("blksize") + 1;
  rrqlen += VasEBoot_strlen ("1024") + 1;
  rrqlen += VasEBoot_strlen ("tsize") + 1;
  rrqlen += VasEBoot_strlen ("0") + 1;

  if (rrqlen >= rrqsize) {
    VasEBoot_free (data);
    return VasEBoot_error (VAS_EBOOT_ERR_BAD_FILENAME, N_("filename too long"));
  }

  VasEBoot_strcpy (rrq, "octet");
  rrq += VasEBoot_strlen ("octet") + 1;

  VasEBoot_strcpy (rrq, "blksize");
  rrq += VasEBoot_strlen ("blksize") + 1;

  VasEBoot_strcpy (rrq, "1024");
  rrq += VasEBoot_strlen ("1024") + 1;

  VasEBoot_strcpy (rrq, "tsize");
  rrq += VasEBoot_strlen ("tsize") + 1;

  VasEBoot_strcpy (rrq, "0");
  rrq += VasEBoot_strlen ("0") + 1;
  hdrlen = sizeof (tftph->opcode) + rrqlen;

  err = VasEBoot_netbuff_unput (&nb, nb.tail - (nb.data + hdrlen));
  if (err)
    {
      VasEBoot_free (data);
      return err;
    }

  file->not_easily_seekable = 1;
  file->data = data;

  err = VasEBoot_net_resolve_address (file->device->net->server, &addr);
  if (err)
    {
      VasEBoot_dprintf ("tftp", "Address resolution failed: %d\n", err);
      VasEBoot_dprintf ("tftp", "file_size is %" PRIuVAS_EBOOT_UINT64_T ", block_size is %" PRIuVAS_EBOOT_UINT32_T "\n",
		    data->file_size, data->block_size);
      VasEBoot_free (data);
      return err;
    }

  data->sock = VasEBoot_net_udp_open (addr,
				  port ? port : TFTP_SERVER_PORT, tftp_receive,
				  file);
  if (!data->sock)
    {
      VasEBoot_free (data);
      return VasEBoot_errno;
    }

  /* Receive OACK packet.  */
  nbd = nb.data;
  for (i = 0; i < VAS_EBOOT_NET_TRIES; i++)
    {
      nb.data = nbd;
      err = VasEBoot_net_send_udp_packet (data->sock, &nb);
      if (err)
	{
	  VasEBoot_net_udp_close (data->sock);
	  VasEBoot_free (data);
	  return err;
	}
      VasEBoot_net_poll_cards (VAS_EBOOT_NET_INTERVAL + (i * VAS_EBOOT_NET_INTERVAL_ADDITION),
                           &data->have_oack);
      if (data->have_oack)
	break;
    }

  if (!data->have_oack)
    VasEBoot_error (VAS_EBOOT_ERR_TIMEOUT, N_("time out opening `%s'"), filename);
  else
    VasEBoot_error_load (&data->save_err);
  if (VasEBoot_errno)
    {
      if (data->sock != NULL)
	{
	  VasEBoot_net_udp_close (data->sock);
	  data->sock = NULL;
	}
      VasEBoot_free (data);
      file->data = NULL;
      return VasEBoot_errno;
    }

  file->size = data->file_size;

  return VAS_EBOOT_ERR_NONE;
}

static VasEBoot_err_t
tftp_close (struct VasEBoot_file *file)
{
  tftp_data_t data = file->data;

  if (data->sock)
    {
      VasEBoot_uint8_t nbdata[512];
      VasEBoot_err_t err;
      struct VasEBoot_net_buff nb_err;
      struct tftphdr *tftph;

      nb_err.head = nbdata;
      nb_err.end = nbdata + sizeof (nbdata);

      VasEBoot_netbuff_clear (&nb_err);
      VasEBoot_netbuff_reserve (&nb_err, 512);
      err = VasEBoot_netbuff_push (&nb_err, sizeof (tftph->opcode)
			       + sizeof (tftph->u.err.errcode)
			       + sizeof ("closed"));
      if (!err)
	{
	  tftph = (struct tftphdr *) nb_err.data;
	  tftph->opcode = VasEBoot_cpu_to_be16_compile_time (TFTP_ERROR);
	  tftph->u.err.errcode = VasEBoot_cpu_to_be16_compile_time (TFTP_EUNDEF);
	  VasEBoot_memcpy (tftph->u.err.errmsg, "closed", sizeof ("closed"));

	  err = VasEBoot_net_send_udp_packet (data->sock, &nb_err);
	}
      if (err)
	VasEBoot_print_error ();
      VasEBoot_net_udp_close (data->sock);
    }
  VasEBoot_free (data);
  file->data = NULL;
  return VAS_EBOOT_ERR_NONE;
}

static VasEBoot_err_t
tftp_packets_pulled (struct VasEBoot_file *file)
{
  tftp_data_t data = file->data;
  if (file->device->net->packs.count >= 50)
    return 0;

  if (!file->device->net->eof)
    file->device->net->stall = 0;
  if (data->ack_sent >= data->block)
    return 0;
  return ack (data, data->block);
}

static struct VasEBoot_net_app_protocol VasEBoot_tftp_protocol =
  {
    .name = "tftp",
    .open = tftp_open,
    .close = tftp_close,
    .packets_pulled = tftp_packets_pulled
  };

VAS_EBOOT_MOD_INIT (tftp)
{
  VasEBoot_net_app_level_register (&VasEBoot_tftp_protocol);
}

VAS_EBOOT_MOD_FINI (tftp)
{
  VasEBoot_net_app_level_unregister (&VasEBoot_tftp_protocol);
}