Repository: prometheus/client_python
Pull request #622 (2021-02-08 18:11:28)

Implement redirect handler option.
This fixes https://github.com/prometheus/client_python/issues/131 and a significant amount of content is discussed in that issue.

This PR adds a new redirect handler utility class that deliberately forwards an unmodified request to the `Location` header provided by a redirect, while preserving the method, headers and data of the original request. As discussed in comments of the PR, HTTP RFC 2616 requires explicit user approval for redirects in certain situations that may alter the meaning or effect of the request. Given this, the default `HTTPRedirectHandler` in `urllib` does not honor redirects for `PUT` and does not preserve all headers and data.

In the case of Prometheus, many users will have some type of in-house "private" deployment of a Prometheus pushgateway server, and e.g. in a deployment infrastructure like Kubernetes where the ultimate IP-based host address could change often, there may be a need for a purely cosmetic redirect that fronts this host address with an easier static URL. In that case, the user would like to "opt in" to trusted pass-through redirects for all method types (including PUT).

This PR provides `passthrough_redirect_handler` to support this specific redirect option. Use of this special redirect handler would constitute the user's explicit approval for RFC2616.

Additionally, this PR adds a new "redirect" HTTPD server to the unit tests for the pushgateway, in order to simulate a real redirect during testing and confirm it is handled as expected. I've locally validated all testing and flake8 instructions - and will work to resolve any requested changes or build issues.


Patch:
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
index 4a79930..0720ce4 100644
--- a/prometheus_client/exposition.py
+++ b/prometheus_client/exposition.py
@@ -17,22 +17,86 @@ try:
 
     from BaseHTTPServer import BaseHTTPRequestHandler
     from SocketServer import ThreadingMixIn
-    from urllib2 import build_opener, HTTPHandler, Request
+    from urllib2 import (
+        build_opener, HTTPError, HTTPHandler, HTTPRedirectHandler, Request,
+    )
     from urlparse import parse_qs, urlparse
 except ImportError:
     # Python 3
     from http.server import BaseHTTPRequestHandler
     from socketserver import ThreadingMixIn
+    from urllib.error import HTTPError
     from urllib.parse import parse_qs, quote_plus, urlparse
-    from urllib.request import build_opener, HTTPHandler, Request
+    from urllib.request import (
+        build_opener, HTTPHandler, HTTPRedirectHandler, Request,
+    )
 
 CONTENT_TYPE_LATEST = str('text/plain; version=0.0.4; charset=utf-8')
 """Content type of the latest text format"""
-
+PYTHON27_OR_OLDER = sys.version_info < (3, )
 PYTHON26_OR_OLDER = sys.version_info < (2, 7)
 PYTHON376_OR_NEWER = sys.version_info > (3, 7, 5)
 
 
