// 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.

// This is an example binary to exercise ulib/tftp. It runs on Linux or macOS.

#define _POSIX_C_SOURCE 200809L
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#include "tftp/tftp.h"

#define BLOCKSZ 1024
#define WINSZ 64

#define DROPRATE 100

#define SCRATCHSZ 2048
static char scratch[SCRATCHSZ];
static char out_scratch[SCRATCHSZ];
static char in_scratch[SCRATCHSZ];

typedef struct connection connection_t;

struct tftp_file {
  int fd;
  size_t size;
};

struct connection {
  int socket;
  struct sockaddr_in out_addr;
  struct sockaddr_in in_addr;
  uint32_t previous_timeout_ms;
};

tftp_status connection_send(void* data, size_t len, void* transport_cookie) {
  connection_t* connection = (connection_t*)transport_cookie;
#if DROPRATE != 0
  if (rand() % DROPRATE == 0) {
    fprintf(stderr, "DROP\n");
    return len;
  }
#endif
  uint8_t* msg = data;
  uint16_t opcode = ntohs(*(uint16_t*)msg);
  fprintf(stderr, "sending opcode=%u\n", opcode);
  ssize_t send_result = sendto(connection->socket, data, len, 0,
                               (struct sockaddr*)&connection->out_addr, sizeof(struct sockaddr_in));
  if (send_result != len) {
    return TFTP_ERR_IO;
  }
  return TFTP_NO_ERROR;
}

int connection_receive(void* data, size_t len, bool block, void* transport_cookie) {
  connection_t* connection = (connection_t*)transport_cookie;
  socklen_t server_len;
  int fl = fcntl(connection->socket, F_GETFL, 0);
  if (fl < 0) {
    int e = errno;
    fprintf(stderr, "could not get socket flags: %d\n", errno);
    errno = e;
    return -1;
  }
  if (block) {
    fl &= ~O_NONBLOCK;
  } else {
    fl |= O_NONBLOCK;
  }
  int ret = fcntl(connection->socket, F_SETFL, fl);
  if (ret < 0) {
    int e = errno;
    fprintf(stderr, "could not set socket flags: %d\n", errno);
    errno = e;
    return -1;
  }
  ssize_t recv_result = recvfrom(connection->socket, data, len, 0,
                                 (struct sockaddr*)&connection->in_addr, &server_len);
  if (recv_result < 0) {
    if ((errno == EAGAIN) || (errno == EWOULDBLOCK))
      return TFTP_ERR_TIMED_OUT;
    fprintf(stderr, "failed during recvfrom: errno=%d\n", (int)errno);
    return TFTP_ERR_INTERNAL;
  }
  return recv_result;
}

