// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#define _POSIX_C_SOURCE 200809L  // for strnlen
#include <arpa/inet.h>
#include <ctype.h>
#include <inttypes.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>

#include "internal.h"

// TODO: update this
// RRQ    ->
//        <- DATA or OACK or ERROR
// ACK(0) -> (to confirm reception of OACK)
// ERROR  -> (on OACK with non requested options)
//        <- DATA(1)
// ACK(1) ->

// WRQ     ->
//         <- ACK or OACK or ERROR
// DATA(1) ->
// ERROR   -> (on OACK with non requested options)
//        <- DATA(2)
// ACK(2) ->

// MODE
static const char* kNetascii = "NETASCII";
static const char* kOctet = "OCTET";
static const char* kMail = "MAIL";
static const size_t kMaxMode = 9;  // strlen(NETASCII) + 1

// TSIZE
// Limit transfer to less than 10GB
static const char* kTsize = "TSIZE";
static const size_t kTsizeLen = 5;      // strlen(kTsize)
static const size_t kMaxTsizeOpt = 17;  // strlen(TSIZE) + 1 + strlen(1000000000) + 1

// BLKSIZE
// Max size is 65535 (max IP datagram)
static const char* kBlkSize = "BLKSIZE";
static const size_t kBlkSizeLen = 7;      // strlen(kBlkSize)
static const size_t kMaxBlkSizeOpt = 15;  // kBlkSizeLen + strlen("!") + 1 + strlen(65535) + 1

// TIMEOUT
// Max is 255 (RFC 2349)
static const char* kTimeout = "TIMEOUT";
static const size_t kTimeoutLen = 7;      // strlen(kTimeout)
static const size_t kMaxTimeoutOpt = 13;  // kTimeoutLen + strlen("!") + 1 + strlen(255) + 1;

// WINDOWSIZE
// Max is 65535 (RFC 7440)
static const char* kWindowSize = "WINDOWSIZE";
static const size_t kWindowSizeLen = 10;  // strlen(kWindowSize);
static const size_t kMaxWindowSizeOpt =
    18;  // kWindowSizeLen + strlen("!") + 1 + strlen(65535) + 1;

// Since RRQ and WRQ come before option negotation, they are limited to max TFTP
// blocksize of 512 (RFC 1350 and 2347).
static const size_t kMaxRequestSize = 512;

#if defined(TFTP_HOSTLIB)
// Host (e.g., netcp, bootserver)
#include <time.h>
#define DEBUG 0
#elif defined(TFTP_USERLIB)
// Fuchsia (e.g., netsvc)
#define DEBUG 0
#elif defined(TFTP_EFILIB)
// Bootloader: use judiciously, since the console can easily become overwhelmed and hang
#define DEBUG 0
#else
#error unable to identify target environment
#endif

#if DEBUG
#define xprintf(args...) fprintf(stderr, args)
#else
#define xprintf(args...)
#endif

#define __ATTR_PRINTF(__fmt, __varargs) __attribute__((__format__(__printf__, __fmt, __varargs)))
#define MIN(x, y) ((x) < (y) ? (x) : (y))

static void append_option_name(char** body, size_t* left, const char* name) {
  size_t offset = strlen(name);
  memcpy(*body, name, offset);
  offset++;
  *body += offset;
  *left -= offset;
}

static void __ATTR_PRINTF(5, 6)
    append_option(char** body, size_t* left, const char* name, bool force, const char* fmt, ...) {
  char* bodyp = *body;
  size_t leftp = *left;

  size_t offset = strlen(name);
  memcpy(bodyp, name, offset);
  if (force) {
    bodyp[offset] = '!';
    offset++;
  }
  offset++;
  bodyp += offset;
  leftp -= offset;
  va_list args;
  va_start(args, fmt);
  offset = vsnprintf(bodyp, leftp - 1, fmt, args);
  va_end(args);
  offset++;
  bodyp += offset;
  leftp -= offset;

  *body = bodyp;
  *left = leftp;
}

#define OPCODE(session, msg, value)                                                    \
  do {                                                                                 \
    if (session->use_opcode_prefix) {                                                  \
      (msg)->opcode = htons((value & 0xff) | ((uint16_t)session->opcode_prefix << 8)); \
    } else {                                                                           \
      (msg)->opcode = htons(value);                                                    \
    }                                                                                  \
  } while (0)

#define TRANSMIT_MORE 1
#define TRANSMIT_WAIT_ON_ACK 2

static size_t next_option(char* buffer, size_t len, char** option, char** value) {
  size_t left = len;
  size_t option_len = strnlen(buffer, left);
  if (option_len == len) {
    return 0;
  }

  *option = buffer;
  xprintf("'%s' %ld\n", *option, option_len);
  buffer += option_len + 1;
  left -= option_len + 1;
  size_t value_len = strnlen(buffer, left);
  if (value_len == left) {
    return 0;
  }
  *value = buffer;
  xprintf("'%s' %ld\n", *value, value_len);
  left -= value_len + 1;
  return len - left;
}

/* Build an err packet in resp_buf and set session state to ERROR

     2 bytes        2 bytes    string   1 byte
   +--------------+----------+---------+------+
   | OPCODE_ERROR | ERR_CODE | ERR_MSG |   0  |
   +--------------+----------+---------+------+
*/
static void set_error(tftp_session* session, uint16_t err_code, void* resp_buf, size_t* resp_len,
                      const char* err_msg) {
  tftp_err_msg* resp = resp_buf;
  OPCODE(session, resp, OPCODE_ERROR);
  resp->err_code = htons(err_code);
  size_t err_msg_len = strlen(err_msg);
  size_t max_msg_sz = *resp_len - (sizeof(tftp_err_msg) + 1);
  if (err_msg_len >= max_msg_sz) {
    memcpy(resp->msg, err_msg, max_msg_sz);
    resp->msg[max_msg_sz] = '\0';
    // *resp_len is unchanged - the whole buffer was used
  } else {
    strcpy(resp->msg, err_msg);
    *resp_len = sizeof(tftp_err_msg) + err_msg_len + 1;
  }
  session->state = ERROR;
}

