lib/marlowe/middleware.rb in marlowe-1.0.3 vs lib/marlowe/middleware.rb in marlowe-2.0

@@ -1,36 +1,97 @@
+# frozen_string_literal: true
+
 require 'rack'
 require 'request_store'
 require 'securerandom'
 
 module Marlowe
-  # Marlowe correlation id middleware. Including this into your
-  # middleware stack will add a correlation id header as an incoming
-  # request, and save that id in a request session variable.
-
-  # Name of the default header to look for and put the correlation id in.
-  CORRELATION_HEADER = 'Correlation-Id'.freeze
-
+  # Marlowe correlation id middleware. Including this into your middleware
+  # stack will add a correlation id header as an incoming request, and save
+  # that id in a request session variable.
   class Middleware
-    # Sets the the rack application to +app+
-    def initialize(app, opts={})
+    # The name of the default header to look for and put the correlation id in.
+    CORRELATION_HEADER = 'X-Request-Id' #:nodoc:
+
+    # Configure the Marlowe middleware to call +app+ with options +opts+.
+    #
+    # === Options
+    #
+    # <tt>:header</tt>:: The name of the header to inspect. Defaults to
+    #                    'X-Request-Id'. Also available as
+    #                    <tt>:correlation_header</tt>.
+    # <tt>:handler</tt>:: The handler for request correlation IDs. Defaults to
+    #                     sanitizing provided request IDs or generating a UUID.
+    #                     If <tt>:simple</tt> is provided, provided request IDs
+    #                     will not be sanitized. A callable (expecting a single
+    #                     input of any possible existing request ID) may be
+    #                     provided to introduce more complex request ID
+    #                     handling.
+    # <tt>:return</tt>:: If +true+ (the default), the request correlation ID
+    #                    will be returned as part of the response headers.
+    # <tt>:action_dispatch</tt>:: If +true+, Marlowe will add code to behave
+    #                             like <tt>ActionDispatch::RequestId</tt>.
+    #                             Depends on <tt>ActionDispatch::Request</tt>.
+    def initialize(app, opts = {})
       @app = app
-      @correlation_header = format_http_header(opts[:correlation_header] || Marlowe::CORRELATION_HEADER)
+      @header, @http_header = format_header_name(
+        opts[:header] || opts[:correlation_header] || CORRELATION_HEADER
+      )
+      @handler = opts.fetch(:handler, :clean)
+      @return = opts.fetch(:return, true)
+      @action_dispatch = opts.fetch(:action_dispatch, false)
     end
 
     # Stores the incoming correlation id from the +env+ hash. If the correlation
     # id has not been sent, a new UUID is generated and the +env+ is modified.
     def call(env)
-      env[@correlation_header] ||= SecureRandom.uuid
-      RequestStore.store[:correlation_id] = env[@correlation_header]
+      req_id = make_request_id(env[@http_header])
+      RequestStore.store[:correlation_id] = env[@http_header] = req_id
 
-      @status, @headers, @response = @app.call(env)
-      [@status, @headers, @response]
+      if @action_dispatch
+        req = ActionDispatch::Request.new(env)
+        req.request_id = req_id
+      end
+
+      @app.call(env).tap { |_status, headers, _body|
+        if @return
+          headers[@header] = if @action_dispatch
+                               req.request_id
+                             else
+                               RequestStore.store[:correlation_id]
+                             end
+        end
+      }
     end
 
     private
 
-    def format_http_header(header)
-      ("HTTP_" + header.gsub(/-/, '_').upcase).freeze
+    def format_header_name(header)
+      [
+        header.to_s.tr('_', '-').freeze,
+        "HTTP_#{header.to_s.tr('-', '_').upcase}"
+      ]
+    end
+
+    def make_request_id(request_id)
+      if @handler == :simple
+        simple(request_id)
+      elsif @handler.kind_of?(Proc)
+        simple(@handler.call(request_id))
+      else
+        clean(request_id)
+      end
+    end
+
+    def clean(request_id)
+      simple(request_id).gsub(/[^\w\-]/, '')[0, 255]
+    end
+
+    def simple(request_id)
+      if request_id && !request_id.empty? && request_id !~ /\A[[:space]]*\z/
+        request_id
+      else
+        SecureRandom.uuid
+      end
     end
   end
 end