Fix duplicate header issue

Author: justmobilize

adafruit_requests.py

@@ -529,6 +529,13 @@ def _get_socket(
         self._socket_free[sock] = False
         return sock
 
+    @staticmethod
+    def _header_suppplied(header, supplied_headers):
+        for supplied_header in supplied_headers:
+            if supplied_header.lower() == header.lower():
+                return True
+        return False
+
     @staticmethod
     def _send(socket: SocketType, data: bytes):
         total_sent = 0
@@ -551,6 +558,16 @@ def _send(socket: SocketType, data: bytes):
                 raise OSError(errno.EIO)
             total_sent += sent
 
+    def _send_as_bytes(self, socket: SocketType, data: str):
+        return self._send(socket, bytes(data, "utf-8"))
+
+    def _send_header(self, socket, header, value):
+        self._send_as_bytes(socket, header)
+        self._send_as_bytes(socket, ": ")
+        self._send_as_bytes(socket, value)
+        self._send_as_bytes(socket, "\r\n")
+
+    # pylint: disable=too-many-arguments
     def _send_request(
         self,
         socket: SocketType,
@@ -561,40 +578,48 @@ def _send_request(
         data: Any,
         json: Any,
     ):
-        # pylint: disable=too-many-arguments
-        self._send(socket, bytes(method, "utf-8"))
-        self._send(socket, b" /")
-        self._send(socket, bytes(path, "utf-8"))
-        self._send(socket, b" HTTP/1.1\r\n")
-        if "Host" not in headers:
-            self._send(socket, b"Host: ")
-            self._send(socket, bytes(host, "utf-8"))
-            self._send(socket, b"\r\n")
-        if "User-Agent" not in headers:
-            self._send(socket, b"User-Agent: Adafruit CircuitPython\r\n")
-        # Iterate over keys to avoid tuple alloc
-        for k in headers:
-            self._send(socket, k.encode())
-            self._send(socket, b": ")
-            self._send(socket, headers[k].encode())
-            self._send(socket, b"\r\n")
+        # Convert data
+        content_type_header = None
+
+        # If json is sent, set content type header and convert to string
         if json is not None:
             assert data is None
+            content_type_header = "application/json"
             data = json_module.dumps(json)
-            self._send(socket, b"Content-Type: application/json\r\n")
-        if data:
-            if isinstance(data, dict):
-                self._send(
-                    socket, b"Content-Type: application/x-www-form-urlencoded\r\n"
-                )
-                _post_data = ""
-                for k in data:
-                    _post_data = "{}&{}={}".format(_post_data, k, data[k])
-                data = _post_data[1:]
-            if isinstance(data, str):
-                data = bytes(data, "utf-8")
-            self._send(socket, b"Content-Length: %d\r\n" % len(data))
+
+        # If data is sent and it's a dict, set content type header and convert to string
+        if data and isinstance(data, dict):
+            content_type_header = "application/x-www-form-urlencoded"
+            _post_data = ""
+            for k in data:
+                _post_data = "{}&{}={}".format(_post_data, k, data[k])
+            # remove first "&" from concatenation
+            data = _post_data[1:]
+
+        # Convert str data to bytes
+        if data and isinstance(data, str):
+            data = bytes(data, "utf-8")
+
+        self._send_as_bytes(socket, method)
+        self._send_as_bytes(socket, " /")
+        self._send_as_bytes(socket, path)
+        self._send_as_bytes(socket, " HTTP/1.1\r\n")
+
+        # Send headers
+        if not self._header_suppplied("Host", headers):
+            self._send_header(socket, "Host", host)
+        if not self._header_suppplied("User-Agent", headers):
+            self._send_header(socket, "User-Agent", "Adafruit CircuitPython")
+        if content_type_header and not self._header_suppplied("Content-Type", headers):
+            self._send_header(socket, "Content-Type", content_type_header)
+        if data and not self._header_suppplied("Content-Length", headers):
+            self._send_header(socket, "Content-Length", str(len(data)))
+        # Iterate over keys to avoid tuple alloc
+        for header in headers:
+            self._send_header(socket, header, headers[header])
         self._send(socket, b"\r\n")
+
+        # Send data
         if data:
             self._send(socket, bytes(data))
 

tests/chunk_test.py

@@ -63,7 +63,8 @@ def do_test_get_text(
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )
@@ -105,7 +106,8 @@ def do_test_close_flush(
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )
@@ -146,7 +148,8 @@ def do_test_get_text_extra_space(
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )

tests/concurrent_test.py

@@ -40,7 +40,8 @@ def test_second_connect_fails_memoryerror():  # pylint: disable=invalid-name
 
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
             mock.call(b"\r\n"),
         ]
@@ -82,7 +83,8 @@ def test_second_connect_fails_oserror():  # pylint: disable=invalid-name
 
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
             mock.call(b"\r\n"),
         ]

