270 lines
8.8 KiB
Python
270 lines
8.8 KiB
Python
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
"""
|
|
Caching HTTP Proxy for use with the Talos pageload tests
|
|
Author: Rob Arnold
|
|
|
|
This file implements a multithreaded caching http 1.1 proxy. HEAD and GET
|
|
methods are supported; POST is not yet.
|
|
|
|
Each incoming request is put onto a new thread; python does not have a thread
|
|
pool library, so a new thread is spawned for each request. I have tried to use
|
|
the python 2.4 standard library wherever possible.
|
|
|
|
Caching:
|
|
The cache is implemented in the Cache class. Items can only be added to the
|
|
cache. The only way to remove items from the cache is to blow it all away,
|
|
either by deleting the file (default: proxy_cache.db) or passing the -c or
|
|
--clear-cache flags on the command line. It is technically possible to remove
|
|
items individually from the cache, but there has been no need to do so so far.
|
|
|
|
The cache is implemented with the shelve module. The key is the combination of
|
|
host, port and request (path + params + fragment) and the values stored are the
|
|
http status code, headers and content that were received from the remote server.
|
|
|
|
Access to the cache is guarded by a semaphore which allows concurrent read
|
|
access. The semaphore is guarded by a simple mutex which prevents a deadlock
|
|
from occuring when two threads try to add an item to the cache at the same time.
|
|
|
|
Memory usage is kept to a minimum by the shelve module; only items in the cache
|
|
that are currently being served stay in memory.
|
|
|
|
Proxy:
|
|
The BaseHTTPServer.BaseHTTPRequestHandler takes care of parsing incoming
|
|
requests and managing the socket connection. See the documentation of the
|
|
BaseHTTPServer module for more information. When do_HEAD or do_GET is called,
|
|
the url that we are supposed to fetch is in self.path.
|
|
|
|
TODO:
|
|
* Implement POST requests. This requires implementing the do_POST method and
|
|
passing the post data along.
|
|
* Implement different cache policies
|
|
* Added an interface to allow administrators to probe the cache and remove
|
|
items from the database and such.
|
|
"""
|
|
|
|
__version__ = "0.1"
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import threading
|
|
import shelve
|
|
from optparse import OptionParser, OptionValueError
|
|
|
|
import SocketServer
|
|
import BaseHTTPServer
|
|
import socket
|
|
import httplib
|
|
from urlparse import urlsplit, urlunsplit
|
|
|
|
class HTTPRequestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
|
|
server_version = "TalosProxy/" + __version__
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def do_GET(self):
|
|
content = self.send_head()
|
|
if content:
|
|
try:
|
|
self.wfile.write(content)
|
|
except socket.error, e:
|
|
if options.verbose:
|
|
print "Got socket error %s" % e
|
|
#self.close_connection = 1
|
|
def do_HEAD(self):
|
|
self.send_head()
|
|
|
|
def getHeaders(self):
|
|
h = {}
|
|
for name in self.headers.keys():
|
|
h[name] = self.headers[name]
|
|
|
|
return h
|
|
|
|
def send_head(self, method="GET"):
|
|
o = urlsplit(self.path)
|
|
|
|
#sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
|
|
headers = self.getHeaders()
|
|
for k in "Proxy-Connection", "Connection":
|
|
if k in headers:
|
|
headers[k] = "Close"
|
|
if "Keep-Alive" in headers:
|
|
del headers["Keep-Alive"]
|
|
|
|
reqstring = urlunsplit(('','',o.path, o.query, o.fragment))
|
|
|
|
if options.no_cache:
|
|
cache_result = None
|
|
else:
|
|
cache_result = cache.get(o.hostname, o.port, reqstring)
|
|
|
|
if not cache_result:
|
|
if options.localonly:
|
|
self.send_error(404, "Object not in cache")
|
|
return None
|
|
else:
|
|
if options.verbose:
|
|
print "Object %s was not in the cache" % self.path
|
|
conn = httplib.HTTPConnection(o.netloc)
|
|
conn.request("GET", reqstring, headers=headers)
|
|
res = conn.getresponse()
|
|
|
|
content = res.read()
|
|
conn.close()
|
|
|
|
status, headers = res.status, res.getheaders()
|
|
|
|
if not options.no_cache:
|
|
cache.add(o.hostname, o.port, reqstring, status, headers, content)
|
|
else:
|
|
status, headers, content = cache_result
|
|
|
|
try:
|
|
self.send_response(status)
|
|
for name, value in headers:
|
|
# kill the transfer-encoding header because we don't support it when
|
|
# we send data to the client
|
|
if name not in ('transfer-encoding',):
|
|
self.send_header(name, value)
|
|
if "Content-Length" not in headers:
|
|
self.send_header("Content-Length", str(len(content)))
|
|
self.end_headers()
|
|
except socket.error, e:
|
|
if options.verbose:
|
|
print "Got socket error %s" % e
|
|
return None
|
|
return content
|
|
def log_message(self, format, *args):
|
|
if options.verbose:
|
|
BaseHTTPServer.BaseHTTPRequestHandler.log_message(self, format, *args)
|
|
|
|
class HTTPServer(SocketServer.ThreadingMixIn, BaseHTTPServer.HTTPServer):
|
|
def __init__(self, address, handler):
|
|
BaseHTTPServer.HTTPServer.__init__(self, address, handler)
|
|
|
|
class Cache(object):
|
|
"""Multithreaded cache uses the shelve module to store pages"""
|
|
# 20 concurrent threads ought to be enough for one browser
|
|
max_concurrency = 20
|
|
def __init__(self, name='', max_concurrency=20):
|
|
name = name or options.cache or "proxy_cache.db"
|
|
self.name = name
|
|
self.max_concurrency = max_concurrency
|
|
self.entries = {}
|
|
self.sem = threading.Semaphore(self.max_concurrency)
|
|
self.semlock = threading.Lock()
|
|
if options.clear_cache:
|
|
flag = 'n'
|
|
else:
|
|
flag = 'c'
|
|
self.db = shelve.DbfilenameShelf(name, flag)
|
|
|
|
def __del__(self):
|
|
if hasattr(self, 'db'):
|
|
self.db.close()
|
|
|
|
def get_key(self, host, port, resource):
|
|
return '%s:%s/%s' % (host, port, resource)
|
|
|
|
def get(self, host, port, resource):
|
|
key = self.get_key(host, port, resource)
|
|
self.semlock.acquire()
|
|
self.sem.acquire()
|
|
self.semlock.release()
|
|
try:
|
|
if not self.db.has_key(key):
|
|
return None
|
|
# returns status, headers, content
|
|
return self.db[key]
|
|
finally:
|
|
self.sem.release()
|
|
def add(self, host, port, resource, status, headers, content):
|
|
key = self.get_key(host, port, resource)
|
|
self.semlock.acquire()
|
|
for i in range(self.max_concurrency):
|
|
self.sem.acquire()
|
|
self.semlock.release()
|
|
try:
|
|
self.db[key] = (status, headers, content)
|
|
self.db.sync()
|
|
finally:
|
|
for i in range(self.max_concurrency):
|
|
self.sem.release()
|
|
|
|
class Options(object):
|
|
port = 8000
|
|
localonly = False
|
|
clear_cache = False
|
|
no_cache = False
|
|
cache = 'proxy_cache.db'
|
|
verbose = False
|
|
|
|
def _parseOptions():
|
|
def port_callback(option, opt, value, parser):
|
|
if value > 0 and value < (2 ** 16 - 1):
|
|
setattr(parser.values, option.dest, value)
|
|
else:
|
|
raise OptionValueError("Port number is out of range")
|
|
|
|
global options
|
|
parser = OptionParser(version="Talos Proxy " + __version__)
|
|
parser.add_option("-p", "--port", dest="port",
|
|
help="The port to run the proxy server on", metavar="PORT", type="int",
|
|
action="callback", callback=port_callback)
|
|
parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
|
|
help="Include additional debugging information")
|
|
parser.add_option("-l", "--localonly", action="store_true", dest="localonly",
|
|
help="Only serve pages from the local database")
|
|
parser.add_option("-c", "--clear", action="store_true", dest="clear_cache",
|
|
help="Clear the cache on startup")
|
|
parser.add_option("-n", "--no-cache", action="store_true", dest="no_cache",
|
|
help="Do not use a cache")
|
|
parser.add_option("-u", "--use-cache", dest="cache",
|
|
help="The filename of the cache to use", metavar="NAME.db")
|
|
parser.set_defaults(verbose=Options.verbose,
|
|
port=Options.port,
|
|
localonly=Options.localonly,
|
|
clear_cache=Options.clear_cache,
|
|
no_cache=Options.no_cache,
|
|
cache=Options.cache)
|
|
options, args = parser.parse_args()
|
|
|
|
"""Configures the proxy server. This should be called before run_proxy. It can be
|
|
called afterwards, but note that it is not threadsafe and some options (namely
|
|
port) will not take effect"""
|
|
def configure_proxy(**kwargs):
|
|
global options
|
|
options = Options()
|
|
for key in kwargs:
|
|
setattr(options, key, kwargs[key])
|
|
|
|
def _run():
|
|
global cache
|
|
cache = Cache()
|
|
server_address = ('', options.port)
|
|
httpd = HTTPServer(server_address, HTTPRequestHandler)
|
|
httpd.serve_forever()
|
|
|
|
"""Starts the proxy; it runs on a separate daemon thread"""
|
|
def run_proxy():
|
|
thr = threading.Thread(target=_run)
|
|
# now when we die, the daemon thread will die too
|
|
thr.setDaemon(1)
|
|
thr.start()
|
|
|
|
if __name__ == '__main__':
|
|
_parseOptions()
|
|
try:
|
|
run_proxy()
|
|
# thr.join() doesn't terminate on keyboard interrupt
|
|
while 1: time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
if options.verbose:
|
|
print "Quittin' time..."
|
|
|
|
__all__ = ['run_proxy', 'configure_proxy']
|