[rubygems/rubygems] Refactor Webauthn listener response - Makes the response class a wrapper around Net::HTTPResponse - Builds a Net::HTTPResponse upon initialization - to_s returns a string representation of the response to send - Adds a Socket Responder class to send responses given a socket

https://github.com/rubygems/rubygems/commit/7513c220b6

Co-authored-by: Jacques Chester <jacques.chester@shopify.com>

lib/rubygems/webauthn_listener.rb

@@ -1,10 +1,6 @@
 # frozen_string_literal: true
 
-require_relative "webauthn_listener/response/response_ok"
-require_relative "webauthn_listener/response/response_no_content"
-require_relative "webauthn_listener/response/response_bad_request"
-require_relative "webauthn_listener/response/response_not_found"
-require_relative "webauthn_listener/response/response_method_not_allowed"
+require_relative "webauthn_listener/response"
 
 ##
 # The WebauthnListener class retrieves an OTP after a user successfully WebAuthns with the Gem host.
@@ -45,24 +41,26 @@ class Gem::WebauthnListener
       method, req_uri, _protocol = request_line.split(" ")
       req_uri = URI.parse(req_uri)
 
+      responder = SocketResponder.new(socket)
+
       unless root_path?(req_uri)
-        ResponseNotFound.send(socket, host)
+        responder.send(NotFoundResponse.for(host))
         raise Gem::WebauthnVerificationError, "Page at #{req_uri.path} not found."
       end
 
       case method.upcase
       when "OPTIONS"
-        ResponseNoContent.send(socket, host)
+        responder.send(NoContentResponse.for(host))
         next # will be GET
       when "GET"
         if otp = parse_otp_from_uri(req_uri)
-          ResponseOk.send(socket, host)
+          responder.send(OkResponse.for(host))
           return otp
         end
-        ResponseBadRequest.send(socket, host)
+        responder.send(BadRequestResponse.for(host))
         raise Gem::WebauthnVerificationError, "Did not receive OTP from #{host}."
       else
-        ResponseMethodNotAllowed.send(socket, host)
+        responder.send(MethodNotAllowedResponse.for(host))
         raise Gem::WebauthnVerificationError, "Invalid HTTP method #{method.upcase} received."
       end
     end
@@ -80,4 +78,15 @@ class Gem::WebauthnListener
     return if uri.query.nil?
     CGI.parse(uri.query).dig("code", 0)
   end
+
+  class SocketResponder
+    def initialize(socket)
+      @socket = socket
+    end
+
+    def send(response)
+      @socket.print response.to_s
+      @socket.close
+    end
+  end
 end

lib/rubygems/webauthn_listener/response.rb

@@ -1,70 +1,161 @@
 # frozen_string_literal: true
 
 ##
-# The WebauthnListener Response class is used by the WebauthnListener to print
-# the specified response to the Gem host using the provided socket. It also closes
-# the socket after printing the response.
+# The WebauthnListener Response class is used by the WebauthnListener to create
+# responses to be sent to the Gem host. It creates a Net::HTTPResponse instance
+# when initialized and can be converted to the appropriate format to be sent by a socket using `to_s`.
+# Net::HTTPResponse instances cannot be directly sent over a socket.
 #
 # Types of response classes:
-#   - ResponseOk
-#   - ResponseNoContent
-#   - ResponseBadRequest
-#   - ResponseNotFound
-#   - ResponseMethodNotAllowed
+#   - OkResponse
+#   - NoContentResponse
+#   - BadRequestResponse
+#   - NotFoundResponse
+#   - MethodNotAllowedResponse
 #
-# Example:
-#   socket = TCPSocket.new(host, port)
-#   Gem::WebauthnListener::ResponseOk.send(socket, host)
+# Example usage:
+#
+#   server = TCPServer.new(0)
+#   socket = server.accept
+#
+#   response = OkResponse.for("https://rubygems.example")
+#   socket.print response.to_s
+#   socket.close
 #
 
 class Gem::WebauthnListener
   class Response