int connection_set_timeout(uint32_t timeout_ms, void* transport_cookie) {
  connection_t* connection = (connection_t*)transport_cookie;
  if (connection->previous_timeout_ms != timeout_ms && timeout_ms > 0) {
    fprintf(stdout, "Setting timeout to %dms\n", timeout_ms);
    connection->previous_timeout_ms = timeout_ms;
    struct timeval tv;
    tv.tv_sec = timeout_ms / 1000;
    tv.tv_usec = 1000 * (timeout_ms - 1000 * tv.tv_sec);
    return setsockopt(connection->socket, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
  }
  return 0;
}

connection_t* create_connection(const char* hostname, int incoming_port, int outgoing_port) {
  connection_t* connection = (connection_t*)malloc(sizeof(connection_t));
  memset(connection, 0, sizeof(connection_t));

  struct hostent* server;

  if ((connection->socket = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
    fprintf(stderr, "Cannot create socket\n");
    goto err;
  }

  if (!(server = gethostbyname(hostname))) {
    fprintf(stderr, "Could not resolve host '%s'\n", hostname);
    goto err;
  }

  memset(&connection->out_addr, 0, sizeof(struct sockaddr_in));
  connection->out_addr.sin_family = AF_INET;
  connection->out_addr.sin_port = htons(outgoing_port);
  void* server_addr = server->h_addr_list[0];
  memcpy(&connection->out_addr.sin_addr.s_addr, server_addr, server->h_length);

  memset(&connection->in_addr, 0, sizeof(struct sockaddr_in));
  connection->in_addr.sin_family = AF_INET;
  connection->in_addr.sin_port = htons(incoming_port);
  memcpy(&connection->in_addr.sin_addr.s_addr, server_addr, server->h_length);
  if (bind(connection->socket, (struct sockaddr*)&connection->in_addr,
           sizeof(struct sockaddr_in)) == -1) {
    fprintf(stderr, "Could not bind\n");
    goto err;
  }

  connection->previous_timeout_ms = 0;
  return connection;

err:
  if (connection->socket)
    close(connection->socket);
  free(connection);
  return NULL;
}

void print_usage(void) {
  fprintf(stdout, "tftp (-s filename|-r)\n");
  fprintf(stdout, "\t -s filename to send the provided file\n");
  fprintf(stdout, "\t -r to receive a file\n");
}

ssize_t open_read_file(const char* filename, void* file_cookie) {
  fprintf(stdout, "Opening %s for reading\n", filename);
  int fd = open(filename, O_RDONLY);
  if (fd < 0) {
    fprintf(stderr, "could not open file: err=%d\n", errno);
    return TFTP_ERR_IO;
  }
  struct tftp_file* f = (struct tftp_file*)file_cookie;
  f->fd = fd;
  struct stat st;
  if (fstat(f->fd, &st) < 0) {
    fprintf(stderr, "could not get file size: err=%d\n", errno);
    return TFTP_ERR_IO;
  }
  return st.st_size;
}

tftp_status open_write_file(const char* filename, size_t size, void* file_cookie) {
  fprintf(stdout, "Opening %s for writing\n", filename);
  int fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH);
  if (fd < 0) {
    fprintf(stderr, "could not open file: err=%d\n", errno);
    return TFTP_ERR_IO;
  }
  struct tftp_file* f = (struct tftp_file*)file_cookie;
  f->fd = fd;
  return TFTP_NO_ERROR;
}

tftp_status read_file(void* data, size_t* length, off_t offset, void* file_cookie) {
  int fd = ((struct tftp_file*)file_cookie)->fd;
  ssize_t n = pread(fd, data, *length, offset);
  if (n < 0) {
    fprintf(stderr, "could not read file: offset %jd, err=%d\n", (intmax_t)offset, errno);
    return n;
  }
  *length = n;
  return TFTP_NO_ERROR;
}

tftp_status write_file(const void* data, size_t* length, off_t offset, void* file_cookie) {
  struct tftp_file* f = file_cookie;
  int fd = f->fd;
  ssize_t n = pwrite(fd, data, *length, offset);
  if (n < 0) {
    fprintf(stderr, "could not write file: offset %jd, err=%d\n", (intmax_t)offset, errno);
    return n;
  }
  *length = n;
  f->size = offset + *length;
  return TFTP_NO_ERROR;
}

void close_file(void* file_cookie) {
  struct tftp_file* f = file_cookie;
  close(f->fd);
}

int tftp_send_file_wrapper(tftp_session* session, connection_t* connection, const char* filename) {
  uint16_t block_size = BLOCKSZ;
  uint16_t window_size = WINSZ;
  tftp_set_options(session, &block_size, NULL, &window_size);

  struct tftp_file file_cookie;
  char err_msg[128];
  tftp_request_opts options = {0};
  options.inbuf = in_scratch;
  options.inbuf_sz = sizeof(in_scratch);
  options.outbuf = out_scratch;
  options.outbuf_sz = sizeof(out_scratch);
  options.err_msg = err_msg;
  options.err_msg_sz = sizeof(err_msg);
  tftp_status send_result =
      tftp_push_file(session, connection, &file_cookie, filename, "zircon.bin", &options);
  if (send_result == TFTP_NO_ERROR) {
    return 0;
  }
  fprintf(stderr, "%s\n", err_msg);
  return -1;
}

int tftp_receive_file_wrapper(tftp_session* session, connection_t* connection) {
  struct tftp_file file_cookie;
  char err_msg[128];
  tftp_handler_opts options = {0};
  options.inbuf = in_scratch;
  options.inbuf_sz = sizeof(in_scratch);
  options.outbuf = out_scratch;
  size_t outbuf_sz = sizeof(out_scratch);
  options.outbuf_sz = &outbuf_sz;
  options.err_msg = err_msg;
  options.err_msg_sz = sizeof(err_msg);

  tftp_status status;
  do {
    status = tftp_service_request(session, connection, &file_cookie, &options);
  } while (status == TFTP_NO_ERROR || status == TFTP_ERR_TIMED_OUT);
  if (status == TFTP_TRANSFER_COMPLETED)
    return 0;
  fprintf(stderr, "%s\n", err_msg);
  return 1;
}

int main(int argc, char* argv[]) {
  const char* hostname = "127.0.0.1";
  int port = 2343;

  srand(time(NULL));

  if (argc < 2) {
    print_usage();
    return 1;
  }

  tftp_session* session = NULL;
  if (SCRATCHSZ < tftp_sizeof_session()) {
    fprintf(stderr, "Need more space for tftp session: %d < %zu\n", SCRATCHSZ,
            tftp_sizeof_session());
    return -1;
  }
  if (tftp_init(&session, scratch, SCRATCHSZ)) {
    fprintf(stderr, "Failed to initialize TFTP Session\n");
    return -1;
  }

  tftp_file_interface file_interface = {open_read_file, open_write_file, read_file, write_file,
                                        close_file};
  tftp_session_set_file_interface(session, &file_interface);
  tftp_transport_interface transport_interface = {connection_send, connection_receive,
                                                  connection_set_timeout};
  tftp_session_set_transport_interface(session, &transport_interface);

  connection_t* connection;
  if (!strncmp(argv[1], "-s", 2)) {
    if (argc < 3) {
      print_usage();
      return 1;
    }
    connection = create_connection(hostname, port, port + 1);
    if (!connection) {
      return -1;
    }
    return tftp_send_file_wrapper(session, connection, argv[2]);
  } else if (!strncmp(argv[1], "-r", 2)) {
    connection = create_connection(hostname, port + 1, port);
    if (!connection) {
      return -1;
    }
    return tftp_receive_file_wrapper(session, connection);
  } else {
    print_usage();
    return 2;
  }
  return 0;
}