tftp_status tx_data(tftp_session* session, tftp_data_msg* resp, size_t* outlen, void* cookie) {
  session->offset = (session->block_number + session->window_index) * session->block_size;
  *outlen = 0;
  if (session->offset <= session->file_size) {
    session->window_index++;
    OPCODE(session, resp, OPCODE_DATA);
    resp->block = htons(session->block_number + session->window_index);
    size_t len = MIN(session->file_size - session->offset, session->block_size);
    xprintf(" -> Copying block #%" PRIu64 " (size:%zu/%d) from %zu/%zu [%d/%d]\n",
            session->block_number + session->window_index, len, session->block_size,
            session->offset, session->file_size, session->window_index, session->window_size);
    void* buf = resp->data;
    size_t len_remaining = len;
    size_t off = session->offset;
    while (len_remaining > 0) {
      // TODO(tkilbourn): assert that these function pointers are set
      size_t rr = len_remaining;
      tftp_status s = session->file_interface.read(buf, &rr, off, cookie);
      if (s < 0) {
        xprintf("Err reading: %d\n", s);
        return s;
      }
      buf += rr;
      off += rr;
      len_remaining -= rr;
    }
    *outlen = sizeof(*resp) + len;

    if (session->window_index < session->window_size) {
      xprintf(" -> TRANSMIT_MORE(%d < %d)\n", session->window_index, session->window_size);
    } else {
      xprintf(" -> TRANSMIT_WAIT_ON_ACK(%d >= %d)\n", session->window_index, session->window_size);
    }
  } else {
    xprintf(" -> TRANSMIT_WAIT_ON_ACK(completed)\n");
  }
  return TFTP_NO_ERROR;
}

size_t tftp_sizeof_session(void) { return sizeof(tftp_session); }

tftp_status tftp_init(tftp_session** session, void* buffer, size_t size) {
  if (buffer == NULL) {
    return TFTP_ERR_INVALID_ARGS;
  }
  if (size < sizeof(tftp_session)) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }
  *session = buffer;
  tftp_session* s = *session;
  memset(s, 0, sizeof(tftp_session));

  // Sensible defaults for non-negotiated values
  s->file_size = DEFAULT_FILESIZE;
  s->mode = DEFAULT_MODE;
  s->max_timeouts = DEFAULT_MAX_TIMEOUTS;
  s->use_opcode_prefix = DEFAULT_USE_OPCODE_PREFIX;

  return TFTP_NO_ERROR;
}

tftp_status tftp_session_set_file_interface(tftp_session* session, tftp_file_interface* callbacks) {
  if (session == NULL) {
    return TFTP_ERR_INVALID_ARGS;
  }

  session->file_interface = *callbacks;
  return TFTP_NO_ERROR;
}

tftp_status tftp_session_set_transport_interface(tftp_session* session,
                                                 tftp_transport_interface* callbacks) {
  if (session == NULL) {
    return TFTP_ERR_INVALID_ARGS;
  }
  session->transport_interface = *callbacks;
  return TFTP_NO_ERROR;
}

bool tftp_session_has_pending(tftp_session* session) {
  return session->direction == SEND_FILE && session->window_index > 0 &&
         session->window_index < session->window_size &&
         ((session->block_number + session->window_index) * session->block_size) <=
             session->file_size;
}

tftp_status tftp_set_options(tftp_session* session, const uint16_t* block_size,
                             const uint8_t* timeout, const uint16_t* window_size) {
  session->options.mask = 0;
  if (block_size) {
    session->options.block_size = *block_size;
    session->options.mask |= BLOCKSIZE_OPTION;
  }
  if (timeout) {
    session->options.timeout = *timeout;
    session->options.mask |= TIMEOUT_OPTION;
  }
  if (window_size) {
    session->options.window_size = *window_size;
    session->options.mask |= WINDOWSIZE_OPTION;
  }
  return TFTP_NO_ERROR;
}

