[rubygems/rubygems] Extract polling logic into its own class

https://github.com/rubygems/rubygems/commit/218b83abed

3 files changed, 205 insertions, 41 deletions

diff --git a/lib/rubygems/gemcutter_utilities.rb b/lib/rubygems/gemcutter_utilities.rb
index c43745c504..fb1a42b5ce 100644
--- a/lib/rubygems/gemcutter_utilities.rb
+++ b/lib/rubygems/gemcutter_utilities.rb
@@ -3,6 +3,7 @@
 require_relative "remote_fetcher"
 require_relative "text"
 require_relative "webauthn_listener"
+require_relative "gemcutter_utilities/webauthn_poller"
 
 ##
 # Utility methods for using the RubyGems API.
@@ -259,7 +260,7 @@ module Gem::GemcutterUtilities
       url_with_port = "#{webauthn_url}?port=#{port}"
       say "You have enabled multi-factor authentication. Please visit #{url_with_port} to authenticate via security device. If you can't verify using WebAuthn but have OTP enabled, you can re-run the gem signin command with the `--otp [your_code]` option."
 
-      threads = [socket_thread(server), poll_thread(webauthn_url, credentials)]
+      threads = [socket_thread(server), WebauthnPoller.poll_thread(options, host, webauthn_url, credentials)]
       otp_thread = wait_for_otp_thread(*threads)
 
       threads.each(&:join)
@@ -302,35 +303,6 @@ module Gem::GemcutterUtilities
     thread
   end
 
-  def poll_thread(webauthn_url, credentials)
-    thread = Thread.new do
-      Timeout.timeout(300) do
-        loop do
-          response = webauthn_verification_poll_response(webauthn_url, credentials)
-          raise Gem::WebauthnVerificationError, response.message unless response.is_a?(Net::HTTPSuccess)
-
-          require "json"
-          parsed_response = JSON.parse(response.body)
-          case parsed_response["status"]
-          when "pending"
-            sleep 5
-          when "success"
-            Thread.current[:otp] = parsed_response["code"]
-            break
-          else
-            raise Gem::WebauthnVerificationError, parsed_response["message"]
-          end
-        end
-      end
-    rescue Gem::WebauthnVerificationError, Timeout::Error => e
-      Thread.current[:error] = e
-    end
-    thread.abort_on_exception = true
-    thread.report_on_exception = false
-
-    thread
-  end
-
   def webauthn_verification_url(credentials)
     response = rubygems_api_request(:post, "api/v1/webauthn_verification") do |request|
       if credentials.empty?
@@ -342,17 +314,6 @@ module Gem::GemcutterUtilities
     response.is_a?(Net::HTTPSuccess) ? response.body : nil
   end
 
-  def webauthn_verification_poll_response(webauthn_url, credentials)
-    webauthn_token = %r{(?<=\/)[^\/]+(?=$)}.match(webauthn_url)[0]
-    rubygems_api_request(:get, "api/v1/webauthn_verification/#{webauthn_token}/status.json") do |request|
-      if credentials.empty?
-        request.add_field "Authorization", api_key
-      else
-        request.basic_auth credentials[:email], credentials[:password]
-      end
-    end
-  end
-
   def pretty_host(host)
     if default_host?
       "RubyGems.org"
