Don't trust send works. Do one recv before creating Response

Author: tannewt

adafruit_requests.py

@@ -87,6 +87,8 @@ def readinto(self, buf):
 class _SendFailed(Exception):
     """Custom exception to abort sending a request."""
 
+class OutOfRetries(Exception):
+    """Raised when requests has retried to make a request unsuccessfully."""
 
 class Response:
     """The response from a request, contains all the headers/content"""
@@ -570,13 +572,27 @@ def request(
         while retry_count < 2:
             retry_count += 1
             socket = self._get_socket(host, port, proto, timeout=timeout)
+            ok = True
             try:
                 self._send_request(socket, host, method, path, headers, data, json)
-                break
             except _SendFailed:
-                self._close_socket(socket)
-                if retry_count > 1:
-                    raise
+                ok = False
+            if ok:
+                # Read the H of "HTTP/1.1" to make sure the socket is alive. send can appear to work
+                # even when the socket is closed.
+                if hasattr(socket, "recv"):
+                    result = socket.recv(1)
+                else:
+                    result = bytearray(1)
+                    socket.recv_into(result)
+                if result == b"H":
+                    # Things seem to be ok so break with socket set.
+                    break
+            self._close_socket(socket)
+            socket = None
+
+        if not socket:
+            raise OutOfRetries()
 
         resp = Response(socket, self)  # our response
         if "location" in resp.headers and 300 <= resp.status_code <= 399:

tests/legacy_test.py

@@ -115,13 +115,13 @@ def test_second_send_fails():
 def test_first_read_fails():
     mocket.getaddrinfo.return_value = ((None, None, None, None, (ip, 80)),)
     sock = mocket.Mocket(b"")
+    sock2 = mocket.Mocket(headers + encoded)
     mocket.socket.call_count = 0  # Reset call count
-    mocket.socket.side_effect = [sock]
+    mocket.socket.side_effect = [sock, sock2]
 
     adafruit_requests.set_socket(mocket, mocket.interface)
 
-    with pytest.raises(RuntimeError):
-        r = adafruit_requests.get("http://" + host + "/testwifi/index.html")
+    r = adafruit_requests.get("http://" + host + "/testwifi/index.html")
 
     sock.send.assert_has_calls(
         [mock.call(b"testwifi/index.html"),]
@@ -131,10 +131,16 @@ def test_first_read_fails():
         [mock.call(b"Host: "), mock.call(host.encode("utf-8")), mock.call(b"\r\n"),]
     )
 
+
+    sock2.send.assert_has_calls(
+        [mock.call(b"Host: "), mock.call(host.encode("utf-8")), mock.call(b"\r\n"),]
+    )
+
     sock.connect.assert_called_once_with((ip, 80))
+    sock2.connect.assert_called_once_with((ip, 80))
     # Make sure that the socket is closed after the first receive fails.
     sock.close.assert_called_once()
-    assert mocket.socket.call_count == 1
+    assert mocket.socket.call_count == 2
 
 
 def test_second_tls_connect_fails():

tests/reuse_test.py

@@ -170,3 +170,33 @@ def test_second_send_fails():
     sock.close.assert_called_once()
     assert sock2.close.call_count == 0
     assert pool.socket.call_count == 2
+
+def test_second_send_lies_recv_fails():
+    pool = mocket.MocketPool()
+    pool.getaddrinfo.return_value = ((None, None, None, None, (ip, 80)),)
+    sock = mocket.Mocket(response)
+    sock2 = mocket.Mocket(response)
+    pool.socket.side_effect = [sock, sock2]
+
+    ssl = mocket.SSLContext()
+
+    s = adafruit_requests.Session(pool, ssl)
+    r = s.get("https://" + host + path)
+
+    sock.send.assert_has_calls(
+        [mock.call(b"testwifi/index.html"),]
+    )
+
+    sock.send.assert_has_calls(
+        [mock.call(b"Host: "), mock.call(b"wifitest.adafruit.com"), mock.call(b"\r\n"),]
+    )
+    assert r.text == str(text, "utf-8")
+
+    s.get("https://" + host + path + "2")
+
+    sock.connect.assert_called_once_with((host, 443))
+    sock2.connect.assert_called_once_with((host, 443))
+    # Make sure that the socket is closed after send fails.
+    sock.close.assert_called_once()
+    assert sock2.close.call_count == 0
+    assert pool.socket.call_count == 2