/* This software was written by Dirk Engling <erdgeist@erdgeist.org>
   It is considered beerware. Prost. Skol. Cheers or whatever.

   $id$ */

/* System */
#include <fcntl.h>
#include <poll.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

/* Opentracker */
#include "ot_batch.h"
#include "test.h"

/* The batch writes to a real socket pair, because what is worth pinning down
   is what it does when the kernel takes only part of what it was offered. */

static char *dup_of(const char *text) {
  size_t len  = strlen(text);
  char  *copy = malloc(len);

  memcpy(copy, text, len);
  return copy;
}

static void drain(int s, char *into, size_t len) {
  size_t got = 0;

  while (got < len) {
    ssize_t r = read(s, into + got, len - got);

    if (r <= 0)
      break;
    got += (size_t)r;
  }
  into[got] = 0;
}

static void case_queue_accounting(void) {
  ot_batch batch;

  memset(&batch, 0, sizeof(batch));

  test_case("a zeroed batch is empty and sends nothing");
  CHECK_INT(batch_pending(&batch), 0);
  CHECK_INT(batch_send(-1, &batch), BATCH_DONE);

  test_case("anything queued makes it pending");
  CHECK_INT(batch_add(&batch, dup_of("hello "), 6), 0);
  CHECK_INT(batch_add_copy(&batch, "world", 5), 0);
  CHECK_INT(batch_pending(&batch), 1);
  CHECK_INT(batch.count, 2);

  batch_reset(&batch);
  test_case("and a reset empties it");
  CHECK_INT(batch_pending(&batch), 0);
  CHECK_PTR(batch.entries, NULL);
}

static void case_send_in_one_go(void) {
  ot_batch batch;
  int      sv[2];
  char     got[32];

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);

  test_case("several buffers arrive as one stream, in order");
  CHECK_INT(batch_add(&batch, dup_of("opentracker"), 11), 0);
  CHECK_INT(batch_add_copy(&batch, " says ", 6), 0);
  CHECK_INT(batch_add(&batch, dup_of("hello"), 5), 0);
  CHECK_INT(batch_send(sv[0], &batch), 22);
  CHECK_INT(batch_pending(&batch), 0);

  drain(sv[1], got, 22);
  CHECK_STR(got, "opentracker says hello");

  test_case("and an empty batch reports it has nothing left");
  CHECK_INT(batch_send(sv[0], &batch), BATCH_DONE);

  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

/* Send everything the batch holds, reading the other end as we go so the
   socket never stays full, and return what arrived. */
static size_t pump(int writer, int reader, ot_batch *batch, char *into, size_t room) {
  size_t got = 0;

  while (batch_pending(batch) || got < room) {
    ssize_t r;

    if (batch_pending(batch)) {
      ssize_t res = batch_send(writer, batch);

      if (res == BATCH_FAILED)
        break;
      if (res > 0)
        continue;
    }

    r = read(reader, into + got, room - got);
    if (r <= 0)
      break;
    got += (size_t)r;
  }
  return got;
}

static void case_partial_write_resumes(void) {
  const size_t size = 200000;
  ot_batch     batch;
  int          sv[2], sndbuf = 4096;
  char        *sent, *got;
  size_t       i, arrived;

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);
  setsockopt(sv[0], SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf));
  fcntl(sv[0], F_SETFL, fcntl(sv[0], F_GETFL, 0) | O_NONBLOCK);

  sent = malloc(size);
  got  = malloc(size);
  for (i = 0; i < size; ++i)
    sent[i] = (char)(i % 251);

  test_case("a buffer the kernel takes in pieces arrives whole and in order");
  CHECK_INT(batch_add(&batch, memcpy(malloc(size), sent, size), size), 0);
  arrived = pump(sv[0], sv[1], &batch, got, size);
  CHECK_INT(arrived, size);
  CHECK_INT(batch_pending(&batch), 0);
  CHECK_MEM(got, sent, size);

  free(sent);
  free(got);
  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

/* writev refuses more than IOV_MAX entries at once, so a long batch has to be
   sent across several calls. */
static void case_more_entries_than_one_writev(void) {
  const size_t entries = 3000, each = 4;
  ot_batch     batch;
  int          sv[2], sndbuf = 4096;
  char        *sent, *got;
  size_t       i, arrived, refused = 0;

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);
  setsockopt(sv[0], SOL_SOCKET, SO_SNDBUF, &sndbuf, sizeof(sndbuf));
  fcntl(sv[0], F_SETFL, fcntl(sv[0], F_GETFL, 0) | O_NONBLOCK);

  sent = malloc(entries * each);
  got  = malloc(entries * each);

  test_case("a batch of more entries than writev accepts still arrives whole");
  for (i = 0; i < entries; ++i) {
    char piece[4];

    piece[0] = (char)(i), piece[1] = (char)(i >> 8), piece[2] = 'x', piece[3] = 'y';
    memcpy(sent + i * each, piece, each);
    if (batch_add(&batch, memcpy(malloc(each), piece, each), each))
      ++refused;
  }
  CHECK_INT(refused, 0);
  CHECK_INT(batch.count, entries);

  arrived = pump(sv[0], sv[1], &batch, got, entries * each);
  CHECK_INT(arrived, entries * each);
  CHECK_MEM(got, sent, entries * each);

  free(sent);
  free(got);
  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

