Level 08

#include "../common/common.c"  
#include <tcr/threaded_cr.h>
#include "crypto_box.h"

/*
 * Design:
 *
 * One stdin-reader
 * Many decryption processes
 * Many command parsers / executors
 * Many encryptor processes
 * One stdout-writer
 */


struct buffer {
  unsigned char *data;
  size_t len;
  CIRCLEQ_ENTRY(buffer) ll;
};

#define free_buf(x) \
  do { \
    if((x)) { \
      if((x)->data) free(x->data); \
      memset((x), 0xdf, sizeof(struct buffer)); \
      free((x)); \
      (x) = NULL; \
    } \
  } while(0)

unsigned char public_key[crypto_box_PUBLICKEYBYTES];
unsigned char secret_key[crypto_box_SECRETKEYBYTES];

static struct tc_waitq decryption_waitq;
static CIRCLEQ_HEAD(, buffer) decryption_queue;
static spinlock_t decryption_sl;

static struct tc_waitq process_waitq;
static CIRCLEQ_HEAD(, buffer) process_queue;
static spinlock_t process_sl;

static struct tc_waitq encryption_waitq;
static CIRCLEQ_HEAD(, buffer) encryption_queue;
static spinlock_t encryption_sl;

static struct tc_waitq output_waitq;
static CIRCLEQ_HEAD(, buffer) output_queue;
static spinlock_t output_sl;

int devurandom_fd = -1;


ssize_t tc_read(struct tc_fd *t, void *buf, ssize_t buflen)
{
  ssize_t ret, len = buflen;
  enum tc_rv rv;

  while (1) {
    ret = read(tc_fd(t), buf, len);
    if (ret > 0) {
      buf += ret;
      len -= ret;
      if (len == 0)
        return buflen;
    } else if (ret < 0) {
      if (errno != EAGAIN)
        return buflen - len ? buflen - len : ret;
    } else /* ret == 0 */
      return buflen - len;

    rv = tc_wait_fd_prio(EPOLLIN, t);
    if (rv != RV_OK) {
      errno = EINTR;
      return -1;
    }
  }
}

ssize_t tc_write(struct tc_fd *t, void *buf, ssize_t buflen)
{
  ssize_t ret, len = buflen;
  enum tc_rv rv;

  while (1) {
    ret = write(tc_fd(t), buf, len);
    if (ret > 0) {
      buf += ret;
      len -= ret;
      if (len == 0)
        return buflen;
    } else if (ret < 0) {
      if (errno != EAGAIN)
        return buflen - len ? buflen - len : ret;
    } else /* ret == 0 */
      return buflen - len;

    rv = tc_wait_fd_prio(EPOLLIN, t);
    if (rv != RV_OK) {
      errno = EINTR;
      return -1;
    }
  }
}

ssize_t randombytes(void *buf, ssize_t len)
{
  struct tc_fd *tcfd;
  ssize_t ret;
  ssize_t max;

  if(devurandom_fd == -1) errx(EXIT_FAILURE, "devurandom_fd has not been setup");
  max = len;
  while(len) {
    ret = read(devurandom_fd, buf, len);
    if(ret == -1) {
      if(errno == EAGAIN || errno == EINTR) continue;
      err(EXIT_FAILURE, "randombytes failed");
    }

    buf += ret;
    len -= ret;
  }

  return max;
}