diff --git a/lib/rubygems/gemcutter_utilities/webauthn_poller.rb b/lib/rubygems/gemcutter_utilities/webauthn_poller.rb
new file mode 100644
index 0000000000..766b38584e
--- /dev/null
+++ b/lib/rubygems/gemcutter_utilities/webauthn_poller.rb
@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+##
+# The WebauthnPoller class retrieves an OTP after a user successfully WebAuthns. An instance
+# polls the Gem host for the OTP code. The polling request (api/v1/webauthn_verification/<webauthn_token>/status.json)
+# is sent to the Gem host every 5 seconds and will timeout after 5 minutes. If the status field in the json response
+# is "success", the code field will contain the OTP code.
+#
+# Example usage:
+#
+#   thread = Gem::WebauthnPoller.poll_thread(
+#     {},
+#     "RubyGems.org",
+#     "https://rubygems.org/api/v1/webauthn_verification/odow34b93t6aPCdY",
+#     { email: "email@example.com", password: "password" }
+#   )
+#   thread.join
+#   otp = thread[:otp]
+#   error = thread[:error]
+#
+
+module Gem::GemcutterUtilities
+  class WebauthnPoller
+    include Gem::GemcutterUtilities
+    TIMEOUT_IN_SECONDS = 300
+
+    attr_reader :options, :host
+
+    def initialize(options, host)
+      @options = options
+      @host = host
+    end
+
+    def self.poll_thread(options, host, webauthn_url, credentials)
+      thread = Thread.new do
+        Thread.current[:otp] = new(options, host).poll_for_otp(webauthn_url, credentials)
+      rescue Gem::WebauthnVerificationError, Timeout::Error => e
+        Thread.current[:error] = e
+      end
+      thread.abort_on_exception = true
+      thread.report_on_exception = false
+
+      thread
+    end
+
+    def poll_for_otp(webauthn_url, credentials)
+      Timeout.timeout(TIMEOUT_IN_SECONDS) do
+        loop do
+          response = webauthn_verification_poll_response(webauthn_url, credentials)
+          raise Gem::WebauthnVerificationError, response.message unless response.is_a?(Net::HTTPSuccess)
+
+          require "json"
+          parsed_response = JSON.parse(response.body)
+          case parsed_response["status"]
+          when "pending"
+            sleep 5
+          when "success"
+            return parsed_response["code"]
+          else
+            raise Gem::WebauthnVerificationError, parsed_response.fetch("message", "Invalid response from server")
+          end
+        end
+      end
+    end
+
+    private
+
+    def webauthn_verification_poll_response(webauthn_url, credentials)
+      webauthn_token = %r{(?<=\/)[^\/]+(?=$)}.match(webauthn_url)[0]
+      rubygems_api_request(:get, "api/v1/webauthn_verification/#{webauthn_token}/status.json") do |request|
+        if credentials.empty?
+          request.add_field "Authorization", api_key
+        else
+          request.basic_auth credentials[:email], credentials[:password]
+        end
+      end
+    end
+  end
+end
diff --git a/test/rubygems/test_webauthn_poller.rb b/test/rubygems/test_webauthn_poller.rb
new file mode 100644
index 0000000000..776deba9cc
--- /dev/null
+++ b/test/rubygems/test_webauthn_poller.rb
@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+require_relative "helper"
+require "rubygems/gemcutter_utilities/webauthn_poller"
+require "rubygems/gemcutter_utilities"
+
+class WebauthnPollerTest < Gem::TestCase
+  def setup
+    super
+
+    @host = Gem.host
+    @webauthn_url = "#{@host}/api/v1/webauthn_verification/odow34b93t6aPCdY"
+    @fetcher = Gem::FakeFetcher.new
+    Gem::RemoteFetcher.fetcher = @fetcher
+    @credentials = {
+      email: "email@example.com",
+      password: "password",
+    }
+  end
+
+  def test_poll_thread_success
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "{\"status\":\"success\",\"code\":\"Uvh6T57tkWuUnWYo\"}",
+      code: 200,
+      msg: "OK"
+    )
+
+    thread = Gem::GemcutterUtilities::WebauthnPoller.poll_thread({}, @host, @webauthn_url, @credentials)
+    thread.join
+
+    assert_equal thread[:otp], "Uvh6T57tkWuUnWYo"
+  end
+
+  def test_poll_thread_webauthn_verification_error
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "HTTP Basic: Access denied.",
+      code: 401,
+      msg: "Unauthorized"
+    )
+
+    thread = Gem::GemcutterUtilities::WebauthnPoller.poll_thread({}, @host, @webauthn_url, @credentials)
+    thread.join
+
+    assert_equal thread[:error].message, "Security device verification failed: Unauthorized"
+  end
+
+  def test_poll_thread_timeout_error
+    raise_error = ->(*_args) { raise Timeout::Error, "execution expired" }
+    Timeout.stub(:timeout, raise_error) do
+      thread = Gem::GemcutterUtilities::WebauthnPoller.poll_thread({}, @host, @webauthn_url, @credentials)
+      thread.join
+      assert_equal thread[:error].message, "execution expired"
+    end
+  end
+
+  def test_poll_for_otp_success
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "{\"status\":\"success\",\"code\":\"Uvh6T57tkWuUnWYo\"}",
+      code: 200,
+      msg: "OK"
+    )
+
+    otp = Gem::GemcutterUtilities::WebauthnPoller.new({}, @host).poll_for_otp(@webauthn_url, @credentials)
+
+    assert_equal otp, "Uvh6T57tkWuUnWYo"
+  end
+
+  def test_poll_for_otp_pending_sleeps
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "{\"status\":\"pending\",\"message\":\"Security device authentication is still pending.\"}",
+      code: 200,
+      msg: "OK"
+    )
+
+    assert_raises Timeout::Error do
+      Timeout.timeout(0.1) do
+        Gem::GemcutterUtilities::WebauthnPoller.new({}, @host).poll_for_otp(@webauthn_url, @credentials)
+      end
+    end
+  end
+
+  def test_poll_for_otp_not_http_success
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "HTTP Basic: Access denied.",
+      code: 401,
+      msg: "Unauthorized"
+    )
+
+    error = assert_raises Gem::WebauthnVerificationError do
+      Gem::GemcutterUtilities::WebauthnPoller.new({}, @host).poll_for_otp(@webauthn_url, @credentials)
+    end
+
+    assert_equal error.message, "Security device verification failed: Unauthorized"
+  end
+
+  def test_poll_for_otp_invalid_format
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "{}",
+      code: 200,
+      msg: "OK"
+    )
+
+    error = assert_raises Gem::WebauthnVerificationError do
+      Gem::GemcutterUtilities::WebauthnPoller.new({}, @host).poll_for_otp(@webauthn_url, @credentials)
+    end
+
+    assert_equal error.message, "Security device verification failed: Invalid response from server"
+  end
+
+  def test_poll_for_otp_invalid_status
+    @fetcher.data["#{@webauthn_url}/status.json"] = Gem::HTTPResponseFactory.create(
+      body: "{\"status\":\"expired\",\"message\":\"The token in the link you used has either expired or been used already.\"}",
+      code: 200,
+      msg: "OK"
+    )
+
+    error = assert_raises Gem::WebauthnVerificationError do
+      Gem::GemcutterUtilities::WebauthnPoller.new({}, @host).poll_for_otp(@webauthn_url, @credentials)
+    end
+
+    assert_equal error.message,
+      "Security device verification failed: The token in the link you used has either expired or been used already."
+  end
+end