-    attr_reader :host
+    attr_reader :http_response
+
+    def self.for(host)
+      new(host)
+    end
 
     def initialize(host)
       @host = host
+
+      build_http_response
     end
 
-    def self.send(socket, host)
-      socket.print new(host).payload
-      socket.close
-    end
+    def to_s
+      status_line = "HTTP/#{@http_response.http_version} #{@http_response.code} #{@http_response.message}\r\n"
+      headers = @http_response.to_hash.map {|header, value| "#{header}: #{value.join(", ")}\r\n" }.join + "\r\n"
+      body = @http_response.body ? "#{@http_response.body}\n" : ""
 
-    def payload
-      status_line_and_connection + access_control_headers + content
+      status_line + headers + body
     end
 
     private
 
-    def status_line_and_connection
-      <<~RESPONSE
-        HTTP/1.1 #{status}
-        Connection: close
-      RESPONSE
+    # Must be implemented in subclasses
+    def code
+      raise NotImplementedError
     end
 
-    def access_control_headers
-      <<~RESPONSE
-        Access-Control-Allow-Origin: #{host}
-        Access-Control-Allow-Methods: POST
-        Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
-      RESPONSE
-    end
-
-    def content
-      return "" unless body
-      <<~RESPONSE
-        Content-Type: text/plain
-        Content-Length: #{body.bytesize}
-
-        #{body}
-      RESPONSE
-    end
-
-    def status
+    def reason_phrase
       raise NotImplementedError
     end
 
     def body; end
+
+    def build_http_response
+      response_class = Net::HTTPResponse::CODE_TO_OBJ[code.to_s]
+      @http_response = response_class.new("1.1", code, reason_phrase)
+      @http_response.instance_variable_set(:@read, true)
+
+      add_connection_header
+      add_access_control_headers
+      add_body
+    end
+
+    def add_connection_header
+      @http_response["connection"] = "close"
+    end
+
+    def add_access_control_headers
+      @http_response["access-control-allow-origin"] = @host
+      @http_response["access-control-allow-methods"] = "POST"
+      @http_response["access-control-allow-headers"] = %w[Content-Type Authorization x-csrf-token]
+    end
+
+    def add_body
+      return unless body
+      @http_response["content-type"] = "text/plain"
+      @http_response["content-length"] = body.bytesize
+      @http_response.instance_variable_set(:@body, body)
+    end
+  end
+
+  class OkResponse < Response
+    private
+
+    def code
+      200
+    end
+
+    def reason_phrase
+      "OK"
+    end
+
+    def body
+      "success"
+    end
+  end
+
+  class NoContentResponse < Response
+    private
+
+    def code
+      204
+    end
+
+    def reason_phrase
+      "No Content"
+    end
+  end
+
+  class BadRequestResponse < Response
+    private
+
+    def code
+      400
+    end
+
+    def reason_phrase
+      "Bad Request"
+    end
+
+    def body
+      "missing code parameter"
+    end
+  end
+
+  class NotFoundResponse < Response
+    private
+
+    def code
+      404
+    end
+
+    def reason_phrase
+      "Not Found"
+    end
+  end
+
+  class MethodNotAllowedResponse < Response
+    private
+
+    def code
+      405
+    end
+
+    def reason_phrase
+      "Method Not Allowed"
+    end
+
+    def add_access_control_headers
+      super
+      @http_response["allow"] = %w[GET OPTIONS]
+    end
   end
 end

lib/rubygems/webauthn_listener/response/response_bad_request.rb

@@ -1,14 +0,0 @@
-# frozen_string_literal: true
-require_relative "../response"
-
-class Gem::WebauthnListener::ResponseBadRequest < Gem::WebauthnListener::Response
-  private
-
-  def status
-    "400 Bad Request"
-  end
-
-  def body
-    "missing code parameter"
-  end
-end