static void read_stdin(void *arg)
{
  struct buffer *buf;

  struct tc_fd *tcfd = (struct tc_fd *)(arg);

  ssize_t rr;
  size_t ltr;
  
  // printf("in read_stdin\n");

  while(1) {
    /* no clean ups since program will exit once this function returns */

    if(tc_read(tcfd, &ltr, sizeof(ltr)) != sizeof(ltr)) {
      fprintf(stderr, "unable to tc_read %d bytes\n", sizeof(ltr));
      break;
    }

    buf = calloc(sizeof(struct buffer), 1);
    if(! buf) {
      fprintf(stderr, "unable to calloc %d bytes\n", sizeof(struct buffer));
      break;
    }

    buf->data = malloc(ltr);
    if(! buf->data) {
      fprintf(stderr, "unable to malloc %d bytes\n", ltr);
      break;
    }

    buf->len = ltr;

    if(tc_read(tcfd, buf->data, ltr) != ltr) {
      fprintf(stderr, "Unable to tc_read %d bytes\n", ltr);
      break;
    }

    /* Insert the buffer into the queue */

    spin_lock(&decryption_sl);
    CIRCLEQ_INSERT_TAIL(&decryption_queue, buf, ll);
    spin_unlock(&decryption_sl);

    /* And wake up one of them */

    tc_waitq_wakeup_one(&decryption_waitq);
  }


}

unsigned char peer_remote_pk[crypto_box_PUBLICKEYBYTES];

static void decryption_worker(void *arg)
{
  // XXX, exit_all
  struct buffer *input_buffer;
  struct buffer *output_buffer;

  size_t tl;
  unsigned char *tmpout;

  while(1) {
    input_buffer = output_buffer = NULL;
    tmpout = NULL;

    // printf("decryption_worker: about to wait\n");

    // XXX, decryption_queue should probably be locked in this case :/

    if(tc_waitq_wait_event(&decryption_waitq, ! CIRCLEQ_EMPTY(&decryption_queue)) != RV_OK) {
      fprintf(stderr, "decryption_worker: failed to wait\n");
      break;
    }

    // printf("decryption_worker: removing from list\n");

    spin_lock(&decryption_sl);
    input_buffer = CIRCLEQ_FIRST(&decryption_queue);
    CIRCLEQ_REMOVE(&decryption_queue, input_buffer, ll);
    spin_unlock(&decryption_sl);

    // printf("decryption_worker: processing\n");

    // XXX, valgrind this :p
    // XXX, more sanity checks? :p

    // calculate how many bytes will be processed.
    tl = input_buffer->len - crypto_box_NONCEBYTES;

    tmpout = malloc(tl);
    if(! tmpout) {
      fprintf(stderr, "decryption_worker: unable to malloc %d bytes for tmpout, skipping\n", tl);
      goto failure;
    }

    output_buffer = calloc(sizeof(struct buffer), 1);
    if(! output_buffer) {
      fprintf(stderr, "decryption_worker: unable to calloc new buffer, skipping\n");
      goto failure;
    }
    output_buffer->len = tl - crypto_box_ZEROBYTES;
    output_buffer->data = malloc(output_buffer->len);
    if(! output_buffer->data) {
      fprintf(stderr, "decryption_worker: unable to malloc new data buffer of "
      "%d bytes, skipping\n", tl - crypto_box_ZEROBYTES);
      goto failure;
    }

    // printf("attempting to decrypt with length of %d bytes\n", tl);

    if(crypto_box_open(tmpout, input_buffer->data + crypto_box_NONCEBYTES, tl, input_buffer->data,
      peer_remote_pk, secret_key) != 0) {
      fprintf(stderr, "decryption_worker: unable to crypto_box_open :s, skipping\n");
      goto failure;
    }

    // printf("decryption_worker: outputting buffer\n");

    memcpy(output_buffer->data, tmpout + crypto_box_ZEROBYTES, output_buffer->len);
    
    /* Insert the buffer into the queue */

    spin_lock(&process_sl);
    CIRCLEQ_INSERT_TAIL(&process_queue, output_buffer, ll);
    spin_unlock(&process_sl);

    /* And wake up one of them */

    tc_waitq_wakeup_one(&process_waitq);

    // printf("decryption_worker: finished, signaled process worker\n");

    goto alright;
failure:
    free_buf(output_buffer);
alright:
    free_buf(input_buffer);
    if(tmpout) free(tmpout);
  }
}

static int sanity_check_name(char *s)
{
  char buf[256];

  int error;
   error = 0;

  error |= (!! strstr(s, "ssh"));
  error |= (!! strstr(s, "/."));

  memset(buf, 0, sizeof(buf));
  if((realpath(s, buf) == NULL) && errno != ENOENT) error |= 1;

  error |= (buf[0] == '.');

  return error;
}