tftp_status tftp_generate_request(tftp_session* session, tftp_file_direction direction,
                                  const char* local_filename, const char* remote_filename,
                                  tftp_mode mode, size_t datalen, const uint16_t* block_size,
                                  const uint8_t* timeout, const uint16_t* window_size,
                                  void* outgoing, size_t* outlen, uint32_t* timeout_ms) {
  if (*outlen < 2) {
    xprintf("outlen too short: %zd\n", *outlen);
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }

  // The actual options are not set until we get a confirmation OACK message. Until then,
  // we have to assume the TFTP defaults.
  session->block_size = DEFAULT_BLOCKSIZE;
  session->timeout = DEFAULT_TIMEOUT;
  session->window_size = DEFAULT_WINDOWSIZE;

  tftp_msg* ack = outgoing;
  OPCODE(session, ack, (direction == SEND_FILE) ? OPCODE_WRQ : OPCODE_RRQ);
  char* body = ack->data;
  memset(body, 0, *outlen - sizeof(*ack));
  size_t left = *outlen - sizeof(*ack);
  size_t remote_filename_len = strlen(remote_filename);
  if (remote_filename_len + 1 > left - kMaxMode) {
    xprintf("filename too long %zd > %zd\n", remote_filename_len, left - kMaxMode);
    return TFTP_ERR_INVALID_ARGS;
  }
  memcpy(body, remote_filename, remote_filename_len);
  body += remote_filename_len + 1;
  left -= remote_filename_len + 1;
  strncpy(session->filename, local_filename, sizeof(session->filename) - 1);
  session->filename[sizeof(session->filename) - 1] = '\0';
  switch (mode) {
    case MODE_NETASCII:
      append_option_name(&body, &left, kNetascii);
      break;
    case MODE_OCTET:
      append_option_name(&body, &left, kOctet);
      break;
    case MODE_MAIL:
      append_option_name(&body, &left, kMail);
      break;
    default:
      return TFTP_ERR_INVALID_ARGS;
  }
  session->mode = mode;

  if (left < kMaxTsizeOpt) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }
  append_option(&body, &left, kTsize, false, "%zu", datalen);
  session->file_size = datalen;
  tftp_options* sent_opts = &session->client_sent_opts;
  sent_opts->mask = 0;

  if (block_size || session->options.mask & BLOCKSIZE_OPTION) {
    if (left < kMaxBlkSizeOpt) {
      return TFTP_ERR_BUFFER_TOO_SMALL;
    }
    bool force_value;
    if (block_size) {
      force_value = true;
      sent_opts->block_size = *block_size;
    } else {
      force_value = false;
      sent_opts->block_size = session->options.block_size;
    }
    append_option(&body, &left, kBlkSize, force_value, "%" PRIu16, sent_opts->block_size);
    sent_opts->mask |= BLOCKSIZE_OPTION;
  }

  if (timeout || session->options.mask & TIMEOUT_OPTION) {
    if (left < kMaxTimeoutOpt) {
      return TFTP_ERR_BUFFER_TOO_SMALL;
    }
    bool force_value;
    if (timeout) {
      force_value = true;
      sent_opts->timeout = *timeout;
    } else {
      force_value = false;
      sent_opts->timeout = session->options.timeout;
    }
    append_option(&body, &left, kTimeout, force_value, "%" PRIu8, sent_opts->timeout);
    sent_opts->mask |= TIMEOUT_OPTION;
  }

  if (window_size || session->options.mask & WINDOWSIZE_OPTION) {
    if (left < kMaxWindowSizeOpt) {
      return TFTP_ERR_BUFFER_TOO_SMALL;
    }
    bool force_value;
    if (window_size) {
      force_value = true;
      sent_opts->window_size = *window_size;
    } else {
      force_value = false;
      sent_opts->window_size = session->options.window_size;
    }
    append_option(&body, &left, kWindowSize, force_value, "%" PRIu16, sent_opts->window_size);
    sent_opts->mask |= WINDOWSIZE_OPTION;
  }

  *outlen = *outlen - left;
  // Nothing has been negotiated yet so use default
  *timeout_ms = 1000 * session->timeout;

  session->direction = direction;
  session->state = REQ_SENT;
  xprintf("Generated %s request, len=%zu\n", (direction == SEND_FILE) ? "write" : "read", *outlen);
  return TFTP_NO_ERROR;
}

