feat: extend server with initialize and websocket_handler_for

lib/prism/server.rb

@@ -8,12 +8,24 @@ module Prism
       422 => 'Unprocessable Entity', 500 => 'Internal Server Error'
     }.freeze
 
-    attr_reader :host, :port
+    attr_reader :host, :port, :websocket_connection_count
 
-    def initialize(app, host: '127.0.0.1', port: 0)
+    # `max_websocket_connections:` is a hard cap on the number of LIVE
+    # WebSocket connections held by this process, across every registered
+    # endpoint and every stream - distinct from
+    # WebSocket::Hub#max_subscribers_per_stream, which only limits how many
+    # connections may subscribe to one particular broadcast stream. Without
+    # this, a process serving N different posts each under their own
+    # per-stream cap could still be driven to opening unbounded sockets by
+    # spreading connections across many streams. `nil` (default) preserves
+    # the old unbounded behavior.
+    def initialize(app, host: '127.0.0.1', port: 0, max_websocket_connections: nil)
       @app = app
       @host = host
       @requested_port = port
+      @max_websocket_connections = max_websocket_connections
+      @websocket_connection_count = 0
+      @websocket_connection_mutex = Mutex.new
     end
 
     def start
@@ -48,6 +60,10 @@ module Prism
       loop do
         env = read_request(socket)
         break unless env
+        if WebSocket::Handshake.request?(env) && (handler = websocket_handler_for(env))
+          handle_websocket(socket, env, handler)
+          break
+        end
         status, headers, body = @app.call(env)
         keep_alive = keep_alive?(env, headers)
         write_response(socket, status, headers, body, keep_alive)
@@ -59,6 +75,72 @@ module Prism
       socket.close unless socket.closed?
     end
 
+    # Looked up on `@app` directly (not through the middleware/Rack `call`
+    # chain) because an Upgrade: websocket request never gets a normal HTTP
+    # response - it either becomes a 101 handshake handed off to
+    # WebSocket::Connection, or it falls through to the ordinary Rack app
+    # (and 404s there) when no handler is registered for the path.
+    def websocket_handler_for(env)
+      return nil unless @app.respond_to?(:find_websocket_handler)
+      @app.find_websocket_handler(env['PATH_INFO'])
+    end
+
+    # Compacted so an unset `max_message_size:` (nil, the common case) is
+    # never passed through to Connection.new at all - passing it explicitly
+    # as nil would override Connection's own DEFAULT_MAX_MESSAGE_SIZE
+    # keyword default with nil instead of leaving it alone.
+    def connection_options(env)
+      return {} unless @app.respond_to?(:find_websocket_options)
+      @app.find_websocket_options(env['PATH_INFO']).compact
+    end
+
+    def handle_websocket(socket, env, handler_class)
+      unless acquire_websocket_slot
+        write_response(socket, 503, { 'Content-Type' => 'text/plain' }, ['too many websocket connections'], false)
+        return
+      end
+
+      begin
+        socket.write("HTTP/1.1 101 Switching Protocols\r\n")
+        WebSocket::Handshake.response_headers(env).each { |name, value| socket.write("#{name}: #{value}\r\n") }
+        socket.write("\r\n")
+
+        connection = WebSocket::Connection.new(socket, env: env, **connection_options(env))
+        handler = handler_class.new
+        handler.on_open(connection) if handler.respond_to?(:on_open)
+        begin
+          while (message = connection.read_message)
+            handler.on_message(connection, message) if handler.respond_to?(:on_message)
+          end
+        rescue WebSocket::ProtocolError
+          nil
+        ensure
+          handler.on_close(connection) if handler.respond_to?(:on_close)
+        end
+      ensure
+        release_websocket_slot
+      end
+    end
+
+    # Checked and incremented under the same mutex so two connections
+    # racing to grab the last slot cannot both succeed - the check and the
+    # reservation are one atomic step, not two.
+    def acquire_websocket_slot
+      return true unless @max_websocket_connections
+
+      @websocket_connection_mutex.synchronize do
+        return false if @websocket_connection_count >= @max_websocket_connections
+        @websocket_connection_count += 1
+        true
+      end
+    end
+
+    def release_websocket_slot
+      return unless @max_websocket_connections
+
+      @websocket_connection_mutex.synchronize { @websocket_connection_count -= 1 }
+    end
+
     def read_request(socket)
       request_line = socket.gets("\r\n")
       return nil if request_line.nil?
