1 files changed, 455 insertions, 0 deletions

diff --git a/src/ot_fullscrape.c b/src/ot_fullscrape.c
new file mode 100644
index 0000000..9afad92
--- /dev/null
+++ b/src/ot_fullscrape.c
@@ -0,0 +1,455 @@
+/* This software was written by Dirk Engling <erdgeist@erdgeist.org>
+   It is considered beerware. Prost. Skol. Cheers or whatever. */
+
+#ifdef WANT_FULLSCRAPE
+
+/* System */
+#include <arpa/inet.h>
+#include <pthread.h>
+#include <stdio.h>
+#include <string.h>
+#ifdef WANT_COMPRESSION_GZIP
+#include <zlib.h>
+#endif
+#ifdef WANT_COMPRESSION_ZSTD
+#include <zstd.h>
+#endif
+
+/* Opentracker */
+#include "ot_fmt.h"
+#include "ot_hex.h"
+#include "ot_fullscrape.h"
+#include "ot_iovec.h"
+#include "ot_mutex.h"
+#include "ot_workqueue.h"
+#include "ot_urlencoded.h"
+#include "trackerlogic.h"
+
+/* Fetch full scrape info for all torrents
+   Full scrapes usually are huge and one does not want to
+   allocate more memory. So lets get them in 512k units
+*/
+#define OT_SCRAPE_CHUNK_SIZE  (1024 * 1024)
+
+/* "d8:completei%zde10:downloadedi%zde10:incompletei%zdee" */
+#define OT_SCRAPE_MAXENTRYLEN 256
+
+/* Forward declaration */
+static void fullscrape_make(int taskid, ot_tasktype mode);
+#ifdef WANT_COMPRESSION_GZIP
+static void fullscrape_make_gzip(int taskid, ot_tasktype mode);
+#endif
+#ifdef WANT_COMPRESSION_ZSTD
+static void fullscrape_make_zstd(int taskid, ot_tasktype mode);
+#endif
+
+/* Converter function from memory to human readable hex strings
+   XXX - Duplicated from ot_stats. Needs fix. */
+/* This is the entry point into this worker thread
+   It grabs tasks from mutex_tasklist and delivers results back
+*/
+static void *fullscrape_worker(void *args) {
+  (void)args;
+
+  while (g_opentracker_running) {
+    ot_tasktype tasktype = TASK_FULLSCRAPE;
+    ot_taskid   taskid   = workqueue_poptask(&tasktype);
+    if (!taskid)
+      break;
+#ifdef WANT_COMPRESSION_ZSTD
+    if (tasktype & TASK_FLAG_ZSTD)
+      fullscrape_make_zstd(taskid, tasktype);
+    else
+#endif
+#ifdef WANT_COMPRESSION_GZIP
+    if (tasktype & TASK_FLAG_GZIP)
+      fullscrape_make_gzip(taskid, tasktype);
+    else
+#endif
+      fullscrape_make(taskid, tasktype);
+    workqueue_pushchunked(taskid, NULL);
+  }
+  return NULL;
+}
+
+static pthread_t thread_id;
+void fullscrape_init(void) {
+  pthread_create( &thread_id, NULL, fullscrape_worker, NULL );
+}
+
+void fullscrape_deinit(void) {
+  workqueue_wakeall();
+  pthread_join( thread_id, NULL );
+}
+
+void fullscrape_deliver( int64_t sock, ot_tasktype tasktype ) {
+  workqueue_pushtask( sock, TASK_FULLSCRAPE | tasktype );
+}
+
+static void fullscrape_write_one( ot_tasktype mode, char **cursor, ot_torrent *torrent, ot_hash *hash ) {
+  /* Callers guarantee this much room: the chunk writer keeps its low
+     watermark that far from the end, the compressors hand over a buffer of
+     exactly that size */
+  char       *r          = *cursor;
+  const char *r_end      = r + OT_SCRAPE_MAXENTRYLEN;
+  size_t      seed_count = torrent->peer_list6->seed_count + torrent->peer_list4->seed_count;
+  size_t      peer_count = torrent->peer_list6->peer_count + torrent->peer_list4->peer_count;
+  size_t      down_count = torrent->peer_list6->down_count + torrent->peer_list4->down_count;
+
+  switch (mode & TASK_TASK_MASK) {
+  case TASK_FULLSCRAPE:
+  default:
+    /* push hash as bencoded string */
+    *r++ = '2';
+    *r++ = '0';
+    *r++ = ':';
+    memcpy(r, hash, sizeof(ot_hash));
+    r += sizeof(ot_hash);
+    /* push rest of the scrape string */
+    fmt_append(&r, r_end, "d8:completei%zde10:downloadedi%zde10:incompletei%zdee", seed_count, down_count, peer_count - seed_count);
+
+    break;
+  case TASK_FULLSCRAPE_TPB_ASCII:
+    hex_from_hash(r, *hash);
+    r += 2 * sizeof(ot_hash);
+    fmt_append(&r, r_end, ":%zd:%zd\n", seed_count, peer_count - seed_count);
+    break;
+  case TASK_FULLSCRAPE_TPB_ASCII_PLUS:
+    hex_from_hash(r, *hash);
+    r += 2 * sizeof(ot_hash);
+    fmt_append(&r, r_end, ":%zd:%zd:%zd\n", seed_count, peer_count - seed_count, down_count);
+    break;
+  case TASK_FULLSCRAPE_TPB_BINARY:
+    memcpy(r, *hash, sizeof(ot_hash));
+    r                    += sizeof(ot_hash);
+    *(uint32_t *)(r + 0)  = htonl((uint32_t)seed_count);
+    *(uint32_t *)(r + 4)  = htonl((uint32_t)(peer_count - seed_count));
+    r                    += 8;
+    break;
+  case TASK_FULLSCRAPE_TPB_URLENCODED:
+    r += urlencoded_format(r, *hash, sizeof(ot_hash));
+    fmt_append(&r, r_end, ":%zd:%zd\n", seed_count, peer_count - seed_count);
+    break;
+  case TASK_FULLSCRAPE_TRACKERSTATE:
+    hex_from_hash(r, *hash);
+    r += 2 * sizeof(ot_hash);
+    fmt_append(&r, r_end, ":%zd:%zd\n", torrent->peer_list6->base, down_count);
+    break;
+  }
+  *cursor = r;
+}
+
+static void fullscrape_make(int taskid, ot_tasktype mode) {
+  size_t bucket;
+  char  *r, *re;
+  struct iovec iovector = {NULL, 0};
+
+  /* Setup return vector... */
+  r = iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+  if (!r)
+    return;
+
+  /* re points to low watermark */
+  re = r + OT_SCRAPE_CHUNK_SIZE - OT_SCRAPE_MAXENTRYLEN;
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE)
+    fmt_append(&r, (char *)iovector.iov_base + OT_SCRAPE_CHUNK_SIZE, "d5:filesd");
+
+  /* For each bucket... */
+  for (bucket = 0; bucket < OT_BUCKET_COUNT; ++bucket) {
+    /* Get exclusive access to that bucket */
+    ot_vector  *torrents_list = mutex_bucket_lock(bucket);
+    ot_torrent *torrents      = (ot_torrent *)(torrents_list->data);
+    size_t      i;
+
+    /* For each torrent in this bucket.. */
+    for (i = 0; i < torrents_list->size; ++i) {
+      fullscrape_write_one(mode, &r, torrents + i, &torrents[i].hash);
+
+      if (r > re) {
+        iovector.iov_len = r - (char *)iovector.iov_base;
+
+        if (workqueue_pushchunked(taskid, &iovector)) {
+          free(iovector.iov_base);
+          mutex_bucket_unlock(bucket, 0);
+          return;
+        }
+        /* Allocate a fresh output buffer */
+        r = iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+        if (!r) {
+          mutex_bucket_unlock(bucket, 0);
+          return;
+        }
+
+        /* re points to low watermark */
+        re = r + OT_SCRAPE_CHUNK_SIZE - OT_SCRAPE_MAXENTRYLEN;
+      }
+    }
+
+    /* All torrents done: release lock on current bucket */
+    mutex_bucket_unlock(bucket, 0);
+
+    /* Parent thread died? */
+    if (!g_opentracker_running)
+      return;
+  }
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE)
+    fmt_append(&r, (char *)iovector.iov_base + OT_SCRAPE_CHUNK_SIZE, "ee");
+
+  /* Send rest of data */
+  iovector.iov_len = r - (char *)iovector.iov_base;
+  if (workqueue_pushchunked(taskid, &iovector))
+    free(iovector.iov_base);
+}
+
+#ifdef WANT_COMPRESSION_GZIP
+
+static void fullscrape_make_gzip(int taskid, ot_tasktype mode) {
+  size_t       bucket;
+  char        *r;
+  struct iovec iovector = {NULL, 0};
+  int          zres;
+  z_stream     strm;
+  /* Setup return vector... */
+  iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+  if (!iovector.iov_base)
+    return;
+
+  memset(&strm, 0, sizeof(strm));
+  strm.next_out  = (uint8_t *)iovector.iov_base;
+  strm.avail_out = OT_SCRAPE_CHUNK_SIZE;
+  if (deflateInit2(&strm, 7, Z_DEFLATED, 31, 9, Z_DEFAULT_STRATEGY) != Z_OK)
+    fprintf(stderr, "not ok.\n");
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE) {
+    strm.next_in  = (uint8_t *)"d5:filesd";
+    strm.avail_in = strlen("d5:filesd");
+    zres          = deflate(&strm, Z_NO_FLUSH);
+  }
+
+  /* For each bucket... */
+  for (bucket = 0; bucket < OT_BUCKET_COUNT; ++bucket) {
+    /* Get exclusive access to that bucket */
+    ot_vector  *torrents_list = mutex_bucket_lock(bucket);
+    ot_torrent *torrents      = (ot_torrent *)(torrents_list->data);
+    size_t      i;
+
+    /* For each torrent in this bucket.. */
+    for (i = 0; i < torrents_list->size; ++i) {
+      char compress_buffer[OT_SCRAPE_MAXENTRYLEN];
+      r             = compress_buffer;
+      fullscrape_write_one(mode, &r, torrents + i, &torrents[i].hash);
+      strm.next_in  = (uint8_t *)compress_buffer;
+      strm.avail_in = r - compress_buffer;
+      zres          = deflate(&strm, Z_NO_FLUSH);
+      if ((zres < Z_OK) && (zres != Z_BUF_ERROR))
+        fprintf(stderr, "deflate() failed while in fullscrape_make().\n");
+
+      /* Check if there still is enough buffer left */
+      while (!strm.avail_out) {
+        iovector.iov_len = (char *)strm.next_out - (char *)iovector.iov_base;
+
+        if (workqueue_pushchunked(taskid, &iovector)) {
+          free(iovector.iov_base);
+          mutex_bucket_unlock(bucket, 0);
+          return;
+        }
+        /* Allocate a fresh output buffer */
+        iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+        if (!iovector.iov_base) {
+          fprintf(stderr, "Out of memory trying to claim ouput buffer\n");
+          deflateEnd(&strm);
+          mutex_bucket_unlock(bucket, 0);
+          return;
+        }
+        strm.next_out  = (uint8_t *)iovector.iov_base;
+        strm.avail_out = OT_SCRAPE_CHUNK_SIZE;
+        zres           = deflate(&strm, Z_NO_FLUSH);
+        if ((zres < Z_OK) && (zres != Z_BUF_ERROR))
+          fprintf(stderr, "deflate() failed while in fullscrape_make().\n");
+      }
+    }
+
+    /* All torrents done: release lock on current bucket */
+    mutex_bucket_unlock(bucket, 0);
+
+    /* Parent thread died? */
+    if (!g_opentracker_running) {
+      deflateEnd(&strm);
+      return;
+    }
+  }
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE) {
+    strm.next_in  = (uint8_t *)"ee";
+    strm.avail_in = strlen("ee");
+  }
+
+  if (deflate(&strm, Z_FINISH) < Z_OK)
+    fprintf(stderr, "deflate() failed while in fullscrape_make()'s endgame.\n");
+
+  iovector.iov_len = (char *)strm.next_out - (char *)iovector.iov_base;
+  if (workqueue_pushchunked(taskid, &iovector)) {
+    free(iovector.iov_base);
+    deflateEnd(&strm);
+    return;
+  }
+
+  /* Check if there's a last batch of data in the zlib buffer */
+  if (!strm.avail_out) {
+      /* Allocate a fresh output buffer */
+      iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+
+      if (!iovector.iov_base) {
+        fprintf(stderr, "Problem with iovec_next_chunk\n");
+        deflateEnd(&strm);
+        return;
+      }
+      strm.next_out  = iovector.iov_base;
+      strm.avail_out = OT_SCRAPE_CHUNK_SIZE;
+      if (deflate(&strm, Z_FINISH) < Z_OK)
+        fprintf(stderr, "deflate() failed while in fullscrape_make()'s endgame.\n");
+
+      /* Only pass the new buffer if there actually was some data left in the buffer */
+      iovector.iov_len = (char *)strm.next_out - (char *)iovector.iov_base;
+      if (!iovector.iov_len || workqueue_pushchunked(taskid, &iovector))
+        free(iovector.iov_base);
+  }
+
+  deflateEnd(&strm);
+}
+/* WANT_COMPRESSION_GZIP */
+#endif
+
+#ifdef WANT_COMPRESSION_ZSTD
+
+static void fullscrape_make_zstd(int taskid, ot_tasktype mode) {
+  int            bucket;
+  char          *r;
+  struct iovec   iovector = {NULL, 0};
+  ZSTD_CCtx     *zstream = ZSTD_createCCtx();
+  ZSTD_inBuffer  inbuf;
+  ZSTD_outBuffer outbuf;
+  size_t         more_bytes;
+
+  if (!zstream)
+    return;
+
+  /* Setup return vector... */
+  iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+  if (!iovector.iov_base) {
+    ZSTD_freeCCtx(zstream);
+    return;
+  }
+
+  /* Working with a compression level 6 is half as fast as level 3, but
+     seems to be the last reasonable bump that's worth extra cpu */
+  ZSTD_CCtx_setParameter(zstream, ZSTD_c_compressionLevel, 6);
+
+  outbuf.dst  = iovector.iov_base;
+  outbuf.size = OT_SCRAPE_CHUNK_SIZE;
+  outbuf.pos  = 0;
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE) {
+    inbuf.src  = (const void *)"d5:filesd";
+    inbuf.size = strlen("d5:filesd");
+    inbuf.pos  = 0;
+    ZSTD_compressStream2(zstream, &outbuf, &inbuf, ZSTD_e_continue);
+  }
+
+  /* For each bucket... */
+  for (bucket = 0; bucket < OT_BUCKET_COUNT; ++bucket) {
+    /* Get exclusive access to that bucket */
+    ot_vector  *torrents_list = mutex_bucket_lock(bucket);
+    ot_torrent *torrents      = (ot_torrent *)(torrents_list->data);
+    size_t      i;
+
+    /* For each torrent in this bucket.. */
+    for (i = 0; i < torrents_list->size; ++i) {
+      char compress_buffer[OT_SCRAPE_MAXENTRYLEN];
+      r             = compress_buffer;
+      fullscrape_write_one(mode, &r, torrents + i, &torrents[i].hash);
+      inbuf.src     = compress_buffer;
+      inbuf.size    = r - compress_buffer;
+      inbuf.pos     = 0;
+      ZSTD_compressStream2(zstream, &outbuf, &inbuf, ZSTD_e_continue);
+
+      /* Check if there still is enough buffer left */
+      while (outbuf.pos + OT_SCRAPE_MAXENTRYLEN > outbuf.size) {
+        iovector.iov_len = outbuf.size;
+
+        if (workqueue_pushchunked(taskid, &iovector)) {
+          free(iovector.iov_base);
+          ZSTD_freeCCtx(zstream);
+          return mutex_bucket_unlock(bucket, 0);
+        }
+        /* Allocate a fresh output buffer */
+        iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+        if (!iovector.iov_base) {
+          fprintf(stderr, "Out of memory trying to claim ouput buffer\n");
+          ZSTD_freeCCtx(zstream);
+          return mutex_bucket_unlock(bucket, 0);
+        }
+
+        outbuf.dst  = iovector.iov_base;
+        outbuf.size = OT_SCRAPE_CHUNK_SIZE;
+        outbuf.pos  = 0;
+
+        ZSTD_compressStream2(zstream, &outbuf, &inbuf, ZSTD_e_continue);
+      }
+    }
+
+    /* All torrents done: release lock on current bucket */
+    mutex_bucket_unlock(bucket, 0);
+
+    /* Parent thread died? */
+    if (!g_opentracker_running)
+      return;
+  }
+
+  if ((mode & TASK_TASK_MASK) == TASK_FULLSCRAPE) {
+    inbuf.src  = (const void *)"ee";
+    inbuf.size = strlen("ee");
+    inbuf.pos  = 0;
+  }
+
+  more_bytes = ZSTD_compressStream2(zstream, &outbuf, &inbuf, ZSTD_e_end);
+
+  iovector.iov_len = outbuf.pos;
+  if (workqueue_pushchunked(taskid, &iovector)) {
+    free(iovector.iov_base);
+    ZSTD_freeCCtx(zstream);
+    return;
+  }
+
+  /* Check if there's a last batch of data in the zlib buffer */
+  if (more_bytes) {
+      /* Allocate a fresh output buffer */
+      iovector.iov_base = malloc(OT_SCRAPE_CHUNK_SIZE);
+
+      if (!iovector.iov_base) {
+        fprintf(stderr, "Problem with iovec_next_chunk\n");
+        ZSTD_freeCCtx(zstream);
+        return;
+      }
+
+      outbuf.dst  = iovector.iov_base;
+      outbuf.size = OT_SCRAPE_CHUNK_SIZE;
+      outbuf.pos  = 0;
+
+      ZSTD_compressStream2(zstream, &outbuf, &inbuf, ZSTD_e_end);
+
+      /* Only pass the new buffer if there actually was some data left in the buffer */
+      iovector.iov_len = outbuf.pos;
+      if (!iovector.iov_len || workqueue_pushchunked(taskid, &iovector))
+        free(iovector.iov_base);
+  }
+
+  ZSTD_freeCCtx(zstream);
+}
+/* WANT_COMPRESSION_ZSTD */
+#endif
+
+/* WANT_FULLSCRAPE */
+#endif