lib/rubygems/webauthn_listener/response/response_method_not_allowed.rb

@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-require_relative "../response"
-
-class Gem::WebauthnListener::ResponseMethodNotAllowed < Gem::WebauthnListener::Response
-  private
-
-  def status
-    "405 Method Not Allowed"
-  end
-
-  def content
-    <<~RESPONSE
-      Allow: GET, OPTIONS
-    RESPONSE
-  end
-end

lib/rubygems/webauthn_listener/response/response_no_content.rb

@@ -1,10 +0,0 @@
-# frozen_string_literal: true
-require_relative "../response"
-
-class Gem::WebauthnListener::ResponseNoContent < Gem::WebauthnListener::Response
-  private
-
-  def status
-    "204 No Content"
-  end
-end

lib/rubygems/webauthn_listener/response/response_not_found.rb

@@ -1,10 +0,0 @@
-# frozen_string_literal: true
-require_relative "../response"
-
-class Gem::WebauthnListener::ResponseNotFound < Gem::WebauthnListener::Response
-  private
-
-  def status
-    "404 Not Found"
-  end
-end

lib/rubygems/webauthn_listener/response/response_ok.rb

@@ -1,14 +0,0 @@
-# frozen_string_literal: true
-require_relative "../response"
-
-class Gem::WebauthnListener::ResponseOk < Gem::WebauthnListener::Response
-  private
-
-  def status
-    "200 OK"
-  end
-
-  def body
-    "success"
-  end
-end

test/rubygems/test_webauthn_listener_response.rb

@@ -1,116 +1,93 @@
 # frozen_string_literal: true
 
 require_relative "helper"
-require "rubygems/webauthn_listener/response/response_ok"
-require "rubygems/webauthn_listener/response/response_no_content"
-require "rubygems/webauthn_listener/response/response_bad_request"
-require "rubygems/webauthn_listener/response/response_not_found"
-require "rubygems/webauthn_listener/response/response_method_not_allowed"
+require "rubygems/webauthn_listener/response"
 
 class WebauthnListenerResponseTest < Gem::TestCase
-  class MockResponse < Gem::WebauthnListener::Response
-    def payload
-      "hello world"
-    end
-  end
-
   def setup
     super
     @host = "rubygems.example"
   end
 
-  def test_ok_response_payload
-    payload = Gem::WebauthnListener::ResponseOk.new(@host).payload
-
-    expected_payload = <<~RESPONSE
-      HTTP/1.1 200 OK
-      Connection: close
-      Access-Control-Allow-Origin: rubygems.example
-      Access-Control-Allow-Methods: POST
-      Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
-      Content-Type: text/plain
-      Content-Length: 7
+  def test_ok_response_to_s
+    to_s = Gem::WebauthnListener::OkResponse.new(@host).to_s
 
+    expected_to_s = <<~RESPONSE
+      HTTP/1.1 200 OK\r
+      connection: close\r
+      access-control-allow-origin: rubygems.example\r
+      access-control-allow-methods: POST\r
+      access-control-allow-headers: Content-Type, Authorization, x-csrf-token\r
+      content-type: text/plain\r
+      content-length: 7\r
+      \r
       success
     RESPONSE
 
-    assert_equal expected_payload, payload
+    assert_equal expected_to_s, to_s
   end
 
-  def test_no_payload_response_payload
-    payload = Gem::WebauthnListener::ResponseNoContent.new(@host).payload
+  def test_no_to_s_response_to_s
+    to_s = Gem::WebauthnListener::NoContentResponse.new(@host).to_s
 
-    expected_payload = <<~RESPONSE
-      HTTP/1.1 204 No Content
-      Connection: close
-      Access-Control-Allow-Origin: rubygems.example
-      Access-Control-Allow-Methods: POST
-      Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
+    expected_to_s = <<~RESPONSE
+      HTTP/1.1 204 No Content\r
+      connection: close\r
+      access-control-allow-origin: rubygems.example\r
+      access-control-allow-methods: POST\r
+      access-control-allow-headers: Content-Type, Authorization, x-csrf-token\r
+      \r
     RESPONSE
 
