From f4cc143dc758402fe49a997088faac63c7cba07d Mon Sep 17 00:00:00 2001 From: Ilya Kreymer Date: Thu, 7 Apr 2016 10:37:40 -0700 Subject: [PATCH] urlrewrite: generalize support for overridable handle_custom_response() callback for handling modifiers (default support top-frame) pass headers to add_custom_params, include error message on error if available headers: use add_header() to support multiple headers with same name is_ajax(): check for X-Pywb-Requested-With header to make as ajax and not pass to upstream --- urlrewrite/rewriteinputreq.py | 3 +++ urlrewrite/rewriterapp.py | 42 +++++++++++++++++++++++------------ 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/urlrewrite/rewriteinputreq.py b/urlrewrite/rewriteinputreq.py index 28879e73..10c929fc 100644 --- a/urlrewrite/rewriteinputreq.py +++ b/urlrewrite/rewriteinputreq.py @@ -42,6 +42,9 @@ class RewriteInputRequest(DirectWSGIInputRequest): if cookie_val: value = cookie_val + elif name == 'HTTP_X_PYWB_REQUESTED_WITH': + continue + elif name == 'HTTP_X_FORWARDED_PROTO': name = 'X-Forwarded-Proto' value = self.splits.scheme diff --git a/urlrewrite/rewriterapp.py b/urlrewrite/rewriterapp.py index 765c39a7..4d9836f4 100644 --- a/urlrewrite/rewriterapp.py +++ b/urlrewrite/rewriterapp.py @@ -48,16 +48,9 @@ class RewriterApp(object): rel_prefix = self.get_rel_prefix() full_prefix = host_prefix + rel_prefix - if self.framed_replay and wb_url.mod == self.frame_mod: - extra_params = self.get_top_frame_params(wb_url, kwargs) - return self.frame_insert_view.get_top_frame(wb_url, - full_prefix, - host_prefix, - request.environ, - self.frame_mod, - self.replay_mod, - coll='', - extra_params=extra_params) + resp = self.handle_custom_response(wb_url, full_prefix, host_prefix, kwargs) + if resp is not None: + return resp urlrewriter = UrlRewriter(wb_url, prefix=full_prefix, @@ -90,12 +83,19 @@ class RewriterApp(object): stream=True) if r.status_code >= 400: + error = None try: + error = r.raw.read() r.raw.close() except: pass - data = dict(url=url, args=kwargs) + if error: + error = error.decode('utf-8') + else: + error = '' + + data = dict(url=url, args=kwargs, error=error) raise HTTPError(r.status_code, exception=data) record = self.loader.parse_record_stream(r.raw) @@ -105,7 +105,7 @@ class RewriterApp(object): cdx['timestamp'] = http_date_to_timestamp(r.headers.get('Memento-Datetime')) cdx['url'] = url - self._add_custom_params(cdx, kwargs) + self._add_custom_params(cdx, r.headers, kwargs) if self.is_ajax(): head_insert_func = None @@ -129,7 +129,7 @@ class RewriterApp(object): response.status = int(status_headers.get_statuscode()) for n, v in status_headers.headers: - response.headers[n] = v + response.add_header(n, v) return gen @@ -158,6 +158,7 @@ class RewriterApp(object): def is_ajax(self): value = request.environ.get('HTTP_X_REQUESTED_WITH') + value = value or request.environ.get('HTTP_X_PYWB_REQUESTED_WITH') if value and value.lower() == 'xmlhttprequest': return True @@ -166,9 +167,22 @@ class RewriterApp(object): def get_upstream_url(self, url, wb_url, closest, kwargs): raise NotImplemented() - def _add_custom_params(self, cdx, kwargs): + def _add_custom_params(self, cdx, headers, kwargs): pass def get_top_frame_params(self, wb_url, kwargs): return None + def handle_custom_response(self, wb_url, full_prefix, host_prefix, kwargs): + if self.framed_replay and wb_url.mod == self.frame_mod: + extra_params = self.get_top_frame_params(wb_url, kwargs) + return self.frame_insert_view.get_top_frame(wb_url, + full_prefix, + host_prefix, + request.environ, + self.frame_mod, + self.replay_mod, + coll='', + extra_params=extra_params) + + return None