+class _PrometheusRedirectHandler(HTTPRedirectHandler):
+    """
+    Allow additional methods (e.g. PUT) and data forwarding in redirects.
+
+    Use of this class constitute a user's explicit agreement to the
+    redirect responses the Prometheus client will receive when using it.
+    You should only use this class if you control or otherwise trust the
+    redirect behavior involved and are certain it is safe to full transfer
+    the original request (method and data) to the redirected URL. For
+    example, if you know there is a cosmetic URL redirect in front of a
+    local deployment of a Prometheus server, and all redirects are safe,
+    this is the class to use to handle redirects in that case.
+
+    The standard HTTPRedirectHandler does not forward request data nor
+    does it allow redirected PUT requests (which Prometheus uses for some
+    operations, for example `push_to_gateway`) because these cannot
+    generically guarantee no violations of HTTP RFC 2616 requirements for
+    the user to explicitly confirm redirects that could have unexpected
+    side effects (such as rendering a PUT request non-idempotent or
+    creating multiple resources not named in the original request).
+    """
+
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        """
+        Apply redirect logic to a request.
+
+        See parent HTTPRedirectHandler.redirect_request for parameter info.
+
+        If the redirect is disallowed, this raises the corresponding HTTP error.
+        If the redirect can't be determined, return None to allow other handlers
+        to try. If the redirect is allowed, return the new request.
+
+        This method specialized for the case when (a) the user knows that the
+        redirect will not cause unacceptable side effects for any request method,
+        and (b) the user knows that any request data should be passed through to
+        the redirect. If either condition is not met, this should not be used.
+        """
+        # note that requests being provided by a handler will use get_method to
+        # indicate the method, by monkeypatching this, instead of setting the
+        # Request object's method attribute.
+        m = getattr(req, "method", req.get_method())
+        if not (code in (301, 302, 303, 307) and m in ("GET", "HEAD")
+                or code in (301, 302, 303) and m in ("POST", "PUT")):
+            raise HTTPError(req.full_url, code, msg, headers, fp)
+        new_request = Request(
+            newurl.replace(' ', '%20'),  # space escaping in new url if needed.
+            headers=req.headers,
+            origin_req_host=req.origin_req_host,
+            unverifiable=True,
+            data=req.data,
+        )
+        if PYTHON27_OR_OLDER:
+            # the `method` attribute did not exist for Request in Python 2.7.
+            new_request.get_method = lambda: m
+        else:
+            new_request.method = m
+        return new_request
+
+
 def _bake_output(registry, accept_header, params):
     """Bake output for metrics output."""
     encoder, content_type = choose_encoder(accept_header)
@@ -141,7 +205,7 @@ def generate_latest(registry=REGISTRY):
             raise
 
         for suffix, lines in sorted(om_samples.items()):
-            output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix, 
+            output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix,
                                                        metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
             output.append('# TYPE {0}{1} gauge\n'.format(metric.name, suffix))
             output.extend(lines)
@@ -205,17 +269,14 @@ def write_to_textfile(path, registry):
     os.rename(tmppath, path)
 
 
-def default_handler(url, method, timeout, headers, data):
-    """Default handler that implements HTTP/HTTPS connections.
-
-    Used by the push_to_gateway functions. Can be re-used by other handlers."""
+def _make_handler(url, method, timeout, headers, data, base_handler):
 
     def handle():
         request = Request(url, data=data)
         request.get_method = lambda: method
         for k, v in headers:
             request.add_header(k, v)
-        resp = build_opener(HTTPHandler).open(request, timeout=timeout)
+        resp = build_opener(base_handler).open(request, timeout=timeout)
         if resp.code >= 400:
             raise IOError("error talking to pushgateway: {0} {1}".format(
                 resp.code, resp.msg))
@@ -223,6 +284,28 @@ def default_handler(url, method, timeout, headers, data):
     return handle
 
 
