From 42860ced1d3ad1deff061f6b7736e53a5fdfc5fe Mon Sep 17 00:00:00 2001
From: James Turley <jamesturley@gocardless.com>
Date: Thu, 30 Jul 2020 08:49:45 +0100
Subject: [PATCH] Add port option to exporter middleware

If the port option is set, all requests for /metrics on other ports will
be forwarded to the app. If it is unset or nil, or the ports match,
export things as usual.

Allows separate mounting of metrics and main app to enforce different
security setups etc.

Signed-off-by: James Turley <jamesturley@gocardless.com>
---
 lib/prometheus/middleware/exporter.rb       |  7 +++++-
 spec/prometheus/middleware/exporter_spec.rb | 28 ++++++++++++++++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/lib/prometheus/middleware/exporter.rb b/lib/prometheus/middleware/exporter.rb
index 5a74d8e9..640a3985 100644
--- a/lib/prometheus/middleware/exporter.rb
+++ b/lib/prometheus/middleware/exporter.rb
@@ -21,11 +21,12 @@ def initialize(app, options = {})
         @app = app
         @registry = options[:registry] || Client.registry
         @path = options[:path] || '/metrics'
+        @port = options[:port]
         @acceptable = build_dictionary(FORMATS, FALLBACK)
       end
 
       def call(env)
-        if env['PATH_INFO'] == @path
+        if metrics_port?(env['SERVER_PORT']) && env['PATH_INFO'] == @path
           format = negotiate(env, @acceptable)
           format ? respond_with(format) : not_acceptable(FORMATS)
         else
@@ -86,6 +87,10 @@ def build_dictionary(formats, fallback)
           memo[format::MEDIA_TYPE] = format
         end
       end
+
+      def metrics_port?(request_port)
+        @port.nil? || @port.to_s == request_port
+      end
     end
   end
 end
diff --git a/spec/prometheus/middleware/exporter_spec.rb b/spec/prometheus/middleware/exporter_spec.rb
index 8916513b..5299fc1e 100644
--- a/spec/prometheus/middleware/exporter_spec.rb
+++ b/spec/prometheus/middleware/exporter_spec.rb
@@ -6,13 +6,14 @@
 describe Prometheus::Middleware::Exporter do
   include Rack::Test::Methods
 
+  let(:options) { { registry: registry } }
   let(:registry) do
     Prometheus::Client::Registry.new
   end
 
   let(:app) do
     app = ->(_) { [200, { 'Content-Type' => 'text/html' }, ['OK']] }
-    described_class.new(app, registry: registry)
+    described_class.new(app, **options)
   end
 
   context 'when requesting app endpoints' do
@@ -96,5 +97,30 @@
 
       include_examples 'ok', { 'HTTP_ACCEPT' => accept }, text
     end
+
+    context 'when a port is specified' do
+      let(:options) { { registry: registry, port: 9999 } }
+
+      context 'when a request is on the specified port' do
+        it 'responds with 200 OK' do
+          registry.counter(:foo, docstring: 'foo counter').increment(by: 9)
+
+          get 'http://example.org:9999/metrics', nil, {}
+
+          expect(last_response.status).to eql(200)
+          expect(last_response.header['Content-Type']).to eql(text::CONTENT_TYPE)
+          expect(last_response.body).to eql(text.marshal(registry))
+        end
+      end
+
+      context 'when a request is not on the specified port' do
+        it 'returns the app response' do
+          get 'http://example.org:8888/metrics', nil, {}
+
+          expect(last_response).to be_ok
+          expect(last_response.body).to eql('OK')
+        end
+      end
+    end
   end
 end