1 files changed, 41 insertions, 16 deletions

diff --git a/ot_http.c b/ot_http.c
index cd2dfc1..a87a02e 100644
--- a/ot_http.c
+++ b/ot_http.c
@@ -4,9 +4,9 @@
    $id$ */
 
 /* System */
+#define _GNU_SOURCE
 #include <arpa/inet.h>
 #include <pthread.h>
-#define _GNU_SOURCE
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -32,7 +32,7 @@
 #include "trackerlogic.h"
 
 #ifdef WANT_NO_AUTO_FREE
-#define OT_IOB_INIT(B) bzero(B, sizeof(io_batch))
+#define OT_IOB_INIT(B) (bzero(B, sizeof(io_batch)), 0)
 #else
 #define OT_IOB_INIT(B) iob_init_autofree(B, 0)
 #endif
@@ -88,7 +88,15 @@ static void http_senddata(const int64 sock, struct ot_workstruct *ws) {
     memcpy(outbuf, ws->reply + written_size, ws->reply_size - written_size);
     if (!cookie->batch) {
       cookie->batch = malloc(sizeof(io_batch));
-      OT_IOB_INIT(cookie->batch);
+      if (!cookie->batch || OT_IOB_INIT(cookie->batch) == -1) {
+        free(cookie->batch);
+        free(outbuf);
+        array_reset(&cookie->request);
+        free(cookie);
+        io_close(sock);
+        return;
+      }
+
       cookie->batches = 1;
     }
 
@@ -159,7 +167,9 @@ ssize_t http_sendiovecdata(const int64 sock, struct ot_workstruct *ws, int iovec
 
   if (iovec_entries) {
 
-    if (cookie->flag & STRUCT_HTTP_FLAG_GZIP)
+    if (cookie->flag & STRUCT_HTTP_FLAG_ZSTD)
+      encoding = "Content-Encoding: zstd\r\n";
+    else if (cookie->flag & STRUCT_HTTP_FLAG_GZIP)
       encoding = "Content-Encoding: gzip\r\n";
     else if (cookie->flag & STRUCT_HTTP_FLAG_BZIP2)
       encoding = "Content-Encoding: bzip2\r\n";
@@ -181,12 +191,12 @@ ssize_t http_sendiovecdata(const int64 sock, struct ot_workstruct *ws, int iovec
 
     if (!cookie->batch) {
       cookie->batch = malloc(sizeof(io_batch));
-      if (!cookie->batch) {
+      if (!cookie->batch || OT_IOB_INIT(cookie->batch) == -1) {
+        free(cookie->batch);
         free(header);
         iovec_free(&iovec_entries, &iovector);
         HTTPERROR_500;
       }
-      OT_IOB_INIT(cookie->batch);
       cookie->batches = 1;
     }
     current = cookie->batch + cookie->batches - 1;
@@ -199,8 +209,8 @@ ssize_t http_sendiovecdata(const int64 sock, struct ot_workstruct *ws, int iovec
         io_batch *new_batch = realloc(cookie->batch, (cookie->batches + 1) * sizeof(io_batch));
         if (new_batch) {
           cookie->batch = new_batch;
-          current       = cookie->batch + cookie->batches++;
-          OT_IOB_INIT(current);
+          if (OT_IOB_INIT(current) != -1)
+            current = cookie->batch + cookie->batches++;
         }
       }
       iob_addbuf_free(current, iovector[i].iov_base, iovector[i].iov_len);
@@ -369,19 +379,34 @@ static ssize_t http_handle_fullscrape(const int64 sock, struct ot_workstruct *ws
   }
 #endif
 
-#ifdef WANT_COMPRESSION_GZIP
+
+#if defined(WANT_COMPRESSION_GZIP) || defined(WANT_COMPRESSION_ZSTD)
   ws->request[ws->request_size - 1] = 0;
-#ifndef WANT_COMPRESSION_GZIP_ALWAYS
+#ifdef WANT_COMPRESSION_GZIP
   if (strstr(ws->request, "gzip")) {
-#endif
     cookie->flag |= STRUCT_HTTP_FLAG_GZIP;
-    format        = TASK_FLAG_GZIP;
-    stats_issue_event(EVENT_FULLSCRAPE_REQUEST_GZIP, 0, (uintptr_t)cookie->ip);
-#ifndef WANT_COMPRESSION_GZIP_ALWAYS
-  } else
+    format       |= TASK_FLAG_GZIP;
+  }
 #endif
+#ifdef WANT_COMPRESSION_ZSTD
+  if (strstr(ws->request, "zstd")) {
+    cookie->flag |= STRUCT_HTTP_FLAG_ZSTD;
+    format       |= TASK_FLAG_ZSTD;
+  }
+#endif
+
+#if defined(WANT_COMPRESSION_ZSTD) && defined(WANT_COMPRESSION_ZSTD_ALWAYS)
+  cookie->flag |= STRUCT_HTTP_FLAG_ZSTD;
+  format       |= TASK_FLAG_ZSTD;
 #endif
-    stats_issue_event(EVENT_FULLSCRAPE_REQUEST, 0, (uintptr_t)cookie->ip);
+
+#if defined(WANT_COMPRESSION_GZIP) && defined(WANT_COMPRESSION_GZIP_ALWAYS)
+  cookie->flag |= STRUCT_HTTP_FLAG_GZIP;
+  format       |= TASK_FLAG_GZIP;
+#endif
+#endif
+
+  stats_issue_event(EVENT_FULLSCRAPE_REQUEST, 0, (uintptr_t)cookie->ip);
 
 #ifdef _DEBUG_HTTPERROR
   fprintf(stderr, "%s", ws->debugbuf);