1#!/usr/bin/env python 2# Copyright 2015 Google Inc. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16 17import httparchive 18import httplib 19import httpproxy 20import threading 21import unittest 22import util 23 24 25class MockCustomResponseHandler(object): 26 def __init__(self, response): 27 """ 28 Args: 29 response: An instance of ArchivedHttpResponse that is returned for each 30 request. 31 """ 32 self._response = response 33 34 def handle(self, request): 35 del request 36 return self._response 37 38 39class MockHttpArchiveFetch(object): 40 def __init__(self, response): 41 """ 42 Args: 43 response: An instance of ArchivedHttpResponse that is returned for each 44 request. 45 """ 46 self.is_record_mode = False 47 self._response = response 48 49 def __call__(self, request): 50 del request # unused 51 return self._response 52 53 54class MockHttpArchiveHandler(httpproxy.HttpArchiveHandler): 55 def handle_one_request(self): 56 httpproxy.HttpArchiveHandler.handle_one_request(self) 57 HttpProxyTest.HANDLED_REQUEST_COUNT += 1 58 59 60class MockRules(object): 61 def Find(self, unused_rule_type_name): # pylint: disable=unused-argument 62 return lambda unused_request, unused_response: None 63 64 65class HttpProxyTest(unittest.TestCase): 66 def setUp(self): 67 self.has_proxy_server_bound_port = False 68 self.has_proxy_server_started = False 69 self.allow_generate_304 = False 70 self.serve_response_by_http_archive = False 71 72 def set_up_proxy_server(self, response): 73 """ 74 Args: 75 response: An instance of ArchivedHttpResponse that is returned for each 76 request. 77 """ 78 HttpProxyTest.HANDLED_REQUEST_COUNT = 0 79 self.host = 'localhost' 80 self.port = 8889 81 custom_handlers = MockCustomResponseHandler( 82 response if not self.serve_response_by_http_archive else None) 83 rules = MockRules() 84 http_archive_fetch = MockHttpArchiveFetch( 85 response if self.serve_response_by_http_archive else None) 86 self.proxy_server = httpproxy.HttpProxyServer( 87 http_archive_fetch, custom_handlers, rules, 88 host=self.host, port=self.port, 89 allow_generate_304=self.allow_generate_304) 90 self.proxy_server.RequestHandlerClass = MockHttpArchiveHandler 91 self.has_proxy_server_bound_port = True 92 93 def tear_down_proxy_server(self): 94 if self.has_proxy_server_started: 95 self.proxy_server.shutdown() 96 if self.has_proxy_server_bound_port: 97 self.proxy_server.server_close() 98 99 def tearDown(self): 100 self.tear_down_proxy_server() 101 102 def serve_requests_forever(self): 103 self.has_proxy_server_started = True 104 self.proxy_server.serve_forever(poll_interval=0.01) 105 106 # Tests that handle_one_request does not leak threads, and does not try to 107 # re-handle connections that are finished. 108 def test_handle_one_request_closes_connection(self): 109 # By default, BaseHTTPServer.py treats all HTTP 1.1 requests as keep-alive. 110 # Intentionally use HTTP 1.0 to prevent this behavior. 111 response = httparchive.ArchivedHttpResponse( 112 version=10, status=200, reason="OK", 113 headers=[], response_data=["bat1"]) 114 self.set_up_proxy_server(response) 115 t = threading.Thread( 116 target=HttpProxyTest.serve_requests_forever, args=(self,)) 117 t.start() 118 119 initial_thread_count = threading.activeCount() 120 121 # Make a bunch of requests. 122 request_count = 10 123 for _ in range(request_count): 124 conn = httplib.HTTPConnection('localhost', 8889, timeout=10) 125 conn.request("GET", "/index.html") 126 res = conn.getresponse().read() 127 self.assertEqual(res, "bat1") 128 conn.close() 129 130 # Check to make sure that there is no leaked thread. 131 util.WaitFor(lambda: threading.activeCount() == initial_thread_count, 2) 132 133 self.assertEqual(request_count, HttpProxyTest.HANDLED_REQUEST_COUNT) 134 135 136 # Tests that keep-alive header works. 137 def test_keep_alive_header(self): 138 response = httparchive.ArchivedHttpResponse( 139 version=11, status=200, reason="OK", 140 headers=[("Connection", "keep-alive")], response_data=["bat1"]) 141 self.set_up_proxy_server(response) 142 t = threading.Thread( 143 target=HttpProxyTest.serve_requests_forever, args=(self,)) 144 t.start() 145 146 initial_thread_count = threading.activeCount() 147 148 # Make a bunch of requests. 149 request_count = 10 150 connections = [] 151 for _ in range(request_count): 152 conn = httplib.HTTPConnection('localhost', 8889, timeout=10) 153 conn.request("GET", "/index.html", headers={"Connection": "keep-alive"}) 154 res = conn.getresponse().read() 155 self.assertEqual(res, "bat1") 156 connections.append(conn) 157 158 # Repeat the same requests. 159 for conn in connections: 160 conn.request("GET", "/index.html", headers={"Connection": "keep-alive"}) 161 res = conn.getresponse().read() 162 self.assertEqual(res, "bat1") 163 164 # Check that the right number of requests have been handled. 165 self.assertEqual(2 * request_count, HttpProxyTest.HANDLED_REQUEST_COUNT) 166 167 # Check to make sure that exactly "request_count" new threads are active. 168 self.assertEqual( 169 threading.activeCount(), initial_thread_count + request_count) 170 171 for conn in connections: 172 conn.close() 173 174 util.WaitFor(lambda: threading.activeCount() == initial_thread_count, 1) 175 176 # Test that opening 400 simultaneous connections does not cause httpproxy to 177 # hit a process fd limit. The default limit is 256 fds. 178 def test_max_fd(self): 179 response = httparchive.ArchivedHttpResponse( 180 version=11, status=200, reason="OK", 181 headers=[("Connection", "keep-alive")], response_data=["bat1"]) 182 self.set_up_proxy_server(response) 183 t = threading.Thread( 184 target=HttpProxyTest.serve_requests_forever, args=(self,)) 185 t.start() 186 187 # Make a bunch of requests. 188 request_count = 400 189 connections = [] 190 for _ in range(request_count): 191 conn = httplib.HTTPConnection('localhost', 8889, timeout=10) 192 conn.request("GET", "/index.html", headers={"Connection": "keep-alive"}) 193 res = conn.getresponse().read() 194 self.assertEqual(res, "bat1") 195 connections.append(conn) 196 197 # Check that the right number of requests have been handled. 198 self.assertEqual(request_count, HttpProxyTest.HANDLED_REQUEST_COUNT) 199 200 for conn in connections: 201 conn.close() 202 203 # Tests that conditional requests return 304. 204 def test_generate_304(self): 205 REQUEST_HEADERS = [ 206 {}, 207 {'If-Modified-Since': 'whatever'}, 208 {'If-None-Match': 'whatever yet again'}] 209 RESPONSE_STATUSES = [200, 204, 304, 404] 210 for allow_generate_304 in [False, True]: 211 self.allow_generate_304 = allow_generate_304 212 for serve_response_by_http_archive in [False, True]: 213 self.serve_response_by_http_archive = serve_response_by_http_archive 214 for response_status in RESPONSE_STATUSES: 215 response = None 216 if response_status != 404: 217 response = httparchive.ArchivedHttpResponse( 218 version=11, status=response_status, reason="OK", headers=[], 219 response_data=["some content"]) 220 self.set_up_proxy_server(response) 221 t = threading.Thread( 222 target=HttpProxyTest.serve_requests_forever, args=(self,)) 223 t.start() 224 for method in ['GET', 'HEAD', 'POST']: 225 for headers in REQUEST_HEADERS: 226 connection = httplib.HTTPConnection('localhost', 8889, timeout=10) 227 connection.request(method, "/index.html", headers=headers) 228 response = connection.getresponse() 229 connection.close() 230 if (allow_generate_304 and 231 serve_response_by_http_archive and 232 method in ['GET', 'HEAD'] and 233 headers and 234 response_status == 200): 235 self.assertEqual(304, response.status) 236 self.assertEqual('', response.read()) 237 else: 238 self.assertEqual(response_status, response.status) 239 self.tear_down_proxy_server() 240 241 242if __name__ == '__main__': 243 unittest.main() 244