Refactoring: Move rack middleware (#34140)

config/application.rb

@@ -34,11 +34,11 @@ require_relative '../lib/paperclip/transcoder'
 require_relative '../lib/paperclip/type_corrector'
 require_relative '../lib/paperclip/response_with_limit_adapter'
 require_relative '../lib/terrapin/multi_pipe_extensions'
+require_relative '../lib/mastodon/middleware/public_file_server'
+require_relative '../lib/mastodon/middleware/socket_cleanup'
 require_relative '../lib/mastodon/snowflake'
 require_relative '../lib/mastodon/feature'
 require_relative '../lib/mastodon/version'
-require_relative '../lib/mastodon/rack_middleware'
-require_relative '../lib/public_file_server_middleware'
 require_relative '../lib/devise/strategies/two_factor_ldap_authenticatable'
 require_relative '../lib/devise/strategies/two_factor_pam_authenticatable'
 require_relative '../lib/elasticsearch/client_extensions'
@@ -88,9 +88,9 @@ module Mastodon
     # We use our own middleware for this
     config.public_file_server.enabled = false
 
-    config.middleware.use PublicFileServerMiddleware if Rails.env.local? || ENV['RAILS_SERVE_STATIC_FILES'] == 'true'
+    config.middleware.use Mastodon::Middleware::PublicFileServer if Rails.env.local? || ENV['RAILS_SERVE_STATIC_FILES'] == 'true'
     config.middleware.use Rack::Attack
-    config.middleware.use Mastodon::RackMiddleware
+    config.middleware.use Mastodon::Middleware::SocketCleanup
 
     config.before_configuration do
       require 'mastodon/redis_configuration'

lib/mastodon/middleware/public_file_server.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'action_dispatch/middleware/static'
+
+module Mastodon
+  module Middleware
+    class PublicFileServer
+      SERVICE_WORKER_TTL = 7.days.to_i
+      CACHE_TTL          = 28.days.to_i
+
+      def initialize(app)
+        @app = app
+        @file_handler = ActionDispatch::FileHandler.new(Rails.application.paths['public'].first)
+      end
+
+      def call(env)
+        file = @file_handler.attempt(env)
+
+        # If the request is not a static file, move on!
+        return @app.call(env) if file.nil?
+
+        status, headers, response = file
+
+        # Set cache headers on static files. Some paths require different cache headers
+        headers['Cache-Control'] = begin
+          request_path = env['REQUEST_PATH']
+
+          if request_path.start_with?('/sw.js')
+            "public, max-age=#{SERVICE_WORKER_TTL}, must-revalidate"
+          elsif request_path.start_with?(paperclip_root_url)
+            "public, max-age=#{CACHE_TTL}, immutable"
+          else
+            "public, max-age=#{CACHE_TTL}, must-revalidate"
+          end
+        end
+
+        # Override the default CSP header set by the CSP middleware
+        headers['Content-Security-Policy'] = "default-src 'none'; form-action 'none'" if request_path.start_with?(paperclip_root_url)
+
+        headers['X-Content-Type-Options'] = 'nosniff'
+
+        [status, headers, response]
+      end
+
+      private
+
+      def paperclip_root_url
+        ENV.fetch('PAPERCLIP_ROOT_URL', '/system')
+      end
+    end
+  end
+end

lib/mastodon/middleware/socket_cleanup.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Mastodon
+  module Middleware
+    class SocketCleanup
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        @app.call(env)
+      ensure
+        clean_up_sockets!
+      end
+
+      private
+
+      def clean_up_sockets!
+        clean_up_redis_socket!
+        clean_up_statsd_socket!
+      end
+
+      def clean_up_redis_socket!
+        RedisConnection.pool.checkin if Thread.current[:redis]
+        Thread.current[:redis] = nil
+      end
+
+      def clean_up_statsd_socket!
+        Thread.current[:statsd_socket]&.close
+        Thread.current[:statsd_socket] = nil
+      end
+    end
+  end
+end

lib/mastodon/rack_middleware.rb

@@ -1,30 +0,0 @@
-# frozen_string_literal: true
-
-class Mastodon::RackMiddleware
-  def initialize(app)
-    @app = app
-  end
-
-  def call(env)
-    @app.call(env)
-  ensure
-    clean_up_sockets!
-  end
-
-  private
-
-  def clean_up_sockets!
-    clean_up_redis_socket!
-    clean_up_statsd_socket!
-  end
-
-  def clean_up_redis_socket!
-    RedisConnection.pool.checkin if Thread.current[:redis]
-    Thread.current[:redis] = nil
-  end
-
-  def clean_up_statsd_socket!
-    Thread.current[:statsd_socket]&.close
-    Thread.current[:statsd_socket] = nil
-  end
-end

lib/public_file_server_middleware.rb

@@ -1,48 +0,0 @@
-# frozen_string_literal: true
-
-require 'action_dispatch/middleware/static'
-
-class PublicFileServerMiddleware
-  SERVICE_WORKER_TTL = 7.days.to_i
-  CACHE_TTL          = 28.days.to_i
-
-  def initialize(app)
-    @app = app
-    @file_handler = ActionDispatch::FileHandler.new(Rails.application.paths['public'].first)
-  end
-
-  def call(env)
-    file = @file_handler.attempt(env)
-
-    # If the request is not a static file, move on!
-    return @app.call(env) if file.nil?
-
-    status, headers, response = file
-
-    # Set cache headers on static files. Some paths require different cache headers
-    headers['Cache-Control'] = begin
-      request_path = env['REQUEST_PATH']
-
-      if request_path.start_with?('/sw.js')
-        "public, max-age=#{SERVICE_WORKER_TTL}, must-revalidate"
-      elsif request_path.start_with?(paperclip_root_url)
-        "public, max-age=#{CACHE_TTL}, immutable"
-      else
-        "public, max-age=#{CACHE_TTL}, must-revalidate"
-      end
-    end
-
-    # Override the default CSP header set by the CSP middleware
-    headers['Content-Security-Policy'] = "default-src 'none'; form-action 'none'" if request_path.start_with?(paperclip_root_url)
-
-    headers['X-Content-Type-Options'] = 'nosniff'
-
-    [status, headers, response]
-  end
-
-  private
-
-  def paperclip_root_url
-    ENV.fetch('PAPERCLIP_ROOT_URL', '/system')
-  end
-end