#define set_str_buffer(tag, buf, string) \
do { \
  size_t *t; \
  buf->len = sizeof(size_t) + strlen(string); \
  buf->data = malloc(buf->len + 1); \
  if(!buf->data) { \
    fprintf(stderr, "process_worker (set_str_buffer): failure to malloc buffer->data\n"); \
    goto failure; \
  } \
  t = (size_t *)(buf->data); \
  t[0] = tag; \
  strcpy(buf->data + 4, string); \
} while(0)



static void process_worker(void *arg)
{
  struct buffer *input_buffer;
  struct buffer *output_buffer;
  size_t *tag;
  int ret;
  mode_t mode;
  unsigned char *f, *g, *h;
  int fd;
  char msg[64];
  size_t offset;
  int len;
  struct tc_fd *tcfd;


  while(1) {
    input_buffer = NULL;

    if(tc_waitq_wait_event(&process_waitq, ! CIRCLEQ_EMPTY(&process_queue)) != RV_OK) {
      fprintf(stderr, "process_worker: failed to wait\n");
      break;
    }

    spin_lock(&process_sl);
    input_buffer = CIRCLEQ_FIRST(&process_queue);
    CIRCLEQ_REMOVE(&process_queue, input_buffer, ll);
    spin_unlock(&process_sl);

    // printf("process_worker: got one!\n");

    output_buffer = calloc(sizeof(struct buffer), 1);
    if(! output_buffer) {
      fprintf(stderr, "process_worker: can't allocate output buffer :(\n");
      goto append_output;
    }

    tag = ((size_t *)(input_buffer->data))[0];

    if(input_buffer->len < 4) {
      set_str_buffer(-1, output_buffer, "input packet length specified");
      goto append_output;
    }

    // printf("process_worker: handling %02x .. input is %s\n", input_buffer->data[4],
    // input_buffer->data + 4);
    
    switch(input_buffer->data[4]) {
      case 'm':
        // make a directory
        mode = strtoul(input_buffer->data + 5, (char **) &f, 8);
        if(f == (input_buffer->data + 5)) {
          printf("invalid mode\n");
          set_str_buffer(tag, output_buffer, "invalid mode specified");
          break;
        }
        if(f[0] == '\0') {
          printf("no filename\n");
          set_str_buffer(tag, output_buffer, "no filename specified");
          break;
        }

        if(sanity_check_name(f) != 0) {
          printf("filename specification\n");
          set_str_buffer(tag, output_buffer, "invalid filename specified");
          break;
        }

        ret = mkdir(f, mode);
        set_str_buffer(tag, output_buffer, ret == -1 ? "error in creating directory" : 
          "successfully created directory");

        break;
      case 'o':
        // open specified file, with mode.
        mode = strtoul(input_buffer->data + 5, (char **) &f, 8);
        if(f == (input_buffer->data + 5)) {
          set_str_buffer(tag, output_buffer, "invalid mode specified");
          break;
        }
        if(f[0] == '\0') {
          printf("input = %s\n", (input_buffer->data + 5));
          set_str_buffer(tag, output_buffer, "no filename specified");
          break;
        }

        if(sanity_check_name(f) != 0) {
          set_str_buffer(tag, output_buffer, "invalid filename specified");
          break;
        }

        fd = open(f, O_CREAT|O_TRUNC|O_WRONLY, mode);
        if(fd == -1) {
          set_str_buffer(tag, output_buffer, "failed to open file");
          break;
        }

        sprintf(msg, "fd is %d", fd);

        set_str_buffer(tag, output_buffer, msg);

        break;

      case 'w':
        // write to file, offset, data

        fd = strtoul(input_buffer->data + 5, (char **)&f, 10);
        if(f == (input_buffer->data + 5)) {
          set_str_buffer(tag, output_buffer, "invalid fd specified");
          break;
        }

        if(f[0] != ',') {
          set_str_buffer(tag, output_buffer, "invalid protocol");
          break;
        }

        f++;  

        offset = strtoul(f, (char **) &g, 10);
        if(f == g) {
          set_str_buffer(tag, output_buffer, "no fd specified");
          break;
        }
        len = input_buffer->len - ((unsigned int)(g) - (unsigned int)(input_buffer->data));

        while(len) {
          size_t ret;

          ret = pwrite(fd, g, len, offset);
          if(ret == -1) {
            if(errno == EAGAIN || errno == EINTR) continue;
            break;

          }

          g += ret;
          len -= ret;
        }

        set_str_buffer(tag, output_buffer, len == 0 ? "successfully wrote contents to fd\x00" : 
          "failed to write contents to fd\x00");

        break;

      case 'c':
        // close file

        fd = strtoul(input_buffer->data + 5, (char **)&f, 10);
        if(f == (input_buffer->data + 5)) {
          set_str_buffer(tag, output_buffer, "no fd specified");
          break;
        }

        if(f[0] != '\0') {
          set_str_buffer(tag, output_buffer, "fd not specified");
          break;
        }

        set_str_buffer(tag, output_buffer, "fd closed");
        close(fd);

        break;
      default:
        set_str_buffer(tag, output_buffer, "malformed input");
        break;

    }

append_output:
    spin_lock(&encryption_sl);
    CIRCLEQ_INSERT_TAIL(&encryption_queue, output_buffer, ll);
    spin_unlock(&encryption_sl);
    tc_waitq_wakeup_one(&encryption_waitq);

failure:
    free_buf(input_buffer);
  }
}