+def default_handler(url, method, timeout, headers, data):
+    """Default handler that implements HTTP/HTTPS connections.
+
+    Used by the push_to_gateway functions. Can be re-used by other handlers."""
+
+    return _make_handler(url, method, timeout, headers, data, HTTPHandler)
+
+
+def passthrough_redirect_handler(url, method, timeout, headers, data):
+    """
+    Handler that automatically trusts redirect responses for all HTTP methods.
+
+    Augments standard HTTPRedirectHandler capability by permitting PUT requests,
+    preserving the method upon redirect, and passing through all headers and
+    data from the original request. Only use this handler if you control or
+    trust the source of redirect responses you encounter when making requests
+    via the Prometheus client. This handler will simply repeat the identical
+    request, including same method and data, to the new redirect URL."""
+
+    return _make_handler(url, method, timeout, headers, data, _PrometheusRedirectHandler)
+
+
 def basic_auth_handler(url, method, timeout, headers, data, username=None, password=None):
     """Handler that implements HTTP/HTTPS connections with Basic Auth.
 


Test patch:
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
index 60ad1be..a6c79ef 100644
--- a/tests/test_exposition.py
+++ b/tests/test_exposition.py
@@ -14,6 +14,7 @@ from prometheus_client import (
 from prometheus_client.core import GaugeHistogramMetricFamily, Timestamp
 from prometheus_client.exposition import (
     basic_auth_handler, default_handler, MetricsHandler,
+    passthrough_redirect_handler,
 )
 
 if sys.version_info < (2, 7):
@@ -208,6 +209,8 @@ ts{foo="f"} 0.0 123000
 
 class TestPushGateway(unittest.TestCase):
     def setUp(self):
+        redirect_flag = 'testFlag'
+        self.redirect_flag = redirect_flag  # preserve a copy for downstream test assertions
         self.registry = CollectorRegistry()
         self.counter = Gauge('g', 'help', registry=self.registry)
         self.requests = requests = []
@@ -216,6 +219,11 @@ class TestPushGateway(unittest.TestCase):
             def do_PUT(self):
                 if 'with_basic_auth' in self.requestline and self.headers['authorization'] != 'Basic Zm9vOmJhcg==':
                     self.send_response(401)
+                elif 'redirect' in self.requestline and redirect_flag not in self.requestline:
+                    # checks for an initial test request with 'redirect' but without the redirect_flag,
+                    # and simulates a redirect to a url with the redirect_flag (which will produce a 201)
+                    self.send_response(301)
+                    self.send_header('Location', getattr(self, 'redirect_address', None))
                 else:
                     self.send_response(201)
                 length = int(self.headers['content-length'])
@@ -225,6 +233,22 @@ class TestPushGateway(unittest.TestCase):
             do_POST = do_PUT
             do_DELETE = do_PUT
 
+        # set up a separate server to serve a fake redirected request.
+        # the redirected URL will have `redirect_flag` added to it,
+        # which will cause the request handler to return 201.
+        httpd_redirect = HTTPServer(('localhost', 0), TestHandler)
+        self.redirect_address = TestHandler.redirect_address = \
+            'http://localhost:{0}/{1}'.format(httpd_redirect.server_address[1], redirect_flag)
+
+        class TestRedirectServer(threading.Thread):
+            def run(self):
+                httpd_redirect.handle_request()
+
+        self.redirect_server = TestRedirectServer()
+        self.redirect_server.daemon = True
+        self.redirect_server.start()
+
+        # set up the normal server to serve the example requests across test cases.
         httpd = HTTPServer(('localhost', 0), TestHandler)
         self.address = 'http://localhost:{0}'.format(httpd.server_address[1])
 
@@ -236,6 +260,7 @@ class TestPushGateway(unittest.TestCase):
         self.server.daemon = True
         self.server.start()
 
+
     def test_push(self):
         push_to_gateway(self.address, "my_job", self.registry)
         self.assertEqual(self.requests[0][0].command, 'PUT')
@@ -330,6 +355,27 @@ class TestPushGateway(unittest.TestCase):
         self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
         self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
 
+    def test_push_with_redirect_handler(self):
+        def my_redirect_handler(url, method, timeout, headers, data):
+            return passthrough_redirect_handler(url, method, timeout, headers, data)
+
+        push_to_gateway(self.address, "my_job_with_redirect", self.registry, handler=my_redirect_handler)
+        self.assertEqual(self.requests[0][0].command, 'PUT')
+        self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job_with_redirect')
+        self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
+        self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
+
+        # ensure the redirect preserved request settings from the initial request.
+        self.assertEqual(self.requests[0][0].command, self.requests[1][0].command)
+        self.assertEqual(
+            self.requests[0][0].headers.get('content-type'),
+            self.requests[1][0].headers.get('content-type')
+        )
+        self.assertEqual(self.requests[0][1], self.requests[1][1])
+
+        # ensure the redirect took place at the expected redirect location.
+        self.assertEqual(self.requests[1][0].path, "/" + self.redirect_flag)
+
     @unittest.skipIf(
         sys.platform == "darwin",
         "instance_ip_grouping_key() does not work on macOS."