From 44fc7f2e5c9428c8b62b3b92d0e4fca7372674d7 Mon Sep 17 00:00:00 2001
From: Lorenz Stechauner <lorenz.stechauner@necronda.net>
Date: Mon, 23 Jan 2023 01:50:50 +0100
Subject: [PATCH] Add flags to sock_splice_chunked

---
 src/lib/proxy.c |  2 +-
 src/lib/sock.c  | 43 +++++++++++++++++++++++++++++++++++--------
 src/lib/sock.h  |  5 ++++-
 src/lib/uri.h   |  3 +--
 4 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/lib/proxy.c b/src/lib/proxy.c
index b72361f..4a73908 100644
--- a/src/lib/proxy.c
+++ b/src/lib/proxy.c
@@ -439,7 +439,7 @@ int proxy_init(proxy_ctx_t **proxy_ptr, http_req *req, http_res *res, http_statu
     if (content_len > 0) {
         ret = sock_splice(&proxy->proxy, client, buffer, sizeof(buffer), content_len);
     } else if (strcontains(transfer_encoding, "chunked")) {
-        ret = sock_splice_chunked(&proxy->proxy, client, buffer, sizeof(buffer));
+        ret = sock_splice_chunked(&proxy->proxy, client, buffer, sizeof(buffer), SOCK_CHUNKED);
     }
 
     if (ret < 0 || (content_len != 0 && ret != content_len)) {
diff --git a/src/lib/sock.c b/src/lib/sock.c
index adde43f..3bc8bbc 100644
--- a/src/lib/sock.c
+++ b/src/lib/sock.c
@@ -90,6 +90,11 @@ int sock_set_timeout(sock *s, double sec) {
 }
 
 long sock_send(sock *s, void *buf, unsigned long len, int flags) {
+    if (s->socket == 0) {
+        errno = ENOTCONN;
+        return -1;
+    }
+
     long ret;
     if (s->enc) {
         ret = SSL_write(s->ssl, buf, (int) len);
@@ -121,6 +126,11 @@ long sock_send_x(sock *s, void *buf, unsigned long len, int flags) {
 }
 
 long sock_recv(sock *s, void *buf, unsigned long len, int flags) {
+    if (s->socket == 0) {
+        errno = ENOTCONN;
+        return -1;
+    }
+
     long ret;
     if (s->enc) {
         int (*func)(SSL*, void*, int) = (flags & MSG_PEEK) ? SSL_peek : SSL_read;
@@ -140,7 +150,7 @@ long sock_recv(sock *s, void *buf, unsigned long len, int flags) {
 
 long sock_recv_x(sock *s, void *buf, unsigned long len, int flags) {
     for (long ret, rcv = 0; rcv < len; rcv += ret) {
-        if ((ret = sock_recv(s, (unsigned char *) buf + rcv, len - rcv, flags)) <= 0) {
+        if ((ret = sock_recv(s, (unsigned char *) buf + rcv, len - rcv, flags | MSG_WAITALL)) <= 0) {
             if (errno == EINTR) {
                 errno = 0, ret = 0;
                 continue;
@@ -169,20 +179,37 @@ long sock_splice(sock *dst, sock *src, void *buf, unsigned long buf_len, unsigne
     return (long) send_len;
 }
 
-long sock_splice_chunked(sock *dst, sock *src, void *buf, unsigned long buf_len) {
+long sock_splice_chunked(sock *dst, sock *src, void *buf, unsigned long buf_len, int flags) {
     long ret;
-    unsigned long send_len = 0;
-    unsigned long next_len;
+    unsigned long send_len = 0, next_len;
 
-    while (1) {
+    while (!(flags & SOCK_SINGLE_CHUNK)) {
         ret = sock_get_chunk_header(src);
         if (ret < 0) return -2;
 
         next_len = ret;
-        if (next_len <= 0) break;
 
-        ret = sock_splice(dst, src, buf, buf_len, next_len);
-        if (ret < 0) return ret;
+        if (flags & SOCK_CHUNKED) {
+            ret = sprintf(buf, "%lX\r\n", next_len);
+            if (sock_send_x(dst, buf, ret, 0) == -1)
+                return -1;
+        }
+
+        if (next_len == 0)
+            break;
+
+        if ((ret = sock_splice(dst, src, buf, buf_len, next_len)) < 0)
+            return ret;
+
+        send_len += ret;
+
+        if (flags & SOCK_CHUNKED) {
+            if (sock_send_x(dst, "\r\n", 2, 0) == -1)
+                return -1;
+        }
+
+        if (sock_recv_x(src, buf, 2, 0) == -1)
+            return -1;
     }
 
     return (long) send_len;
diff --git a/src/lib/sock.h b/src/lib/sock.h
index 886c6a2..8371b3d 100644
--- a/src/lib/sock.h
+++ b/src/lib/sock.h
@@ -13,6 +13,9 @@
 #include <sys/socket.h>
 #include <arpa/inet.h>
 
+#define SOCK_CHUNKED 1
+#define SOCK_SINGLE_CHUNK 2
+
 typedef struct {
     unsigned int enc:1;
     int socket;
@@ -48,7 +51,7 @@ long sock_recv_x(sock *s, void *buf, unsigned long len, int flags);
 
 long sock_splice(sock *dst, sock *src, void *buf, unsigned long buf_len, unsigned long len);
 
-long sock_splice_chunked(sock *dst, sock *src, void *buf, unsigned long buf_len);
+long sock_splice_chunked(sock *dst, sock *src, void *buf, unsigned long buf_len, int flags);
 
 int sock_close(sock *s);
 
diff --git a/src/lib/uri.h b/src/lib/uri.h
index 5fdf0e9..783ef84 100644
--- a/src/lib/uri.h
+++ b/src/lib/uri.h
@@ -38,8 +38,7 @@ typedef struct {
     char *filename;       // "/account/index.php"
     char *uri;            // "/account/login?username=test"
     metadata_t *meta;
-    unsigned int is_static:1;
-    unsigned int is_dir:1;
+    unsigned int is_static:1, is_dir:1;
 } http_uri;