static void encryption_worker(void *arg)
{
  struct buffer *input_buffer;
  struct buffer *output_buffer;
  unsigned char *tmp;

  while(1) {
    tmp = NULL;

    // dequeue, and remove
    if(tc_waitq_wait_event(&encryption_waitq, ! CIRCLEQ_EMPTY(&encryption_queue)) != RV_OK) {
      fprintf(stderr, "encryption_worker: failed to wait\n");
      break;
    }

    spin_lock(&encryption_sl);
    input_buffer = CIRCLEQ_FIRST(&encryption_queue);
    CIRCLEQ_REMOVE(&encryption_queue, input_buffer, ll);
    spin_unlock(&encryption_sl);

    output_buffer = calloc(sizeof(struct buffer), 1);
    if(output_buffer == NULL) goto failure;
    output_buffer->len = input_buffer->len + crypto_box_ZEROBYTES + crypto_box_NONCEBYTES;
    output_buffer->data = malloc(output_buffer->len);
    if(output_buffer->data == NULL) goto failure;

    randombytes(output_buffer->data, crypto_box_NONCEBYTES);

    tmp = malloc(input_buffer->len + crypto_box_ZEROBYTES);
    if(! tmp) goto failure;


    // printf("input_buffer is %08x\n", input_buffer);
    // printf("input_buffer->data = %s, input_buffer->len = %d\n", input_buffer->data, input_buffer->len);

    memset(tmp, 0, crypto_box_ZEROBYTES);
    memcpy(tmp + crypto_box_ZEROBYTES, input_buffer->data, input_buffer->len);    

    if(crypto_box(output_buffer->data + crypto_box_NONCEBYTES, tmp, crypto_box_ZEROBYTES + input_buffer->len,
      output_buffer->data, peer_remote_pk, secret_key) != 0) {
      fprintf(stderr, "crypto_box failed\n");
      goto failure;
    }

    spin_lock(&output_sl);
    CIRCLEQ_INSERT_TAIL(&output_queue, output_buffer, ll);
    spin_unlock(&output_sl);
    tc_waitq_wakeup_one(&output_waitq);
    
    goto alright;
failure:
    free_buf(output_buffer);
alright:
    free_buf(input_buffer);
    if(tmp) { free(tmp); tmp = NULL; }
  }
}