-    assert_equal expected_payload, payload
+    assert_equal expected_to_s, to_s
   end
 
-  def test_method_not_allowed_response_payload
-    payload = Gem::WebauthnListener::ResponseMethodNotAllowed.new(@host).payload
+  def test_method_not_allowed_response_to_s
+    to_s = Gem::WebauthnListener::MethodNotAllowedResponse.new(@host).to_s
 
-    expected_payload = <<~RESPONSE
-      HTTP/1.1 405 Method Not Allowed
-      Connection: close
-      Access-Control-Allow-Origin: rubygems.example
-      Access-Control-Allow-Methods: POST
-      Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
-      Allow: GET, OPTIONS
+    expected_to_s = <<~RESPONSE
+      HTTP/1.1 405 Method Not Allowed\r
+      connection: close\r
+      access-control-allow-origin: rubygems.example\r
+      access-control-allow-methods: POST\r
+      access-control-allow-headers: Content-Type, Authorization, x-csrf-token\r
+      allow: GET, OPTIONS\r
+      \r
     RESPONSE
 
-    assert_equal expected_payload, payload
+    assert_equal expected_to_s, to_s
   end
 
-  def test_method_not_found_response_payload
-    payload = Gem::WebauthnListener::ResponseNotFound.new(@host).payload
+  def test_method_not_found_response_to_s
+    to_s = Gem::WebauthnListener::NotFoundResponse.new(@host).to_s
 
-    expected_payload = <<~RESPONSE
-      HTTP/1.1 404 Not Found
-      Connection: close
-      Access-Control-Allow-Origin: rubygems.example
-      Access-Control-Allow-Methods: POST
-      Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
+    expected_to_s = <<~RESPONSE
+      HTTP/1.1 404 Not Found\r
+      connection: close\r
+      access-control-allow-origin: rubygems.example\r
+      access-control-allow-methods: POST\r
+      access-control-allow-headers: Content-Type, Authorization, x-csrf-token\r
+      \r
     RESPONSE
 
-    assert_equal expected_payload, payload
+    assert_equal expected_to_s, to_s
   end
 
-  def test_bad_request_response_payload
-    payload = Gem::WebauthnListener::ResponseBadRequest.new(@host).payload
-
-    expected_payload = <<~RESPONSE
-      HTTP/1.1 400 Bad Request
-      Connection: close
-      Access-Control-Allow-Origin: rubygems.example
-      Access-Control-Allow-Methods: POST
-      Access-Control-Allow-Headers: Content-Type, Authorization, x-csrf-token
-      Content-Type: text/plain
-      Content-Length: 22
+  def test_bad_request_response_to_s
+    to_s = Gem::WebauthnListener::BadRequestResponse.new(@host).to_s
 
+    expected_to_s = <<~RESPONSE
+      HTTP/1.1 400 Bad Request\r
+      connection: close\r
+      access-control-allow-origin: rubygems.example\r
+      access-control-allow-methods: POST\r
+      access-control-allow-headers: Content-Type, Authorization, x-csrf-token\r
+      content-type: text/plain\r
+      content-length: 22\r
+      \r
       missing code parameter
     RESPONSE
 
-    assert_equal expected_payload, payload
-  end
-
-  def test_send_response
-    server = TCPServer.new "localhost", 5678
-    thread = Thread.new do
-      receive_socket = server.accept
-      Thread.current[:payload] = receive_socket.read
-      receive_socket.close
-    end
-
-    send_socket = TCPSocket.new "localhost", 5678
-    MockResponse.send(send_socket, @host)
-
-    thread.join
-    assert_equal "hello world", thread[:payload]
-    assert_predicate send_socket, :closed?
+    assert_equal expected_to_s, to_s
   end
 end