@@ -82,6 +164,7 @@ module Prism
         'QUERY_STRING' => query.to_s,
         'SERVER_NAME' => @host,
         'SERVER_PORT' => @port.to_s,
+        'REMOTE_ADDR' => remote_addr(socket),
         'rack.url_scheme' => 'http'
       }
       headers.each do |name, value|
@@ -94,23 +177,66 @@ module Prism
       env
     end
 
+    # The client's numeric IP, the standard Rack env key that
+    # Middleware::RateLimiting (and anything else keying off the caller's
+    # address) relies on - `peeraddr(false)` skips the reverse DNS lookup,
+    # since only the numeric address (index 3) is ever used.
+    def remote_addr(socket)
+      socket.peeraddr(false)[3]
+    rescue IOError, Errno::ENOTCONN, SocketError
+      nil
+    end
+
     def keep_alive?(env, headers)
       return false if headers['Connection'].to_s.downcase == 'close'
       env['HTTP_CONNECTION'].to_s.downcase != 'close'
     end
 
     def write_response(socket, status, headers, body, keep_alive)
+      if body.is_a?(StreamingBody)
+        write_chunked_response(socket, status, headers, body, keep_alive)
+      else
+        write_buffered_response(socket, status, headers, body, keep_alive)
+      end
+    end
+
+    def write_buffered_response(socket, status, headers, body, keep_alive)
       payload = body.join
-      out = +"HTTP/1.1 #{status} #{status_text(status)}\r\n"
       response_headers = headers.dup
       response_headers['Content-Length'] = payload.bytesize.to_s
+      socket.write(status_line_and_headers(status, response_headers, keep_alive))
+      socket.write(payload) unless payload.empty?
+    end
+
+    # The body's total size is unknowable up front - that is the entire
+    # reason to stream instead of building one payload String first - so
+    # this sends `Transfer-Encoding: chunked` (RFC 7230 SS4.1) instead of a
+    # `Content-Length`, and writes each chunk `body#each` produces as its
+    # own `<hex size>\r\n<data>\r\n` frame the moment it is produced,
+    # followed by the zero-length terminating chunk once the producer block
+    # returns. A `write` call happening as fast, or as slowly, as the
+    # action generates data is the entire point - nothing here buffers
+    # more than one chunk at a time.
+    def write_chunked_response(socket, status, headers, body, keep_alive)
+      response_headers = headers.dup
+      response_headers.delete('Content-Length')
+      response_headers['Transfer-Encoding'] = 'chunked'
+      socket.write(status_line_and_headers(status, response_headers, keep_alive))
+      body.each do |chunk|
+        next if chunk.empty?
+        socket.write("#{chunk.bytesize.to_s(16)}\r\n#{chunk}\r\n")
+      end
+      socket.write("0\r\n\r\n")
+    end
+
+    def status_line_and_headers(status, response_headers, keep_alive)
+      out = +"HTTP/1.1 #{status} #{status_text(status)}\r\n"
       response_headers['Connection'] = keep_alive ? 'keep-alive' : 'close'
       response_headers.each do |name, value|
         value.to_s.split("\n").each { |entry| out << "#{name}: #{entry}\r\n" }
       end
       out << "\r\n"
-      socket.write(out)
+      out
-      socket.write(payload) unless payload.empty?
     end
 
     def status_text(status)

lib/prism/session.rb

@@ -9,11 +9,20 @@ module Prism
       "#{payload}#{Separator}#{sign(payload, secret)}"
     end
 
-    def decode(cookie, secret)
+    # `previous_secrets:` lets an app rotate its session secret (the
+    # Middleware::Session equivalent of Application#enable_attribute_
+    # encryption's `previous_secrets:`) without logging out every signed-in
+    # user the moment it deploys: a cookie signed under any retired secret
+    # still decodes here, while #encode (called once per request that
+    # touches the session) always signs under the current one, so every
+    # session is transparently re-signed onto the new secret the next time
+    # it changes at all.
+    def decode(cookie, secret, previous_secrets: [])
       return {} if cookie.nil? || cookie.empty?
       payload, signature = cookie.split(Separator, 2)
       return {} if payload.nil? || signature.nil?
-      return {} unless secure_compare(signature, sign(payload, secret))
+      candidates = [secret, *previous_secrets]
+      return {} unless candidates.any? { |candidate| secure_compare(signature, sign(payload, candidate)) }
       JSON.parse(decode64(payload))
     rescue JSON::ParserError, ArgumentError
       {}