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

#include <tftp/tftp.h>
#include <unittest/unittest.h>

#include <limits.h>
#include <pthread.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <atomic>

// This test simulates a tftp file transfer by running two threads. Both the
// file and transport interfaces are implemented in memory buffers.

typedef enum { DIR_SEND, DIR_RECEIVE } xfer_dir_t;

struct test_params {
  xfer_dir_t direction;
  uint32_t filesz;
  uint16_t winsz;
  uint16_t blksz;
};

static uint8_t* src_file;
static uint8_t* dst_file;

/* FAUX FILES INTERFACE */

typedef struct {
  uint8_t* buf;
  char filename[PATH_MAX + 1];
  size_t filesz;
} file_info_t;

// Allocate our src and dst buffers, filling both with random values.
int initialize_files(struct test_params* tp) {
  src_file = reinterpret_cast<uint8_t*>(malloc(tp->filesz));
  dst_file = reinterpret_cast<uint8_t*>(malloc(tp->filesz));

  if (!src_file || !dst_file) {
    return 1;
  }

  int* src_as_ints = (int*)src_file;
  int* dst_as_ints = (int*)dst_file;

  size_t ndx;
  srand(0);
  for (ndx = 0; ndx < tp->filesz / sizeof(int); ndx++) {
    src_as_ints[ndx] = rand();
    dst_as_ints[ndx] = rand();
  }
  for (ndx = (tp->filesz / sizeof(int)) * sizeof(int); ndx < tp->filesz; ndx++) {
    src_file[ndx] = static_cast<uint8_t>(rand());
    dst_file[ndx] = static_cast<uint8_t>(rand());
  }

  return 0;
}

int compare_files(size_t filesz) { return memcmp(src_file, dst_file, filesz); }

const char* file_get_filename(file_info_t* file_info) { return file_info->filename; }

ssize_t file_open_read(const char* filename, void* file_cookie) {
  auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
  file_info->buf = src_file;
  strncpy(file_info->filename, filename, PATH_MAX);
  file_info->filename[PATH_MAX] = '\0';
  return file_info->filesz;
}

tftp_status file_open_write(const char* filename, size_t size, void* file_cookie) {
  auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
  file_info->buf = dst_file;
  file_info->filename[PATH_MAX] = '\0';
  strncpy(file_info->filename, filename, PATH_MAX);
  return TFTP_NO_ERROR;
}

// Every SHORT_READ_FREQ writes we will read a smaller amount instead, to verify behavior when
// a read operation returns a length less than the requested amount.
#define SHORT_READ_FREQ 10

tftp_status file_read(void* data, size_t* length, off_t offset, void* file_cookie) {
  auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
  if ((size_t)offset > file_info->filesz) {
    // Something has gone wrong in libtftp
    return TFTP_ERR_INTERNAL;
  }
  if ((offset + *length) > file_info->filesz) {
    *length = file_info->filesz - offset;
  }
  static size_t read_count = 0;
  if (read_count++ % SHORT_READ_FREQ == 0) {
    *length /= 2;
  }
  memcpy(data, &file_info->buf[offset], *length);
  return TFTP_NO_ERROR;
}

// Every SHORT_WRITE_FREQ writes we will write a smaller amount instead, to verify behavior when
// a write operation returns a length less than the requested amount.
#define SHORT_WRITE_FREQ 10

tftp_status file_write(const void* data, size_t* length, off_t offset, void* file_cookie) {
  auto* file_info = reinterpret_cast<file_info_t*>(file_cookie);
  if (((size_t)offset > file_info->filesz) || ((offset + *length) > file_info->filesz)) {
    // Something has gone wrong in libtftp
    return TFTP_ERR_INTERNAL;
  }
  static size_t write_count = 0;
  if (write_count++ % SHORT_WRITE_FREQ == 0) {
    *length /= 2;
  }
  memcpy(&file_info->buf[offset], data, *length);
  return TFTP_NO_ERROR;
}

void file_close(void* file_cookie) {}

/* FAUX SOCKET INTERFACE */

