Coverage for src/icap_request_handler.py : 87%

# vim:ts=3:sw=3:expandtab 

import tornado.gen 

from icap_filter import message_factory 

# NoModICAPFilter is loaded dynamically. Without this import, 

# globals() does not see it. But, pylint complains. Disable pylint error. 

# pylint: disable=unused-import 

from pe_icap_filter import PEICAPFilter 

from file_download_icap_filter import FileDownloadICAPFilter 

from nomod_icap_filter import NoModICAPFilter 

from ops_logging.logger import get_logger 

from pyicap import BaseICAPRequestHandler 

logger = get_logger('request-handler') 

class ICAPRequestHandler(BaseICAPRequestHandler): 

"""Subclasses pyicap implementation and invoke icap_filter chain 

Builds a icap_filter chain and invokes the filters in the specified order. 

If any icap_filter decides to handle the icap_message, it needs to set the 

icap_response. All subsequent icap_filters will not see the icap transaction. 

""" 

def __init__(self, icap_connection, config, tenant_config, 

close_connection=False): 

"""Initializes the icap_filter chain""" 

BaseICAPRequestHandler.__init__(self, icap_connection, close_connection) 

self.config = config 

self.tenant_config = tenant_config 

self.preview_body = b'' 

def mkinst(cls, *args, **kwargs): 

try: 

return globals()[cls](*args, **kwargs) 

except Exception: 

raise NameError("Class %s is not defined" % cls) 

# Order specific. Will be created with self.config and self.tenant_config 

rq = self.config.get('icap_server', 'reqmod_classes') 

rs = self.config.get('icap_server', 'respmod_classes') 

reqmod_classes = [x.strip() for x in rq.split(',')] 

respmod_classes = [x.strip() for x in rs.split(',')] 

self.reqmod_filters = [mkinst(klass, self.config, self.tenant_config) 

for klass in reqmod_classes] 

self.respmod_filters = [mkinst(klass, self.config, self.tenant_config) 

for klass in respmod_classes] 

@tornado.gen.coroutine 

def read_preview(self): 

"""Helper function to read icap preview data""" 

if self.has_body and self.preview: 

while True: 

data = yield self.read_chunk() 

self.preview_body += data 

if not data: 

break 

def set_enc_status(self, status): 

"""Sets the ICAP response's status line and response code""" 

entry = self._responses.get(status, '') 

if entry and entry[0]: 

st = str(status).encode('utf-8') 

self.enc_status = b'HTTP/1.1 ' + st + b' ' + entry[0] 

else: 

super(ICAPRequestHandler, self).set_enc_status(status) 

def update_icap_headers(self, icap_resp): 

for h in icap_resp.headers: 

for v in icap_resp.headers[h]: 

self.set_icap_header(h, v) 

@tornado.gen.coroutine 

def send_icap_response(self, icap_resp, http_resp): 

"""Write the icap response to squid""" 

self.set_icap_response(icap_resp.status_code) 

if http_resp.status_code: 

self.set_enc_status(http_resp.status_code) 

self.enc_req = None 

for h in http_resp.headers: 

for v in http_resp.headers[h]: 

self.set_enc_header(h, v) 

hb = bool(http_resp.body) 

yield self.send_headers(has_body=hb) 

if hb: 

if callable(http_resp.body): 

yield http_resp.body(self.write_chunk) 

else: 

yield self.write_chunk(http_resp.body) 

yield self.write_chunk('') 

# Finally flush the connection to ensure there is no data left. 

# In the event of a transfer error the Squid <-> ICAP connection 

# MUST be flushed or data will be left and interpreted as a new 

# ICAP REQ in the connection. 

# Flush is not done at this point of error to prevent a request with no 

# content-length header from sending an unlimited amount of data. 

# Closing the client <-> squid connection by sending a valid ICAP 

# response first, will close the squid <-> upstream server connection, 

# and therefore only leave a finite number of bytes to flush. 

# By this point the ICAP response will have completed (successful or 

# not), so the flush can be performed without getting stuck in an 

# infinite loop. 

yield http_resp.flush_all() 

@tornado.gen.coroutine 

def icap_options(self): 

"""Handles ICAP OPTIONS request on behalf of all icap_filters""" 

self.set_icap_response(200) 

self.set_icap_header('Methods', 'REQMOD') 

self.set_icap_header('Methods', 'RESPMOD') 

self.set_icap_header('Service', 'Safeview PyICAP Server 1.0') 

self.set_icap_header('Allow', '204') 

self.set_icap_header('Preview', '4096') 

self.set_icap_header('Transfer-Preview', '*') 

yield self.send_headers(has_body=False) 

@tornado.gen.coroutine 

def invoke_filters(self, filters, mod_type): 

"""On REQMOD or RESPMOD requests, invokes the icap_filters""" 

try: 

icap_req, http_req, http_resp, icap_resp = message_factory(self) 

for f in filters: 

if mod_type == 'reqmod': 

yield f.icap_reqmod(icap_req, http_req, http_resp, icap_resp) 

else: 

yield f.icap_respmod(icap_req, http_req, http_resp, icap_resp) 

if not icap_resp.status_code: 

continue 

logger.info({'uri': http_req.uri, 

'mod_type': mod_type, 

'status_code': icap_resp.status_code}, 

event='icap-mod-response') 

if icap_resp.status_code == 204: 

self.update_icap_headers(icap_resp) 

yield self.no_adaptation_required() 

return 

if icap_resp.status_code == 200: 

yield self.send_icap_response(icap_resp, http_resp) 

return 

logger.info({'status_code': icap_resp.status_code}, 

event='unexpected-icap-mod-response') 

logger.info({'mod_type': mod_type}, 

event='no-icap-mod-response-any-filter') 

yield self.no_adaptation_required() 

except tornado.iostream.StreamClosedError: 

# The stream close error is raised again so that icap_connection 

# can get out of the infinite loop. 

raise 

except Exception as ex: 

logger.exception({'error': ex, 

'mod_type': mod_type}, 

event='icap_request_handler_error') 

try: 

# Give 1 last try to send no_adaptation response so that squid does 

# not flag this transaction as faulty. If squid sees too many faulty 

# transactions, it will take out the icap_server for a while (a few 

# minutes). 

yield self.no_adaptation_required() 

except Exception: 

pass 

@tornado.gen.coroutine 

def icap_reqmod(self): 

"""Handles REQMOD icap message""" 

yield self.invoke_filters(self.reqmod_filters, 'reqmod') 

@tornado.gen.coroutine 

def icap_respmod(self): 

"""Handles RESPMOD icap message""" 

yield self.read_preview() 

yield self.invoke_filters(self.respmod_filters, 'respmod')