tftp_status tftp_handle_request(tftp_session* session, tftp_file_direction direction, tftp_msg* req,
                                size_t req_len, tftp_msg* resp, size_t* resp_len,
                                uint32_t* timeout_ms, void* cookie) {
  // We could be in REQ_RECEIVED if our OACK was dropped.
  if (session->state != NONE && session->state != REQ_RECEIVED) {
    xprintf("Invalid state transition %d -> %d\n", session->state, REQ_RECEIVED);
    set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "invalid state transition");
    return TFTP_ERR_BAD_STATE;
  }
  // opcode, filename, 0, mode, 0, opt1, 0, value1 ... optN, 0, valueN, 0
  // Max length is 512 no matter
  if (req_len > kMaxRequestSize) {
    xprintf("Write request is too large\n");
    set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "write request is too large");
    return TFTP_ERR_INTERNAL;
  }
  // Skip opcode
  size_t left = req_len - sizeof(*resp);
  char* cur = req->data;
  char *option, *value;
  // filename, 0, mode, 0 can be interpreted like option, 0, value, 0
  size_t offset = next_option(cur, left, &option, &value);
  if (!offset) {
    xprintf("No options\n");
    set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no options");
    return TFTP_ERR_INTERNAL;
  }
  left -= offset;

  xprintf("filename = '%s', mode = '%s'\n", option, value);

  strncpy(session->filename, option, sizeof(session->filename) - 1);
  session->filename[sizeof(session->filename) - 1] = '\0';
  char* mode = value;
  if (!strncasecmp(mode, kNetascii, strlen(kNetascii))) {
    session->mode = MODE_NETASCII;
  } else if (!strncasecmp(mode, kOctet, strlen(kOctet))) {
    session->mode = MODE_OCTET;
  } else if (!strncasecmp(mode, kMail, strlen(kMail))) {
    session->mode = MODE_MAIL;
  } else {
    xprintf("Unknown write request mode\n");
    set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "unknown write request mode");
    return TFTP_ERR_INTERNAL;
  }

  // Initialize the values to TFTP defaults
  session->block_size = DEFAULT_BLOCKSIZE;
  session->timeout = DEFAULT_TIMEOUT;
  session->window_size = DEFAULT_WINDOWSIZE;

  // TODO(tkilbourn): refactor option handling code to share with
  // tftp_handle_oack
  cur += offset;
  bool file_size_seen = false;
  tftp_options requested_options = {.mask = 0};
  tftp_options* override_opts = &session->options;
  while (offset > 0 && left > 0) {
    offset = next_option(cur, left, &option, &value);
    if (!offset) {
      xprintf("No more options\n");
      set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no more options");
      return TFTP_ERR_INTERNAL;
    }

    if (!strncasecmp(option, kTsize, kTsizeLen)) {  // RFC 2349
      if (direction == RECV_FILE) {
        long val = atol(value);
        if (val < 0) {
          xprintf("invalid file size\n");
          set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid file size");
          return TFTP_ERR_INTERNAL;
        }
        session->file_size = val;
      }
      file_size_seen = true;
    } else if (!strncasecmp(option, kBlkSize, kBlkSizeLen)) {  // RFC 2348
      bool force_block_size = (option[kBlkSizeLen] == '!');
      // Valid values range between "8" and "65464" octets, inclusive
      long val = atol(value);
      // TODO(tkilbourn): with an MTU of 1500, shouldn't be more than 1428
      if (val < 8 || val > 65464) {
        xprintf("invalid block size\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid block size");
        return TFTP_ERR_INTERNAL;
      }
      requested_options.block_size = val;
      requested_options.mask |= BLOCKSIZE_OPTION;
      if (force_block_size || !(override_opts->mask & BLOCKSIZE_OPTION)) {
        session->block_size = val;
      } else {
        session->block_size = override_opts->block_size;
      }
    } else if (!strncasecmp(option, kTimeout, kTimeoutLen)) {  // RFC 2349
      bool force_timeout_val = (option[kTimeoutLen] == '!');
      // Valid values range between "1" and "255" seconds inclusive.
      long val = atol(value);
      if (val < 1 || val > 255) {
        xprintf("invalid timeout\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid timeout");
        return TFTP_ERR_INTERNAL;
      }
      requested_options.timeout = val;
      requested_options.mask |= TIMEOUT_OPTION;
      if (force_timeout_val || !(override_opts->mask & TIMEOUT_OPTION)) {
        session->timeout = val;
      } else {
        session->timeout = override_opts->timeout;
      }
    } else if (!strncasecmp(option, kWindowSize, kWindowSizeLen)) {  // RFC 7440
      bool force_window_size = (option[kWindowSizeLen] == '!');
      // The valid values range MUST be between 1 and 65535 blocks, inclusive.
      long val = atol(value);
      if (val < 1 || val > 65535) {
        xprintf("invalid window size\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid window size");
        return TFTP_ERR_INTERNAL;
      }
      requested_options.window_size = val;
      requested_options.mask |= WINDOWSIZE_OPTION;
      if (force_window_size || !(override_opts->mask & WINDOWSIZE_OPTION)) {
        session->window_size = val;
      } else {
        session->window_size = override_opts->window_size;
      }
    } else {
      // Options which the server does not support should be omitted from the
      // OACK; they should not cause an ERROR packet to be generated.
    }

    cur += offset;
    left -= offset;
  }

  char* body = resp->data;
  memset(body, 0, *resp_len - sizeof(*resp));
  left = *resp_len - sizeof(*resp);

  OPCODE(session, resp, OPCODE_OACK);

  // Open file, if we haven't already
  if (session->state == NONE) {
    if (direction == RECV_FILE) {
      if (!session->file_interface.open_write) {
        xprintf("Unable to service write request: no open_write implementation\n");
        set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "internal error");
        return TFTP_ERR_BAD_STATE;
      }
      switch (session->file_interface.open_write(session->filename, session->file_size,
                                                 session->timeout, cookie)) {
        case TFTP_ERR_SHOULD_WAIT:
          // The open_write() callback can return an ERR_SHOULD_WAIT response if it isn't
          // prepared to service another request at the moment and the client should retry
          // later.
          xprintf("Denying write request received when not ready\n");
          set_error(session, TFTP_ERR_CODE_BUSY, resp, resp_len, "not ready to receive");
          session->state = NONE;
          return TFTP_ERR_SHOULD_WAIT;
        case TFTP_NO_ERROR:
          break;
        default:
          xprintf("Could not open file on write request\n");
          set_error(session, TFTP_ERR_CODE_ACCESS_VIOLATION, resp, resp_len,
                    "could not open file for writing");
          return TFTP_ERR_BAD_STATE;
      }
    } else {
      ssize_t file_size;
      if (!session->file_interface.open_read) {
        xprintf("Unable to service read request: no open_read implementation\n");
        set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "internal error");
        return TFTP_ERR_BAD_STATE;
      }

      file_size = session->file_interface.open_read(session->filename, session->timeout, cookie);
      if (file_size == TFTP_ERR_SHOULD_WAIT) {
        // The open_read() callback can return an ERR_SHOULD_WAIT response if it isn't
        // prepared to service another request at the moment and the client should retry
        // later.
        xprintf("Denying read request received when not ready\n");
        set_error(session, TFTP_ERR_CODE_BUSY, resp, resp_len, "not ready to send");
        session->state = NONE;
        return TFTP_ERR_SHOULD_WAIT;
      }
      if (file_size < 0) {
        xprintf("Unable to open file %s for reading\n", session->filename);
        set_error(session, TFTP_ERR_CODE_FILE_NOT_FOUND, resp, resp_len,
                  "could not open file for reading");
        return TFTP_ERR_BAD_STATE;
      }
      session->file_size = file_size;
    }
  }

  if (file_size_seen) {
    append_option(&body, &left, kTsize, false, "%zu", session->file_size);
  } else {
    xprintf("No TSIZE option specified\n");
    set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no TSIZE option");
    if (session->file_interface.close) {
      session->file_interface.close(cookie);
    }
    return TFTP_ERR_BAD_STATE;
  }
  if (requested_options.mask & BLOCKSIZE_OPTION) {
    // TODO(jpoichet) Make sure this block size is possible. Need API upwards to
    // request allocation of block size * window size memory
    append_option(&body, &left, kBlkSize, false, "%d", session->block_size);
  }
  if (requested_options.mask & TIMEOUT_OPTION) {
    // TODO(jpoichet) Make sure this timeout is possible. Need API upwards to
    // request allocation of block size * window size memory
    append_option(&body, &left, kTimeout, false, "%d", session->timeout);
    *timeout_ms = 1000 * session->timeout;
  }
  if (requested_options.mask & WINDOWSIZE_OPTION) {
    append_option(&body, &left, kWindowSize, false, "%d", session->window_size);
  }
  *resp_len = *resp_len - left;
  session->state = REQ_RECEIVED;
  session->direction = direction;

  printf("%s Request Parsed\n", (direction == SEND_FILE) ? "Read" : "Write");
  printf("    Mode       : %s\n", session->mode == MODE_NETASCII ? "netascii"
                                  : session->mode == MODE_OCTET  ? "octet"
                                  : session->mode == MODE_MAIL   ? "mail"
                                                                 : "unrecognized");
  printf("    File Size  : %zu\n", session->file_size);
  printf("Options requested: %08x\n", requested_options.mask);
  printf("    Block Size : %d\n", requested_options.block_size);
  printf("    Timeout    : %d\n", requested_options.timeout);
  printf("    Window Size: %d\n", requested_options.window_size);
  printf("Using options\n");
  printf("    Block Size : %d\n", session->block_size);
  printf("    Timeout    : %d\n", session->timeout);
  printf("    Window Size: %d\n", session->window_size);

  return TFTP_NO_ERROR;
}