#define FAKE_SOCK_BUF_SZ 65536
typedef struct {
  uint8_t buf[FAKE_SOCK_BUF_SZ];
  size_t size = FAKE_SOCK_BUF_SZ;
  std::atomic<size_t> read_ndx;
  std::atomic<size_t> write_ndx;
} fake_socket_t;
static fake_socket_t client_out_socket;
static fake_socket_t server_out_socket;

typedef struct {
  fake_socket_t* in_sock;
  fake_socket_t* out_sock;
} transport_info_t;

void clear_sockets(void) {
  client_out_socket.read_ndx.store(0);
  client_out_socket.write_ndx.store(0);
  server_out_socket.read_ndx.store(0);
  server_out_socket.write_ndx.store(0);
}

// Initialize "sockets" for either client or server.
void transport_init(transport_info_t* transport_info, bool is_server) {
  if (is_server) {
    transport_info->in_sock = &client_out_socket;
    transport_info->out_sock = &server_out_socket;
  } else {
    transport_info->in_sock = &server_out_socket;
    transport_info->out_sock = &client_out_socket;
  }
}

// Write to our circular message buffer.
void write_to_buf(fake_socket_t* sock, void* data, size_t size) {
  uint8_t* in_buf = reinterpret_cast<uint8_t*>(data);
  uint8_t* out_buf = sock->buf;
  size_t curr_offset = sock->write_ndx.load() % sock->size;
  if (curr_offset + size <= sock->size) {
    memcpy(&out_buf[curr_offset], in_buf, size);
  } else {
    size_t first_size = sock->size - curr_offset;
    size_t second_size = size - first_size;
    memcpy(out_buf + curr_offset, in_buf, first_size);
    memcpy(out_buf, in_buf + first_size, second_size);
  }
  sock->write_ndx.fetch_add(size);
}

// Send a message. Note that the buffer's read_ndx and write_ndx don't wrap,
// which makes it easier to recognize underflow.
tftp_status transport_send(void* data, size_t len, void* transport_cookie) {
  auto* transport_info = reinterpret_cast<transport_info_t*>(transport_cookie);
  fake_socket_t* sock = transport_info->out_sock;
  while ((sock->write_ndx.load() + sizeof(len) + len - sock->read_ndx.load()) > sock->size) {
    // Wait for the other thread to catch up
    usleep(10);
  }
  write_to_buf(sock, &len, sizeof(len));
  write_to_buf(sock, data, len);
  return TFTP_NO_ERROR;
}

// Read from our circular message buffer. If |move_ptr| is false, just peeks at
// the data (reads without updating the read pointer).
void read_from_buf(fake_socket_t* sock, void* data, size_t size, bool move_ptr) {
  uint8_t* in_buf = sock->buf;
  uint8_t* out_buf = reinterpret_cast<uint8_t*>(data);
  size_t curr_offset = sock->read_ndx.load() % sock->size;
  if (curr_offset + size <= sock->size) {
    memcpy(out_buf, &in_buf[curr_offset], size);
  } else {
    size_t first_size = sock->size - curr_offset;
    size_t second_size = size - first_size;
    memcpy(out_buf, in_buf + curr_offset, first_size);
    memcpy(out_buf + first_size, in_buf, second_size);
  }
  if (move_ptr) {
    sock->read_ndx.fetch_add(size);
  }
}

// Receive a message. Note that the buffer's read_ndx and write_ndx don't
// wrap, which makes it easier to recognize underflow.
int transport_recv(void* data, size_t len, bool block, void* transport_cookie) {
  auto* transport_info = reinterpret_cast<transport_info_t*>(transport_cookie);
  if (block) {
    while ((transport_info->in_sock->read_ndx.load() + sizeof(size_t)) >=
           transport_info->in_sock->write_ndx.load()) {
      usleep(10);
    }
  } else if ((transport_info->in_sock->read_ndx.load() + sizeof(size_t)) >=
             transport_info->in_sock->write_ndx.load()) {
    return TFTP_ERR_TIMED_OUT;
  }
  size_t block_len;
  read_from_buf(transport_info->in_sock, &block_len, sizeof(block_len), false);
  if (block_len > len) {
    return TFTP_ERR_BUFFER_TOO_SMALL;
  }
  transport_info->in_sock->read_ndx.fetch_add(sizeof(block_len));
  read_from_buf(transport_info->in_sock, data, block_len, true);
  return static_cast<int>(block_len);
}