tests/header_test.py

@@ -12,7 +12,75 @@
 RESPONSE_HEADERS = b"HTTP/1.0 200 OK\r\nContent-Length: 0\r\n\r\n"
 
 
-def test_json():
+def test_host():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
+    sock = mocket.Mocket(RESPONSE_HEADERS)
+    pool.socket.return_value = sock
+    sent = []
+
+    def _send(data):
+        sent.append(data)  # pylint: disable=no-member
+        return len(data)
+
+    sock.send.side_effect = _send
+
+    requests_session = adafruit_requests.Session(pool)
+    headers = {}
+    requests_session.get("http://" + HOST + "/get", headers=headers)
+
+    sock.connect.assert_called_once_with((IP, 80))
+    sent = b"".join(sent)
+    assert b"Host: httpbin.org\r\n" in sent
+
+
+def test_host_replace():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
+    sock = mocket.Mocket(RESPONSE_HEADERS)
+    pool.socket.return_value = sock
+    sent = []
+
+    def _send(data):
+        sent.append(data)  # pylint: disable=no-member
+        return len(data)
+
+    sock.send.side_effect = _send
+
+    requests_session = adafruit_requests.Session(pool)
+    headers = {"host": IP}
+    requests_session.get("http://" + HOST + "/get", headers=headers)
+
+    sock.connect.assert_called_once_with((IP, 80))
+    sent = b"".join(sent)
+    assert b"host: 1.2.3.4\r\n" in sent
+    assert b"Host: httpbin.org\r\n" not in sent
+    assert sent.lower().count(b"host:") == 1
+
+
+def test_user_agent():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
+    sock = mocket.Mocket(RESPONSE_HEADERS)
+    pool.socket.return_value = sock
+    sent = []
+
+    def _send(data):
+        sent.append(data)  # pylint: disable=no-member
+        return len(data)
+
+    sock.send.side_effect = _send
+
+    requests_session = adafruit_requests.Session(pool)
+    headers = {}
+    requests_session.get("http://" + HOST + "/get", headers=headers)
+
+    sock.connect.assert_called_once_with((IP, 80))
+    sent = b"".join(sent)
+    assert b"User-Agent: Adafruit CircuitPython\r\n" in sent
+
+
+def test_user_agent_replace():
     pool = mocket.MocketPool()
     pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
     sock = mocket.Mocket(RESPONSE_HEADERS)
@@ -30,7 +98,55 @@ def _send(data):
     requests_session.get("http://" + HOST + "/get", headers=headers)
 
     sock.connect.assert_called_once_with((IP, 80))
-    sent = b"".join(sent).lower()
+    sent = b"".join(sent)
     assert b"user-agent: blinka/1.0.0\r\n" in sent