tftp_status tftp_handle_wrq(tftp_session* session, tftp_msg* wrq, size_t wrq_len, tftp_msg* resp,
                            size_t* resp_len, uint32_t* timeout_ms, void* cookie) {
  return tftp_handle_request(session, RECV_FILE, wrq, wrq_len, resp, resp_len, timeout_ms, cookie);
}

tftp_status tftp_handle_rrq(tftp_session* session, tftp_msg* rrq, size_t rrq_len, tftp_msg* resp,
                            size_t* resp_len, uint32_t* timeout_ms, void* cookie) {
  return tftp_handle_request(session, SEND_FILE, rrq, rrq_len, resp, resp_len, timeout_ms, cookie);
}

static void tftp_prepare_ack(tftp_session* session, tftp_msg* msg, size_t* msg_len) {
  tftp_data_msg* ack_data = (tftp_data_msg*)msg;
  xprintf(" -> Ack %" PRIu64 "\n", session->block_number);
  session->window_index = 0;
  OPCODE(session, ack_data, OPCODE_ACK);
  ack_data->block = htons(session->block_number & 0xffff);
  *msg_len = sizeof(*ack_data);
}

tftp_status tftp_handle_data(tftp_session* session, tftp_msg* msg, size_t msg_len, tftp_msg* resp,
                             size_t* resp_len, uint32_t* timeout_ms, void* cookie) {
  if ((session->direction == RECV_FILE) &&
      ((session->state == REQ_RECEIVED) || (session->state == FIRST_DATA) ||
       (session->state == RECEIVING_DATA))) {
    session->state = RECEIVING_DATA;
  } else {
    set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "internal error: bad state");
    return TFTP_ERR_INTERNAL;
  }

  tftp_data_msg* data = (tftp_data_msg*)msg;

  uint16_t block_num = ntohs(data->block);

  // The block field of the message is only 16 bits wide. To support large files
  // (> 65535 * blocksize bytes), we allow the block number to wrap. We use signed modulo
  // math to determine the relative location of the block to our current position.
  int16_t block_delta = block_num - (uint16_t)session->block_number;
  xprintf(" <- Block %" PRIu64 " (Last = %" PRIu64 ", Offset = %" PRIu64
          ", Size = %zd, Left = %" PRIu64 ")\n",
          session->block_number + block_delta, session->block_number,
          session->block_number * session->block_size, session->file_size,
          session->file_size - session->block_number * session->block_size);
  if (block_delta == 1) {
    xprintf("Advancing normally + 1\n");
    session->metrics.inorder_blocks++;
    void* buf = data->data;
    size_t len = msg_len - sizeof(tftp_data_msg);
    size_t off = session->block_number * session->block_size;
    while (len > 0) {
      tftp_status ret;
      // TODO(tkilbourn): assert that these function pointers are set
      size_t wr = len;
      ret = session->file_interface.write(buf, &wr, off, cookie);
      if (ret < 0) {
        xprintf("Error writing: %d\n", ret);
        set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "internal error: write failed");
        return ret;
      }
      session->metrics.inorder_bytes += wr;
      buf += wr;
      off += wr;
      len -= wr;
    }
    session->block_number++;
    session->window_index++;
  } else if (block_delta > 1) {
    session->metrics.outoforder_blocks++;
    // Force sending a ACK with the last block_number we received
    xprintf("Skipped: got %" PRIu64 ", expected %" PRIu64 "\n", session->block_number + block_delta,
            session->block_number + 1);
    session->window_index = session->window_size;
    // It's possible that a previous ACK wasn't received, increment the prefix
    if (session->use_opcode_prefix) {
      session->opcode_prefix++;
    }
  }

  if (session->window_index == session->window_size ||
      session->block_number * session->block_size > session->file_size) {
    tftp_prepare_ack(session, resp, resp_len);
    if (block_delta > 1) {
      session->metrics.nacks_sent++;
    } else {
      session->metrics.acks_sent++;
    }
    if (session->block_number * session->block_size > session->file_size) {
      return TFTP_TRANSFER_COMPLETED;
    }
  } else {
    // Nothing to send
    *resp_len = 0;
  }
  return TFTP_NO_ERROR;
}

tftp_status tftp_handle_ack(tftp_session* session, tftp_msg* ack, size_t ack_len, tftp_msg* resp,
                            size_t* resp_len, uint32_t* timeout_ms, void* cookie) {
  if ((session->direction != SEND_FILE) ||
      ((session->state != FIRST_DATA) && (session->state != REQ_RECEIVED) &&
       (session->state != SENDING_DATA))) {
    set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "internal error: bad state");
    return TFTP_ERR_INTERNAL;
  }
  // Need to move forward in data and send it
  tftp_data_msg* ack_data = (void*)ack;
  tftp_data_msg* resp_data = (void*)resp;

  uint16_t ack_block = ntohs(ack_data->block);
  xprintf(" <- Ack %d\n", ack_block);

  // Since we track blocks in 32 bits, but the packets only support 16 bits, calculate the
  // signed 16 bit offset to determine the adjustment to the current position.
  int16_t block_offset = ack_block - (uint16_t)session->block_number;

  if (session->state != FIRST_DATA && session->state != REQ_RECEIVED && block_offset == 0) {
    session->metrics.sas_events++;
    // Don't acknowledge duplicate ACKs, avoiding the "Sorcerer's Apprentice Syndrome"
    *resp_len = 0;
    return TFTP_NO_ERROR;
  }

  if (block_offset < session->window_size) {
    // If it looks like some of our data might have been dropped, modify the prefix
    // before resending.
    if (session->use_opcode_prefix) {
      session->opcode_prefix++;
    }
  }
  session->state = SENDING_DATA;
  session->block_number += block_offset;
  session->window_index = 0;

  if (session->block_number * session->block_size > session->file_size) {
    *resp_len = 0;
    return TFTP_TRANSFER_COMPLETED;
  }

  tftp_status ret = tx_data(session, resp_data, resp_len, cookie);
  if (ret < 0) {
    set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "could not transmit data");
  }
  return ret;
}