int transport_timeout_set(uint32_t timeout_ms, void* transport_cookie) { return 0; }

/// SEND THREAD

bool run_client_test(struct test_params* tp) {
  BEGIN_HELPER;

  // Configure TFTP session
  tftp_session* session;
  size_t session_size = tftp_sizeof_session();
  void* session_buf = malloc(session_size);
  ASSERT_NONNULL(session_buf, "memory allocation failed");

  tftp_status status = tftp_init(&session, session_buf, session_size);
  ASSERT_EQ(status, TFTP_NO_ERROR, "unable to initialize a tftp session");

  // Configure file interface
  file_info_t file_info;
  file_info.filesz = tp->filesz;
  tftp_file_interface file_callbacks = {file_open_read, file_open_write, file_read, file_write,
                                        file_close};
  status = tftp_session_set_file_interface(session, &file_callbacks);
  ASSERT_EQ(status, TFTP_NO_ERROR, "could not set file interface");

  // Configure transport interface
  transport_info_t transport_info;
  transport_init(&transport_info, false);

  tftp_transport_interface transport_callbacks = {transport_send, transport_recv,
                                                  transport_timeout_set};
  status = tftp_session_set_transport_interface(session, &transport_callbacks);
  ASSERT_EQ(status, TFTP_NO_ERROR, "could not set transport interface");

  // Allocate intermediate buffers
  size_t buf_sz = tp->blksz > PATH_MAX ? tp->blksz + 2 : PATH_MAX + 2;
  char* msg_in_buf = reinterpret_cast<char*>(malloc(buf_sz));
  ASSERT_NONNULL(msg_in_buf, "memory allocation failure");
  char* msg_out_buf = reinterpret_cast<char*>(malloc(buf_sz));
  ASSERT_NONNULL(msg_out_buf, "memory allocation failure");

  char err_msg_buf[128];

  // Set our preferred transport options
  tftp_set_options(session, &tp->blksz, NULL, &tp->winsz);

  tftp_request_opts opts = {};
  opts.inbuf = msg_in_buf;
  opts.inbuf_sz = buf_sz;
  opts.outbuf = msg_out_buf;
  opts.outbuf_sz = buf_sz;
  opts.err_msg = err_msg_buf;
  opts.err_msg_sz = sizeof(err_msg_buf);

  if (tp->direction == DIR_SEND) {
    status = tftp_push_file(session, &transport_info, &file_info, "abc.txt", "xyz.txt", &opts);
    EXPECT_GE(status, 0, "failed to send file");
  } else {
    status = tftp_pull_file(session, &transport_info, &file_info, "abc.txt", "xyz.txt", &opts);
    EXPECT_GE(status, 0, "failed to receive file");
  }

  free(session);
  END_HELPER;
}

void* tftp_client_main(void* arg) {
  auto* tp = reinterpret_cast<test_params*>(arg);
  run_client_test(tp);
  pthread_exit(NULL);
}

/// RECV THREAD

