[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>

7 files changed, 148 insertions, 112 deletions

diff --git a/lib/rubygems/webauthn_listener.rb b/lib/rubygems/webauthn_listener.rb
index ab80281ef5..22f7ea2011 100644
--- a/lib/rubygems/webauthn_listener.rb
+++ b/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
diff --git a/lib/rubygems/webauthn_listener/response.rb b/lib/rubygems/webauthn_listener/response.rb
index 8596e7bd69..baa769c4ae 100644
--- a/lib/rubygems/webauthn_listener/response.rb
+++ b/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 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" : ""
+
+      status_line + headers + body
+    end
+
+    private
+
+    # Must be implemented in subclasses
+    def code
+      raise NotImplementedError
+    end
+
+    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 self.send(socket, host)
-      socket.print new(host).payload
-      socket.close
+    def code
+      200
     end
 
-    def payload
-      status_line_and_connection + access_control_headers + content
+    def reason_phrase
+      "OK"
     end
 
+    def body
+      "success"
+    end
+  end
+
+  class NoContentResponse < Response
     private
 
-    def status_line_and_connection
-      <<~RESPONSE
-        HTTP/1.1 #{status}
-        Connection: close
-      RESPONSE
+    def code
+      204
     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
+    def reason_phrase
+      "No Content"
     end
+  end
 
-    def content
-      return "" unless body
-      <<~RESPONSE
-        Content-Type: text/plain
-        Content-Length: #{body.bytesize}
+  class BadRequestResponse < Response
+    private
 
-        #{body}
-      RESPONSE
+    def code
+      400
     end
 
-    def status
-      raise NotImplementedError
+    def reason_phrase
+      "Bad Request"
     end
 
-    def body; 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
diff --git a/lib/rubygems/webauthn_listener/response/response_bad_request.rb b/lib/rubygems/webauthn_listener/response/response_bad_request.rb
deleted file mode 100644
index 031c72e08e..0000000000
--- a/lib/rubygems/webauthn_listener/response/response_bad_request.rb
+++ /dev/null
@@ -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
diff --git a/lib/rubygems/webauthn_listener/response/response_method_not_allowed.rb b/lib/rubygems/webauthn_listener/response/response_method_not_allowed.rb
deleted file mode 100644
index ae071fc242..0000000000
--- a/lib/rubygems/webauthn_listener/response/response_method_not_allowed.rb
+++ /dev/null
@@ -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
diff --git a/lib/rubygems/webauthn_listener/response/response_no_content.rb b/lib/rubygems/webauthn_listener/response/response_no_content.rb
deleted file mode 100644
index feb6cade8c..0000000000
--- a/lib/rubygems/webauthn_listener/response/response_no_content.rb
+++ /dev/null
@@ -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
diff --git a/lib/rubygems/webauthn_listener/response/response_not_found.rb b/lib/rubygems/webauthn_listener/response/response_not_found.rb
deleted file mode 100644
index c1207cea36..0000000000
--- a/lib/rubygems/webauthn_listener/response/response_not_found.rb
+++ /dev/null
@@ -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
diff --git a/lib/rubygems/webauthn_listener/response/response_ok.rb b/lib/rubygems/webauthn_listener/response/response_ok.rb
deleted file mode 100644
index 83966b58b8..0000000000
--- a/lib/rubygems/webauthn_listener/response/response_ok.rb
+++ /dev/null
@@ -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