tftp_status tftp_handle_error(tftp_session* session, tftp_err_msg* err, size_t err_len,
                              tftp_msg* resp, size_t* resp_len, uint32_t* timeout_ms,
                              void* cookie) {
  if (err_len < sizeof(tftp_err_msg)) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }
  uint16_t err_code = ntohs(err->err_code);

  // There's no need to respond to an error
  *resp_len = 0;

  if (err_code == TFTP_ERR_CODE_BUSY) {
    xprintf("Target busy\n");
    session->state = NONE;
    return TFTP_ERR_SHOULD_WAIT;
  }
  xprintf("Target sent error %d\n", err_code);
  session->state = ERROR;
  return TFTP_ERR_INTERNAL;
}

tftp_status tftp_handle_oack(tftp_session* session, tftp_msg* oack, size_t oack_len, tftp_msg* resp,
                             size_t* resp_len, uint32_t* timeout_ms, void* cookie) {
  xprintf("Option Ack\n");
  if (session->state == REQ_SENT || session->state == FIRST_DATA) {
    session->state = FIRST_DATA;
  }

  size_t left = oack_len - sizeof(*oack);
  char* cur = oack->data;
  size_t offset;
  char *option, *value;

  while (left > 0) {
    offset = next_option(cur, left, &option, &value);
    if (!offset) {
      set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid option format");
      return TFTP_ERR_INTERNAL;
    }

    if (!strncasecmp(option, kTsize, kTsizeLen)) {  // RFC 2349
      if (session->direction == RECV_FILE) {
        session->file_size = atol(value);
      }
      // If we are sending the file, we don't care what value the server wrote in here
    } else if (!strncasecmp(option, kBlkSize, kBlkSizeLen)) {  // RFC 2348
      if (!(session->client_sent_opts.mask & BLOCKSIZE_OPTION)) {
        xprintf("block size not requested\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no block size");
        return TFTP_ERR_INTERNAL;
      }
      // Valid values range between "8" and "65464" octets, inclusive
      long val = atol(value);
      if (val < 8 || val > 65464) {
        xprintf("invalid block size\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid block size");
        return TFTP_ERR_INTERNAL;
      }
      // TODO(tkilbourn): with an MTU of 1500, shouldn't be more than 1428
      session->block_size = val;
    } else if (!strncasecmp(option, kTimeout, kTimeoutLen)) {  // RFC 2349
      if (!(session->client_sent_opts.mask & TIMEOUT_OPTION)) {
        xprintf("timeout not requested\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no timeout");
        return TFTP_ERR_INTERNAL;
      }
      // Valid values range between "1" and "255" seconds inclusive.
      long val = atol(value);
      if (val < 1 || val > 255) {
        xprintf("invalid timeout\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid timeout");
        return TFTP_ERR_INTERNAL;
      }
      session->timeout = val;
    } else if (!strncasecmp(option, kWindowSize, kWindowSizeLen)) {  // RFC 7440
      if (!(session->client_sent_opts.mask & WINDOWSIZE_OPTION)) {
        xprintf("window size not requested\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "no window size");
        return TFTP_ERR_INTERNAL;
      }
      // The valid values range MUST be between 1 and 65535 blocks, inclusive.
      long val = atol(value);
      if (val < 1 || val > 65535) {
        xprintf("invalid window size\n");
        set_error(session, TFTP_ERR_CODE_BAD_OPTIONS, resp, resp_len, "invalid window size");
        return TFTP_ERR_INTERNAL;
      }
      session->window_size = val;
    } else {
      // Options which the server does not support should be omitted from the
      // OACK; they should not cause an ERROR packet to be generated.
    }

    cur += offset;
    left -= offset;
  }
  *timeout_ms = 1000 * session->timeout;

  printf("Options negotiated\n");
  printf("    File Size  : %zu\n", session->file_size);
  printf("    Block Size : %d\n", session->block_size);
  printf("    Timeout    : %d\n", session->timeout);
  printf("    Window Size: %d\n", session->window_size);

  session->offset = 0;
  session->block_number = 0;
  session->window_index = 0;

  if (session->direction == SEND_FILE) {
    tftp_data_msg* resp_data = (void*)resp;
    tftp_status ret = tx_data(session, resp_data, resp_len, cookie);
    if (ret < 0) {
      set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "failure to transmit data");
    }
    return ret;
  } else {
    if (!session->file_interface.open_write ||
        session->file_interface.open_write(session->filename, session->file_size, session->timeout,
                                           cookie)) {
      xprintf("Could not open file on write request\n");
      set_error(session, TFTP_ERR_CODE_UNDEF, resp, resp_len, "could not open file for writing");
      return TFTP_ERR_BAD_STATE;
    }
    tftp_prepare_ack(session, resp, resp_len);
    session->metrics.acks_sent++;
    return TFTP_NO_ERROR;
  }
}