bool run_server_test(struct test_params* tp) {
  BEGIN_HELPER;

  // Configure TFTP session
  tftp_session* session;
  size_t session_size = tftp_sizeof_session();
  void* session_buf = malloc(session_size);
  ASSERT_NONNULL(session_buf, "memory allocation failed");

  tftp_status status = tftp_init(&session, session_buf, session_size);
  ASSERT_EQ(status, TFTP_NO_ERROR, "unable to initiate a tftp session");

  // Configure file interface
  file_info_t file_info;
  file_info.filesz = tp->filesz;
  tftp_file_interface file_callbacks = {file_open_read, file_open_write, file_read, file_write,
                                        file_close};
  status = tftp_session_set_file_interface(session, &file_callbacks);
  ASSERT_EQ(status, TFTP_NO_ERROR, "could not set file interface");

  // Configure transport interface
  transport_info_t transport_info;
  transport_init(&transport_info, true);
  tftp_transport_interface transport_callbacks = {transport_send, transport_recv,
                                                  transport_timeout_set};
  status = tftp_session_set_transport_interface(session, &transport_callbacks);
  ASSERT_EQ(status, TFTP_NO_ERROR, "could not set transport interface");

  // Allocate intermediate buffers
  size_t buf_sz = tp->blksz > PATH_MAX ? tp->blksz + 2 : PATH_MAX + 2;
  char* msg_in_buf = reinterpret_cast<char*>(malloc(buf_sz));
  ASSERT_NONNULL(msg_in_buf, "memory allocation failure");
  char* msg_out_buf = reinterpret_cast<char*>(malloc(buf_sz));
  ASSERT_NONNULL(msg_out_buf, "memory allocation failure");

  char err_msg_buf[128];
  tftp_handler_opts opts = {.inbuf = msg_in_buf,
                            .inbuf_sz = buf_sz,
                            .outbuf = msg_out_buf,
                            .outbuf_sz = &buf_sz,
                            .err_msg = err_msg_buf,
                            .err_msg_sz = sizeof(err_msg_buf)};
  do {
    status = tftp_service_request(session, &transport_info, &file_info, &opts);
  } while (status == TFTP_NO_ERROR);
  EXPECT_EQ(status, TFTP_TRANSFER_COMPLETED, "failed to receive file");
  free(session);
  END_HELPER;
}

void* tftp_server_main(void* arg) {
  auto* tp = reinterpret_cast<test_params*>(arg);
  run_server_test(tp);
  pthread_exit(NULL);
}

bool run_one_test(struct test_params* tp) {
  BEGIN_TEST;
  int init_result = initialize_files(tp);
  ASSERT_EQ(init_result, 0, "failure to initialize state");

  clear_sockets();

  pthread_t client_thread, server_thread;
  pthread_create(&client_thread, NULL, tftp_client_main, tp);
  pthread_create(&server_thread, NULL, tftp_server_main, tp);

  pthread_join(client_thread, NULL);
  pthread_join(server_thread, NULL);

  int compare_result = compare_files(tp->filesz);
  EXPECT_EQ(compare_result, 0, "output file mismatch");
  END_TEST;
}

bool test_tftp_send_file(void) {
  struct test_params tp = {.direction = DIR_SEND, .filesz = 1000000, .winsz = 20, .blksz = 1000};
  return run_one_test(&tp);
}

bool test_tftp_send_file_wrapping_block_count(void) {
  // Wraps block count 4 times
  struct test_params tp = {.direction = DIR_SEND, .filesz = 2100000, .winsz = 9999, .blksz = 8};
  return run_one_test(&tp);
}

bool test_tftp_send_file_lg_window(void) {
  // Make sure that a window size > 255 works properly
  struct test_params tp = {.direction = DIR_SEND, .filesz = 1000000, .winsz = 1024, .blksz = 1024};
  return run_one_test(&tp);
}

bool test_tftp_receive_file(void) {
  struct test_params tp = {.direction = DIR_RECEIVE, .filesz = 1000000, .winsz = 20, .blksz = 1000};
  return run_one_test(&tp);
}

bool test_tftp_receive_file_wrapping_block_count(void) {
  // Wraps block count 4 times
  struct test_params tp = {.direction = DIR_RECEIVE, .filesz = 2100000, .winsz = 8192, .blksz = 8};
  return run_one_test(&tp);
}

bool test_tftp_receive_file_lg_window(void) {
  // Make sure that a window size > 255 works properly
  struct test_params tp = {
      .direction = DIR_RECEIVE, .filesz = 1000000, .winsz = 1024, .blksz = 1024};
  return run_one_test(&tp);
}

BEGIN_TEST_CASE(tftp_transfer_file)
RUN_TEST(test_tftp_send_file)
RUN_TEST(test_tftp_send_file_wrapping_block_count)
RUN_TEST(test_tftp_send_file_lg_window)
RUN_TEST(test_tftp_receive_file)
RUN_TEST(test_tftp_receive_file_wrapping_block_count)
RUN_TEST(test_tftp_receive_file_lg_window)
END_TEST_CASE(tftp_transfer_file)