-    # The current implementation sends two user agents. Fix it, and uncomment below.
-    # assert sent.count(b"user-agent:") == 1
+    assert b"User-Agent: Adafruit CircuitPython\r\n" not in sent
+    assert sent.lower().count(b"user-agent:") == 1
+
+
+def test_content_type():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
+    sock = mocket.Mocket(RESPONSE_HEADERS)
+    pool.socket.return_value = sock
+    sent = []
+
+    def _send(data):
+        sent.append(data)  # pylint: disable=no-member
+        return len(data)
+
+    sock.send.side_effect = _send
+
+    requests_session = adafruit_requests.Session(pool)
+    headers = {}
+    data = {"test": True}
+    requests_session.post("http://" + HOST + "/get", data=data, headers=headers)
+
+    sock.connect.assert_called_once_with((IP, 80))
+    sent = b"".join(sent)
+    assert b"Content-Type: application/x-www-form-urlencoded\r\n" in sent
+
+
+def test_content_type_replace():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (IP, 80)),)
+    sock = mocket.Mocket(RESPONSE_HEADERS)
+    pool.socket.return_value = sock
+    sent = []
+
+    def _send(data):
+        sent.append(data)  # pylint: disable=no-member
+        return len(data)
+
+    sock.send.side_effect = _send
+
+    requests_session = adafruit_requests.Session(pool)
+    headers = {"content-type": "application/test"}
+    data = {"test": True}
+    requests_session.post("http://" + HOST + "/get", data=data, headers=headers)
+
+    sock.connect.assert_called_once_with((IP, 80))
+    sent = b"".join(sent)
+    assert b"content-type: application/test\r\n" in sent
+    assert b"Content-Type: application/x-www-form-urlencoded\r\n" not in sent
+    assert sent.lower().count(b"content-type:") == 1

tests/post_test.py

@@ -38,7 +38,8 @@ def test_method():
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"httpbin.org"),
         ]
     )
@@ -64,10 +65,18 @@ def test_form():
     pool.socket.return_value = sock
 
     requests_session = adafruit_requests.Session(pool)
-    data = {"Date": "July 25, 2019"}
+    data = {"Date": "July 25, 2019", "Time": "12:00"}
     requests_session.post("http://" + HOST + "/post", data=data)
     sock.connect.assert_called_once_with((IP, 80))
-    sock.send.assert_called_with(b"Date=July 25, 2019")
+    sock.send.assert_has_calls(
+        [
+            mock.call(b"Content-Type"),
+            mock.call(b": "),
+            mock.call(b"application/x-www-form-urlencoded"),
+            mock.call(b"\r\n"),
+        ]
+    )
+    sock.send.assert_called_with(b"Date=July 25, 2019&Time=12:00")
 
 
 def test_json():
@@ -77,7 +86,15 @@ def test_json():
     pool.socket.return_value = sock
 
     requests_session = adafruit_requests.Session(pool)
-    json_data = {"Date": "July 25, 2019"}
+    json_data = {"Date": "July 25, 2019", "Time": "12:00"}
     requests_session.post("http://" + HOST + "/post", json=json_data)
     sock.connect.assert_called_once_with((IP, 80))
-    sock.send.assert_called_with(b'{"Date": "July 25, 2019"}')
+    sock.send.assert_has_calls(
+        [
+            mock.call(b"Content-Type"),
+            mock.call(b": "),
+            mock.call(b"application/json"),
+            mock.call(b"\r\n"),
+        ]
+    )
+    sock.send.assert_called_with(b'{"Date": "July 25, 2019", "Time": "12:00"}')

tests/protocol_test.py

@@ -49,7 +49,8 @@ def test_get_https_text():
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )
@@ -80,7 +81,8 @@ def test_get_http_text():
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )
@@ -109,7 +111,8 @@ def test_get_close():
     )
     sock.send.assert_has_calls(
         [
-            mock.call(b"Host: "),
+            mock.call(b"Host"),
+            mock.call(b": "),
             mock.call(b"wifitest.adafruit.com"),
         ]
     )