Skip to content

Commit

Permalink
Rewrote http.py to be faster, mainly by using select() instead of tim…
Browse files Browse the repository at this point in the history
…eouts

Split from PR #66, part 1
  • Loading branch information
theseion committed Mar 10, 2022
1 parent 8204a00 commit 5509b0e
Showing 1 changed file with 66 additions and 50 deletions.
116 changes: 66 additions & 50 deletions ftw/http.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@
import socket
import ssl
import sys
import time
import zlib
import select

import brotli
from IPy import IP
Expand Down Expand Up @@ -263,7 +263,6 @@ def __init__(self):
'ADH-AES256-SHA:ECDHE-ECDSA-AES128-GCM-SHA256:' \
'ECDHE-RSA-AES128-GCM-SHA256:AES128-GCM-SHA256:AES128-SHA256:HIGH:'
self.CRLF = '\r\n'
self.HTTP_TIMEOUT = .3
self.RECEIVE_BYTES = 8192
self.SOCKET_TIMEOUT = 5

Expand Down Expand Up @@ -463,38 +462,77 @@ def get_response(self):
"""
Get the response from the socket
"""
self.sock.setblocking(0)
our_data = []
# Beginning time
begin = time.time()
self.sock.setblocking(False)
try:
our_data = self.read_response_from_socket()
finally:
try:
self.sock.shutdown(socket.SHUT_WR)
self.sock.close()
except OSError as err:
raise errors.TestError(
'We were unable to close the socket as expected.',
{
'msg': err,
'function': 'http.HttpUA.get_response'
})
else:
self.response_object = HttpResponse(b''.join(our_data), self)
finally:
if not b''.join(our_data):
raise errors.TestError(
'No response from server.'
+ ' Request likely timed out.',
{
'host': self.request_object.dest_addr,
'port': self.request_object.port,
'proto': self.request_object.protocol,
'msg': 'Please send the request and check'
+ ' Wireshark',
'function': 'http.HttpUA.get_response'
})

def read_response_from_socket(self):
# wait for socket to become ready
ready_sock, _, _ = select.select(
[self.sock], [], [self.sock], self.SOCKET_TIMEOUT)
if not ready_sock:
raise errors.TestError(
f'No response from server within {self.SOCKET_TIMEOUT}s',
{
'host': self.request_object.dest_addr,
'port': self.request_object.port,
'proto': self.request_object.protocol,
'msg': 'Please send the request and check Wireshark',
'function': 'http.HttpUA.get_response'
})

our_data = []
while True:
# If we have data then if we're passed the timeout break
if our_data and time.time() - begin > self.HTTP_TIMEOUT:
break
# If we're dataless wait just a bit
elif time.time() - begin > self.HTTP_TIMEOUT * 2:
break
# Recv data
try:
data = self.sock.recv(self.RECEIVE_BYTES)
if data:
our_data.append(util.ensure_binary(data))
begin = time.time()
else:
# Sleep for sometime to indicate a gap
time.sleep(self.HTTP_TIMEOUT)
except socket.error as err:
# Check if we got a timeout
if err.errno == errno.EAGAIN:
pass
if len(data) == 0:
# we're done
break
our_data.append(util.ensure_binary(data))
except BlockingIOError as e:
# If we can't handle the error here, pass it on
if e.errno == socket.EAGAIN or e.errno == socket.EWOULDBLOCK:
# we're done
break
except OSError as err:
# SSL will return SSLWantRead instead of EAGAIN
elif sys.platform == 'win32' and \
err.errno == errno.WSAEWOULDBLOCK:
if (sys.platform == 'win32' and
err.errno == errno.WSAEWOULDBLOCK):
pass
elif (self.request_object.protocol == 'https' and
err.args[0] == ssl.SSL_ERROR_WANT_READ):
continue
# If we didn't it's an error
err.args[0] == ssl.SSL_ERROR_WANT_READ):
ready_sock, _, _ = select.select(
[self.sock], [], [self.sock], .3)
if not ready_sock:
break
# It's an error
else:
raise errors.TestError(
'Failed to connect to server',
Expand All @@ -505,26 +543,4 @@ def get_response(self):
'message': err,
'function': 'http.HttpUA.get_response'
})
try:
self.sock.shutdown(socket.SHUT_WR)
self.sock.close()
except socket.error as err:
raise errors.TestError(
'We were unable to close the socket as expected.',
{
'msg': err,
'function': 'http.HttpUA.get_response'
})
else:
self.response_object = HttpResponse(b''.join(our_data), self)
finally:
if not b''.join(our_data):
raise errors.TestError(
'No response from server. Request likely timed out.',
{
'host': self.request_object.dest_addr,
'port': self.request_object.port,
'proto': self.request_object.protocol,
'msg': 'Please send the request and check Wireshark',
'function': 'http.HttpUA.get_response'
})
return our_data

0 comments on commit 5509b0e

Please sign in to comment.