Handle 304 (and other responses without length) correctly

1 files changed, 24 insertions, 9 deletions

diff --git a/src/proxy.cc b/src/proxy.cc
index 293e9be..e102b81 100644
--- a/src/proxy.cc
+++ b/src/proxy.cc
@@ -1004,7 +1004,7 @@ void ProxyImpl::client_event(size_t index, int fd, uint8_t events) {
   client_empty_input(index);
 }
 
-bool setup_content(Http const* http, Content* content) {
+bool setup_common_content(Http const* http, Content* content) {
   assert(content->type == CONTENT_NONE);
   auto iter = http->header_tokens("transfer-encoding");
   bool chunked = false;
@@ -1037,6 +1037,24 @@ bool setup_content(Http const* http, Content* content) {
   return true;
 }
 
+bool setup_request_content(HttpRequest const* http, Content* content) {
+  if (!setup_common_content(http, content)) return false;
+  if (content->type == CONTENT_CLOSE) {
+    content->type = CONTENT_NONE;
+  }
+  return true;
+}
+
+bool setup_response_content(HttpResponse const* http, Content* content) {
+  auto status = http->status_code();
+  if ((status >= 100 && status < 200) || status == 204 || status == 304) {
+    assert(content->type == CONTENT_NONE);
+    content->type = CONTENT_NONE;
+    return true;
+  }
+  return setup_common_content(http, content);
+}
+
 void ProxyImpl::client_empty_input(size_t index) {
   auto& client = clients_[index];
   while (true) {
@@ -1142,14 +1160,11 @@ void ProxyImpl::client_empty_input(size_t index) {
           client_error(index, 400, "Bad request");
           return;
         }
-        if (!setup_content(client.request.get(), &client.content)) {
+        if (!setup_request_content(client.request.get(), &client.content)) {
           logger_->out(Logger::INFO, "%zu: Client bad content-length", index);
           client_error(index, 400, "Bad request");
           return;
         }
-        if (client.content.type == CONTENT_CLOSE) {
-          client.content.type = CONTENT_NONE;
-        }
       }
       if (!client_request(index)) {
         client.content.type = CONTENT_NONE;
@@ -1393,7 +1408,7 @@ void ProxyImpl::client_remote_event(size_t index, int fd, uint8_t events) {
           client_remote_error(index, 502);
           return;
         }
-        if (!setup_content(response.get(), &client.remote.content)) {
+        if (!setup_response_content(response.get(), &client.remote.content)) {
           logger_->out(Logger::INFO, "%zu: Client remote bad content-length",
                        index);
           client.remote.content.type = CONTENT_CLOSE;
@@ -1628,7 +1643,7 @@ void ProxyImpl::monitor_event(size_t index, int fd, uint8_t events) {
     }
     Content content;
     content.type = CONTENT_NONE;
-    if (!setup_content(request.get(), &content)) {
+    if (!setup_request_content(request.get(), &content)) {
       monitor_error(index, 400, "Bad request");
       return;
     }
@@ -2057,7 +2072,7 @@ void ProxyImpl::MitmMonitor::local2remote(void const* data, size_t size) {
         req_buffer_.clear();
         break;
       }
-      if (!setup_content(req.get(), &request_)) {
+      if (!setup_request_content(req.get(), &request_)) {
         state_ = UNKNOWN;
         req_buffer_.clear();
         break;
@@ -2124,7 +2139,7 @@ void ProxyImpl::MitmMonitor::remote2local(void const* data, size_t size) {
         resp_buffer_.clear();
         break;
       }
-      if (!setup_content(resp.get(), &response_)) {
+      if (!setup_response_content(resp.get(), &response_)) {
         resp_buffer_.clear();
         break;
       }