tftp_status tftp_process_msg(tftp_session* session, void* incoming, size_t inlen, void* outgoing,
                             size_t* outlen, uint32_t* timeout_ms, void* cookie) {
  if (inlen < sizeof(tftp_msg)) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }

  tftp_msg* msg = incoming;
  tftp_msg* resp = outgoing;

  // Decode opcode
  uint16_t opcode = ntohs(msg->opcode) & 0xff;
  xprintf("handle_msg opcode=%u length=%d\n", opcode, (int)inlen);

  // Set default timeout
  *timeout_ms = 1000 * session->timeout;

  // Reset timeout count
  session->consecutive_timeouts = 0;

  switch (opcode) {
    case OPCODE_RRQ:
      return tftp_handle_rrq(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    case OPCODE_WRQ:
      return tftp_handle_wrq(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    case OPCODE_DATA:
      return tftp_handle_data(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    case OPCODE_ACK:
      return tftp_handle_ack(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    case OPCODE_ERROR:
      return tftp_handle_error(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    case OPCODE_OACK:
      return tftp_handle_oack(session, incoming, inlen, resp, outlen, timeout_ms, cookie);
    default:
      xprintf("Unknown opcode\n");
      session->state = ERROR;
      return TFTP_ERR_INTERNAL;
  }
}

tftp_status tftp_prepare_data(tftp_session* session, void* outgoing, size_t* outlen, void* cookie) {
  tftp_data_msg* resp_data = outgoing;

  if ((session->block_number + session->window_index) * session->block_size > session->file_size) {
    *outlen = 0;
    return TFTP_TRANSFER_COMPLETED;
  }

  tftp_status ret = tx_data(session, resp_data, outlen, cookie);
  if (ret < 0) {
    set_error(session, TFTP_ERR_CODE_UNDEF, outgoing, outlen, "failure to transmit data");
  }
  return ret;
}

void tftp_session_set_max_timeouts(tftp_session* session, uint16_t max_timeouts) {
  session->max_timeouts = max_timeouts;
}

void tftp_session_set_opcode_prefix_use(tftp_session* session, bool enable) {
  session->use_opcode_prefix = enable;
}

tftp_status tftp_timeout(tftp_session* session, void* msg_buf, size_t* msg_len, size_t buf_sz,
                         void* file_cookie) {
  xprintf("Timeout\n");
  session->metrics.timeouts++;
  if (++session->consecutive_timeouts > session->max_timeouts) {
    return TFTP_ERR_TIMED_OUT;
  }
  // It's possible our previous transmission was dropped because of checksum errors.
  // Use a different opcode prefix when we resend.
  if (session->use_opcode_prefix) {
    session->opcode_prefix++;
  }
  if (session->state == REQ_SENT || session->state == REQ_RECEIVED) {
    // Resend previous message
    return TFTP_NO_ERROR;
  }
  *msg_len = buf_sz;
  if (session->direction == SEND_FILE) {
    // Reset back to the last-acknowledged block
    session->window_index = 0;
    return tftp_prepare_data(session, msg_buf, msg_len, file_cookie);
  } else {
    // ACK up to the last block read
    tftp_prepare_ack(session, msg_buf, msg_len);
    return TFTP_NO_ERROR;
  }
}

#define REPORT_ERR(opts, ...)                               \
  if (opts->err_msg) {                                      \
    snprintf(opts->err_msg, opts->err_msg_sz, __VA_ARGS__); \
  }

typedef struct {
  void* incoming;
  size_t in_buf_sz;
  void* outgoing;
  size_t out_buf_sz;
  char* err_msg;
  size_t err_msg_sz;
} tftp_msg_loop_opts;

static tftp_status tftp_msg_loop(tftp_session* session, void* transport_cookie, void* file_cookie,
                                 tftp_msg_loop_opts* opts, uint32_t timeout_ms) {
  tftp_status ret;
  size_t out_sz = 0;

  do {
    tftp_status send_status;
    int result = session->transport_interface.timeout_set(timeout_ms, transport_cookie);
    if (result < 0) {
      REPORT_ERR(opts, "failed during transport timeout set callback");
      return TFTP_ERR_INTERNAL;
    }

    bool pending = tftp_session_has_pending(session);
    int n = session->transport_interface.recv(opts->incoming, opts->in_buf_sz, !pending,
                                              transport_cookie);
    if (n == TFTP_ERR_TIMED_OUT) {
      if (pending) {
        out_sz = opts->out_buf_sz;
        ret = tftp_prepare_data(session, opts->outgoing, &out_sz, file_cookie);
        if (out_sz) {
          send_status = session->transport_interface.send(opts->outgoing, out_sz, transport_cookie);
          if (send_status != TFTP_NO_ERROR) {
            REPORT_ERR(opts, "failed during transport send callback");
            return send_status;
          }
        }
        if (ret < 0) {
          REPORT_ERR(opts, "failed to prepare data to send");
          return ret;
        }
      } else if (session->state != NONE) {
        ret = tftp_timeout(session, opts->outgoing, &out_sz, opts->out_buf_sz, file_cookie);
        if (ret == TFTP_ERR_TIMED_OUT) {
          REPORT_ERR(opts, "too many consecutive timeouts, aborting");
          return ret;
        }
        if (ret < 0) {
          REPORT_ERR(opts, "failed during timeout processing");
          return ret;
        }
        if (out_sz) {
          send_status = session->transport_interface.send(opts->outgoing, out_sz, transport_cookie);
          if (send_status != TFTP_NO_ERROR) {
            REPORT_ERR(opts, "failed during transport send callback");
            return n;
          }
        }
      }
      continue;
    } else if (n < 0) {
      REPORT_ERR(opts, "failed during transport recv callback");
      return n;
    }

    out_sz = opts->out_buf_sz;
    ret = tftp_process_msg(session, opts->incoming, n, opts->outgoing, &out_sz, &timeout_ms,
                           file_cookie);
    if (out_sz) {
      send_status = session->transport_interface.send(opts->outgoing, out_sz, transport_cookie);
      if (send_status != TFTP_NO_ERROR) {
        REPORT_ERR(opts, "failed during transport send callback");
        return send_status;
      }
    }
    if (ret < 0) {
      REPORT_ERR(opts, "failed to handle message");
      return ret;
    } else if (ret == TFTP_TRANSFER_COMPLETED) {
      return ret;
    }
  } while (1);
}

static tftp_status transfer_file(tftp_session* session, void* transport_cookie, void* file_cookie,
                                 tftp_file_direction xfer_direction, const char* local_filename,
                                 const char* remote_filename, tftp_request_opts* opts) {
  if (!opts || !opts->inbuf || !opts->inbuf_sz || !opts->outbuf || !opts->outbuf_sz) {
    return TFTP_ERR_INVALID_ARGS;
  }

  tftp_status status;

  ssize_t file_size = 0;
  if (xfer_direction == SEND_FILE) {
    file_size = session->file_interface.open_read(local_filename, session->timeout, file_cookie);
    if (file_size < 0) {
      REPORT_ERR(opts, "failed during file open callback");
      return file_size;
    }
  }

  tftp_mode mode = opts->mode ? *opts->mode : TFTP_DEFAULT_CLIENT_MODE;

  size_t out_sz = opts->outbuf_sz;
  uint32_t timeout_ms;
  status = tftp_generate_request(session, xfer_direction, local_filename, remote_filename, mode,
                                 file_size, opts->block_size, opts->timeout, opts->window_size,
                                 opts->outbuf, &out_sz, &timeout_ms);

  const char* xfer_direction_str = (xfer_direction == SEND_FILE) ? "write" : "read";
  if (status < 0) {
    REPORT_ERR(opts, "failed to generate %s request", xfer_direction_str);
    goto done;
  }
  if (!out_sz) {
    REPORT_ERR(opts, "no %s request generated", xfer_direction_str);
    status = TFTP_ERR_INTERNAL;
    goto done;
  }

  status = session->transport_interface.send(opts->outbuf, out_sz, transport_cookie);
  if (status != TFTP_NO_ERROR) {
    REPORT_ERR(opts, "failed during transport send callback");
    goto done;
  }

  tftp_msg_loop_opts msg_loop_opts = {.incoming = opts->inbuf,
                                      .in_buf_sz = opts->inbuf_sz,
                                      .outgoing = opts->outbuf,
                                      .out_buf_sz = opts->outbuf_sz,
                                      .err_msg = opts->err_msg,
                                      .err_msg_sz = opts->err_msg_sz};
  status = tftp_msg_loop(session, transport_cookie, file_cookie, &msg_loop_opts, timeout_ms);

done:
  if ((xfer_direction == SEND_FILE) || (session->state != NONE)) {
    if (session->file_interface.close) {
      session->file_interface.close(file_cookie);
    }
  }
  return status;
}

tftp_status tftp_push_file(tftp_session* session, void* transport_cookie, void* file_cookie,
                           const char* local_filename, const char* remote_filename,
                           tftp_request_opts* opts) {
  return transfer_file(session, transport_cookie, file_cookie, SEND_FILE, local_filename,
                       remote_filename, opts);
}

tftp_status tftp_pull_file(tftp_session* session, void* transport_cookie, void* file_cookie,
                           const char* local_filename, const char* remote_filename,
                           tftp_request_opts* opts) {
  return transfer_file(session, transport_cookie, file_cookie, RECV_FILE, local_filename,
                       remote_filename, opts);
}

tftp_status tftp_service_request(tftp_session* session, void* transport_cookie, void* file_cookie,
                                 tftp_handler_opts* opts) {
  if (!opts || !opts->inbuf || !opts->outbuf || !opts->outbuf_sz) {
    return TFTP_ERR_INVALID_ARGS;
  }
  tftp_msg_loop_opts msg_loop_opts = {.incoming = opts->inbuf,
                                      .in_buf_sz = opts->inbuf_sz,
                                      .outgoing = opts->outbuf,
                                      .out_buf_sz = *opts->outbuf_sz,
                                      .err_msg = opts->err_msg,
                                      .err_msg_sz = opts->err_msg_sz};
  uint32_t timeout_ms = session->timeout * 1000;
  tftp_status status =
      tftp_msg_loop(session, transport_cookie, file_cookie, &msg_loop_opts, timeout_ms);
  if ((session->state != NONE) && session->file_interface.close) {
    session->file_interface.close(file_cookie);
  }
  return status;
}

tftp_status tftp_handle_msg(tftp_session* session, void* transport_cookie, void* file_cookie,
                            tftp_handler_opts* opts) {
  if (!opts || !opts->inbuf || !opts->outbuf || !opts->outbuf_sz) {
    return TFTP_ERR_INVALID_ARGS;
  }
  uint32_t timeout_ms;
  tftp_status ret;
  ret = tftp_process_msg(session, opts->inbuf, opts->inbuf_sz, opts->outbuf, opts->outbuf_sz,
                         &timeout_ms, file_cookie);
  if (*opts->outbuf_sz) {
    tftp_status send_status =
        session->transport_interface.send(opts->outbuf, *opts->outbuf_sz, transport_cookie);
    if (send_status != TFTP_NO_ERROR) {
      REPORT_ERR(opts, "failed during transport send callback");
      return send_status;
    }
  }
  if (ret == TFTP_ERR_SHOULD_WAIT) {
    REPORT_ERR(opts, "request received, host is busy");
  } else if (ret < 0) {
    REPORT_ERR(opts, "handling tftp request failed (file might not exist)");
  } else if (ret == TFTP_TRANSFER_COMPLETED) {
    if (session->file_interface.close) {
      session->file_interface.close(file_cookie);
    }
  } else {
    ret = session->transport_interface.timeout_set(timeout_ms, transport_cookie);
    if (ret < 0) {
      REPORT_ERR(opts, "failed during transport timeout set callback");
    }
  }
  return ret;
}

tftp_status tftp_get_metrics(tftp_session* session, char* buf, size_t buf_sz) {
  snprintf(buf, buf_sz,
           "{"
           "\"inorderblks\": %u,"
           "\"oooblks\": %u,"
           "\"ack\": %u,"
           "\"nack\": %u,"
           "\"timeouts\": %u,"
           "\"sas\": %u,"
           "\"inorderbytes\": %" PRIu64 "}",
           session->metrics.inorder_blocks, session->metrics.outoforder_blocks,
           session->metrics.acks_sent, session->metrics.nacks_sent, session->metrics.timeouts,
           session->metrics.sas_events, session->metrics.inorder_bytes);
  if (strlen(buf) == buf_sz - 1) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }
  return TFTP_NO_ERROR;
}

void tftp_clear_metrics(tftp_session* session) {
  memset(&session->metrics, 0, sizeof(session->metrics));
}