static void write_stdout(void *arg)
{
  struct tc_fd *tcfd = (struct tc_fd *)(arg);
  struct buffer *buffer;
  size_t len;

  while(1) {
    if(tc_waitq_wait_event(&output_waitq, ! CIRCLEQ_EMPTY(&output_queue)) != RV_OK) {
      fprintf(stderr, "output_worker: failed to wait\n");
      break;
    }

    spin_lock(&output_sl);
    buffer = CIRCLEQ_FIRST(&output_queue);
    CIRCLEQ_REMOVE(&output_queue, buffer, ll);
    spin_unlock(&output_sl);

    len = buffer->len;

    if(tc_write(tcfd, &len, sizeof(size_t)) != sizeof(size_t)) {
      errx(1, "failed to write");
      // xxx, return to tc_main
    }

    if(tc_write(tcfd, buffer->data, buffer->len) != buffer->len) {
      errx(1, "failed to write data");
    }

    free_buf(buffer);

  }
}

static void tc_main(void *arg)
{
  struct tc_thread *reader;
  struct tc_thread_pool decryption_pool;
  struct tc_thread_pool process_pool;
  struct tc_thread_pool encryption_pool;
  struct tc_thread *writer;
  struct tc_fd *tcfd;

  CIRCLEQ_INIT(&decryption_queue);
  spin_lock_init(&decryption_sl);
  tc_waitq_init(&decryption_waitq);

  CIRCLEQ_INIT(&process_queue);
  spin_lock_init(&process_sl);
  tc_waitq_init(&process_waitq);

  CIRCLEQ_INIT(&encryption_queue);
  spin_lock_init(&encryption_sl);
  tc_waitq_init(&encryption_waitq);

  CIRCLEQ_INIT(&output_queue);
  spin_lock_init(&output_sl);
  tc_waitq_init(&output_waitq);

  tcfd = tc_register_fd(fileno(stdin));

  if(write(fileno(stdin), public_key, crypto_box_PUBLICKEYBYTES) != crypto_box_PUBLICKEYBYTES)
    err(EXIT_FAILURE, "writing public key");

  if(tc_read(tcfd, peer_remote_pk, sizeof(peer_remote_pk)) != sizeof(peer_remote_pk)) {
    fprintf(stderr, "failed to read public key\n");
    goto failure;
  }

  tc_thread_pool_new(&decryption_pool, decryption_worker, NULL, "decryption_worker_%d", 0);
  tc_thread_pool_new(&process_pool, process_worker, NULL, "process_worker_%d", 0);
  tc_thread_pool_new(&encryption_pool, encryption_worker, NULL, "encryption_worker_%d", 0);

  if((reader = tc_thread_new(read_stdin, tcfd, "stdin_reader")) == NULL) {
    errx(1, "tc_thread_new failed to create stdin_reader");
  }

  if((writer = tc_thread_new(write_stdout, tcfd, "stdout_writer")) == NULL) {
    errx(1, "tc_thread_new failed to create stdout_writer");
  }

  tc_thread_wait(reader);
    
failure:
  tc_unregister_fd(tcfd);

  // signal exit

  // printf("waited on read_stdin\n");

}

int main(int argc, char **argv, char **envp)
{
  int fd;
  char *p;

  if(is_restarted_process(argc, argv)) {
    devurandom_fd = open("/dev/urandom", O_RDONLY);

    crypto_box_keypair(public_key, secret_key);

    if(chroot("/home/level08") == -1) {
      err(1, "Unable to chroot");
    }

    if(chdir("/") == -1) {
      err(1, "chdir(\"/\")");
    }

    drop_privileges(UID, GID);

    tc_run(tc_main, NULL, "tc_main", sysconf(_SC_NPROCESSORS_ONLN));
    return 0;    
  }

  if(daemon(0, 0) == -1) {
    err(EXIT_FAILURE, "Unable to daemon(0, 0)");
  }

  fd = serve_forever(PORT);
  set_io(fd);
  restart_process(NAME);

}