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