lib/eeny-meeny/middleware.rb in eeny-meeny-2.1.1 vs lib/eeny-meeny/middleware.rb in eeny-meeny-2.1.2

@@ -6,57 +6,54 @@
 require 'eeny-meeny/models/cookie'
 
 module EenyMeeny
   class Middleware
 
+    # Headers
+    HTTP_COOKIE    = 'HTTP_COOKIE'.freeze
+    REQUEST_METHOD = 'REQUEST_METHOD'.freeze
+    QUERY_STRING   = 'QUERY_STRING'.freeze
+
     def initialize(app)
       @app = app
       @experiments = EenyMeeny::Experiment.find_all
       @cookie_config = EenyMeeny.config.cookies
     end
 
     def call(env)
-      request = Rack::Request.new(env)
-      cookies = request.cookies
-      now = Time.zone.now
-      new_cookies = {}
-      existing_set_cookie_header = env['Set-Cookie']
+      cookies          = Rack::Utils.parse_query(env[HTTP_COOKIE],';,')  { |s| Rack::Utils.unescape(s) rescue s }
+      query_parameters = query_hash(env)
+      now              = Time.zone.now
+      new_cookies      = {}
       # Prepare for experiments.
       @experiments.each do |experiment|
         # Skip inactive experiments
         next unless experiment.active?(now)
-        # Trigger experiment through query parmeters
+        # Trigger experiment through query parameters
         cookie_name = EenyMeeny::Cookie.cookie_name(experiment)
-        has_experiment_trigger = EenyMeeny.config.query_parameters[:experiment] && request.params.has_key?(cookie_name)
+        has_experiment_trigger = EenyMeeny.config.query_parameters[:experiment] && query_parameters.key?(cookie_name)
         # skip experiments that already have a cookie
-        if has_experiment_trigger || !cookies.has_key?(cookie_name)
-          cookie = if has_experiment_trigger
-                     # Trigger experiment variation through query parameter.
-                     EenyMeeny::Cookie.create_for_experiment_variation(experiment, request.params[cookie_name].to_sym, @cookie_config)
-                   else
-                     EenyMeeny::Cookie.create_for_experiment(experiment, @cookie_config)
-                   end
-          # Set HTTP_COOKIE header to enable experiment on first pageview
-          env = add_http_cookie(env, cookie, precede: has_experiment_trigger)
-          new_cookies[cookie.name] = cookie
-        end
+        next unless has_experiment_trigger || !cookies.key?(cookie_name)
+        cookie = if has_experiment_trigger
+                   # Trigger experiment variation through query parameter.
+                   EenyMeeny::Cookie.create_for_experiment_variation(experiment, query_parameters[cookie_name].to_sym, @cookie_config)
+                 else
+                   EenyMeeny::Cookie.create_for_experiment(experiment, @cookie_config)
+                 end
+        # Set HTTP_COOKIE header to enable experiment on first pageview
+        env = add_or_replace_http_cookie(env, cookie)
+        new_cookies[cookie.name] = cookie
       end
       # Prepare smoke tests (if enabled through query parameters)
       if EenyMeeny.config.query_parameters[:smoke_test]
-        if request.params.has_key?('smoke_test_id') && (request.params['smoke_test_id'] =~ /[A-Za-z_]+/)
+        if query_parameters.key?('smoke_test_id') && (query_parameters['smoke_test_id'] =~ /\A[A-Za-z_]+\z/)
           # Set HTTP_COOKIE header to enable smoke test on first pageview
-          cookie = EenyMeeny::Cookie.create_for_smoke_test(request.params['smoke_test_id'])
-          env = add_http_cookie(env, cookie, precede: true)
+          cookie = EenyMeeny::Cookie.create_for_smoke_test(query_parameters['smoke_test_id'])
+          env = add_or_replace_http_cookie(env, cookie)
           new_cookies[cookie.name] = cookie
         end
       end
-      # Clean up 'Set-Cookie' header.
-      if existing_set_cookie_header.nil?
-        env.delete('Set-Cookie')
-      else
-        env['Set-Cookie'] = existing_set_cookie_header
-      end
       # Delegate to app
       status, headers, body = @app.call(env)
       response = Rack::Response.new(body, status, headers)
       # Add new cookies to 'Set-Cookie' header
       new_cookies.each do |key, value|
@@ -64,22 +61,25 @@
       end
       response.finish
     end
 
     private
-    def add_http_cookie(env, cookie, precede: false)
-      env['Set-Cookie'] = ''
-      Rack::Utils.set_cookie_header!(env,
-                                     cookie.name,
-                                     cookie.to_h)
-      env['HTTP_COOKIE'] = '' if env['HTTP_COOKIE'].nil?
-      if precede
-        # Prepend cookie to the 'HTTP_COOKIE' header. This ensures it overwrites existing cookies when present.
-        env['HTTP_COOKIE'] = env['Set-Cookie'] + '; ' + env['HTTP_COOKIE']
-      else
-        env['HTTP_COOKIE'] += '; ' unless env['HTTP_COOKIE'].empty?
-        env['HTTP_COOKIE'] += env['Set-Cookie']
-      end
+
+    def query_hash(env)
+      # Query Params are only relevant if EenyMeeny.config have them enabled.
+      return {} unless EenyMeeny.config.query_parameters[:experiment] || EenyMeeny.config.query_parameters[:smoke_test]
+      # Query Params are only relevant to HTTP GET requests.
+      return {} unless env[REQUEST_METHOD] == 'GET'
+      Rack::Utils.parse_query(env[QUERY_STRING], '&;')
+    end
+
+    def add_or_replace_http_cookie(env, cookie)
+      cookie_name_escaped = Rack::Utils.escape(cookie.name)
+      cookie_string = "#{cookie_name_escaped}=#{Rack::Utils.escape(cookie.value)}"
+      env[HTTP_COOKIE] = '' if env[HTTP_COOKIE].nil?
+      return env if env[HTTP_COOKIE].sub!(/#{Regexp.escape(cookie_name_escaped)}=[^;]+/, cookie_string)
+      env[HTTP_COOKIE] += '; ' unless env[HTTP_COOKIE].empty?
+      env[HTTP_COOKIE] += cookie_string
       env
     end
   end
 end