Tests: parsing of "Transfer-Encoding: chunked" responses.

1 files changed, 56 insertions, 6 deletions

diff --git a/test/unit/http.py b/test/unit/http.py
index c7e3e36d..d59c7b56 100644
--- a/test/unit/http.py
+++ b/test/unit/http.py
@@ -1,5 +1,6 @@
 import re
 import time
+import json
 import socket
 import select
 from unit.main import TestUnit
@@ -86,23 +87,24 @@ class TestHTTP(TestUnit):
 
         sock.sendall(req)
 
+        encoding = 'utf-8' if 'encoding' not in kwargs else kwargs['encoding']
+
         if TestUnit.detailed:
             print('>>>')
             try:
-                print(req.decode('utf-8', 'ignore'))
+                print(req.decode(encoding, 'ignore'))
             except UnicodeEncodeError:
                 print(req)
 
         resp = ''
 
         if 'no_recv' not in kwargs:
-            enc = 'utf-8' if 'encoding' not in kwargs else kwargs['encoding']
             read_timeout = (
                 30 if 'read_timeout' not in kwargs else kwargs['read_timeout']
             )
             resp = self.recvall(
                 sock, read_timeout=read_timeout, buff_size=read_buffer_size
-            ).decode(enc)
+            ).decode(encoding)
 
         if TestUnit.detailed:
             print('<<<')
@@ -114,6 +116,12 @@ class TestHTTP(TestUnit):
         if 'raw_resp' not in kwargs:
             resp = self._resp_to_dict(resp)
 
+            headers = resp.get('headers')
+            if headers and headers.get('Transfer-Encoding') == 'chunked':
+                resp['body'] = self._parse_chunked_body(resp['body']).decode(
+                    encoding
+                )
+
         if 'start' not in kwargs:
             sock.close()
             return resp
@@ -151,7 +159,7 @@ class TestHTTP(TestUnit):
         return data
 
     def _resp_to_dict(self, resp):
-        m = re.search('(.*?\x0d\x0a?)\x0d\x0a?(.*)', resp, re.M | re.S)
+        m = re.search(r'(.*?\x0d\x0a?)\x0d\x0a?(.*)', resp, re.M | re.S)
 
         if not m:
             return {}
@@ -162,12 +170,12 @@ class TestHTTP(TestUnit):
         headers_lines = p.findall(headers_text)
 
         status = re.search(
-            '^HTTP\/\d\.\d\s(\d+)|$', headers_lines.pop(0)
+            r'^HTTP\/\d\.\d\s(\d+)|$', headers_lines.pop(0)
         ).group(1)
 
         headers = {}
         for line in headers_lines:
-            m = re.search('(.*)\:\s(.*)', line)
+            m = re.search(r'(.*)\:\s(.*)', line)
 
             if m.group(1) not in headers:
                 headers[m.group(1)] = m.group(2)
@@ -180,6 +188,48 @@ class TestHTTP(TestUnit):
 
         return {'status': int(status), 'headers': headers, 'body': body}
 
+    def _parse_chunked_body(self, raw_body):
+        if isinstance(raw_body, str):
+            raw_body = bytes(raw_body.encode())
+
+        crlf = b'\r\n'
+        chunks = raw_body.split(crlf)
+
+        if len(chunks) < 3:
+            self.fail('Invalid chunked body')
+
+        if chunks.pop() != b'':
+            self.fail('No CRLF at the end of the body')
+
+        try:
+            last_size = int(chunks[-2], 16)
+        except:
+            self.fail('Invalid zero size chunk')
+
+        if last_size != 0 or chunks[-1] != b'':
+            self.fail('Incomplete body')
+
+        body = b''
+        while len(chunks) >= 2:
+            try:
+                size = int(chunks.pop(0), 16)
+            except:
+                self.fail('Invalid chunk size %s' % str(size))
+
+            if size == 0:
+                self.assertEqual(len(chunks), 1, 'last zero size')
+                break
+
+            temp_body = crlf.join(chunks)
+
+            body += temp_body[:size]
+
+            temp_body = temp_body[size + len(crlf) :]
+
+            chunks = temp_body.split(crlf)
+
+        return body
+
     def waitforsocket(self, port):
         ret = False