static void case_reset_frees_the_unsent(void) {
  ot_batch batch;
  int      sv[2];

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);

  test_case("a batch torn down mid flight keeps no bytes outstanding");
  CHECK_INT(batch_add(&batch, dup_of("first"), 5), 0);
  CHECK_INT(batch_add(&batch, dup_of("second"), 6), 0);
  CHECK_INT(batch_send(sv[0], &batch), 11);
  CHECK_INT(batch_add(&batch, dup_of("never sent"), 10), 0);
  CHECK_INT(batch_pending(&batch), 1);

  /* Whether the unsent buffer is freed is for the sanitizer to say */
  batch_reset(&batch);
  CHECK_INT(batch_pending(&batch), 0);
  CHECK_INT(batch.count, 0);

  close(sv[0]);
  close(sv[1]);
}

static void case_closed_socket_is_fatal(void) {
  ot_batch batch;
  int      sv[2];

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);
  close(sv[1]);

  test_case("writing to a socket whose peer is gone reports a dead connection");
  CHECK_INT(batch_add(&batch, dup_of("nobody is listening"), 19), 0);
  CHECK_INT(batch_send(sv[0], &batch), BATCH_FAILED);

  batch_reset(&batch);
  close(sv[0]);
}

static void case_vector_is_taken_whole(void) {
  ot_batch      batch;
  struct iovec *vector = malloc(3 * sizeof(struct iovec));
  int           sv[2];
  char          got[32];

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);

  test_case("a whole vector is queued in one go, in order");
  vector[0].iov_base = dup_of("one "), vector[0].iov_len = 4;
  vector[1].iov_base = dup_of("two "), vector[1].iov_len = 4;
  vector[2].iov_base = dup_of("three"), vector[2].iov_len = 5;
  CHECK_INT(batch_add_vector(&batch, vector, 3), 0);
  free(vector);
  CHECK_INT(batch.count, 3);
  CHECK_INT(batch_send(sv[0], &batch), 13);

  drain(sv[1], got, 13);
  CHECK_STR(got, "one two three");

  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

/* The entry array starts small and doubles, so crossing that boundary one
   buffer at a time is worth its own case. */
static void case_growing_one_entry_at_a_time(void) {
  const size_t entries = 100;
  ot_batch     batch;
  int          sv[2];
  char         sent[400], got[401];
  size_t       i, refused = 0;

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);

  test_case("a batch grown one buffer at a time keeps them all, in order");
  for (i = 0; i < entries; ++i) {
    char piece[4];

    piece[0] = (char)('a' + i % 26), piece[1] = (char)('0' + i % 10);
    piece[2] = (char)('A' + i % 26), piece[3] = '.';
    memcpy(sent + i * 4, piece, 4);
    if (batch_add(&batch, memcpy(malloc(4), piece, 4), 4))
      ++refused;
  }
  CHECK_INT(refused, 0);
  CHECK_INT(batch.count, entries);
  CHECK_INT(batch_send(sv[0], &batch), (ssize_t)(entries * 4));

  drain(sv[1], got, entries * 4);
  CHECK_MEM(got, sent, entries * 4);

  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

/* Appending while earlier entries are still on their way is what a chunked
   transfer does. Rather than depend on when a kernel decides to block, this
   takes entries off the front the way a completed send would and checks that
   the array reuses the space instead of growing. */
static void case_append_while_sending(void) {
  const size_t per_round = 12, each = 8;
  ot_batch     batch;
  int          sv[2];
  char         expected[16], got[17];
  size_t       i, round, refused = 0, space_after_first = 0;

  memset(&batch, 0, sizeof(batch));
  CHECK_INT(socketpair(AF_UNIX, SOCK_STREAM, 0, sv), 0);

  test_case("an array whose front has gone out is reused, not grown");
  for (round = 0; round < 8; ++round) {
    for (i = 0; i < per_round; ++i) {
      /* Distinct per entry, so a move from the wrong index shows up */
      char *piece = malloc(each);

      memset(piece, (int)('a' + (round * per_round + i) % 26), each);
      piece[0] = (char)('0' + i);
      piece[1] = (char)('0' + round);
      if (batch_add(&batch, piece, each))
        ++refused;
    }

    /* All but the last entry have left and been freed */
    while (batch.count - batch.first > 1)
      free(batch.entries[batch.first++].iov_base);

    if (!round)
      space_after_first = batch.space;
  }
  CHECK_INT(refused, 0);
  CHECK_INT(batch.space, space_after_first);

  test_case("and what is still queued survives the move, in order");
  CHECK_INT(batch.count - batch.first, 1);
  CHECK_INT(batch_add(&batch, memset(malloc(each), 'Z', each), each), 0);
  CHECK_INT(batch_send(sv[0], &batch), (ssize_t)(2 * each));

  memset(expected, (int)('a' + (7 * per_round + per_round - 1) % 26), each);
  expected[0] = (char)('0' + per_round - 1);
  expected[1] = '7';
  memset(expected + each, 'Z', each);
  drain(sv[1], got, 2 * each);
  CHECK_MEM(got, expected, 2 * each);

  batch_reset(&batch);
  close(sv[0]);
  close(sv[1]);
}

void test_batch(void) {
  case_queue_accounting();
  case_send_in_one_go();
  case_vector_is_taken_whole();
  case_growing_one_entry_at_a_time();
  case_append_while_sending();
  case_partial_write_resumes();
  case_more_entries_than_one_writev();
  case_reset_frees_the_unsent();
  case_closed_socket_is_fatal();
}