From 175356470edf3e8e8a46ce7d7e56ac0b162713c0 Mon Sep 17 00:00:00 2001
From: nick evans <nick@rubinick.dev>
Date: Sun, 30 Jun 2024 14:49:31 -0400
Subject: [PATCH] =?UTF-8?q?=F0=9F=A5=85=20Cancel=20AUTHENTICATE=20if=20pro?=
 =?UTF-8?q?cess=20raises=20an=20error=20[=F0=9F=9A=A7=20server=20error]?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The exception will be re-raised after the protocol cancel response has
been sent.
---
 lib/net/imap/sasl/authentication_exchange.rb | 21 ++++++++++++++++----
 test/net/imap/fake_server/command_router.rb  |  5 +++--
 test/net/imap/test_imap.rb                   | 19 ++++++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/lib/net/imap/sasl/authentication_exchange.rb b/lib/net/imap/sasl/authentication_exchange.rb
index 5fb48b18..3b628476 100644
--- a/lib/net/imap/sasl/authentication_exchange.rb
+++ b/lib/net/imap/sasl/authentication_exchange.rb
@@ -7,8 +7,6 @@ module SASL
       # AuthenticationExchange is used internally by Net::IMAP#authenticate.
       # But the API is still *experimental*, and may change.
       #
-      # TODO: catch exceptions in #process and send #cancel_response.
-      # TODO: raise an error if the command succeeds after being canceled.
       # TODO: use with more clients, to verify the API can accommodate them.
       # TODO: pass ClientAdapter#service to SASL.authenticator
       #
@@ -81,6 +79,9 @@ def self.build(client, mechanism, *args, sasl_ir: true, **kwargs, &block)
 
         attr_reader :mechanism, :authenticator
 
+        # An exception that has been raised by <tt>authenticator.process</tt>.
+        attr_reader :process_error
+
         def initialize(client, mechanism, authenticator, sasl_ir: true)
           @client = client
           @mechanism = Authenticators.normalize_name(mechanism)
@@ -93,8 +94,17 @@ def initialize(client, mechanism, authenticator, sasl_ir: true)
         # using #authenticator.  Authentication failures will raise an
         # exception.  Any exceptions other than AuthenticationCanceled or those
         # in <tt>client.response_errors</tt> will drop the connection.
+        #
+        # When <tt>authenticator.process</tt> raises any StandardError
+        # (including AuthenticationCanceled), the authentication exchange will
+        # be canceled before re-raising the exception.  The server will usually
+        # respond with an error response, and the client will most likely raise
+        # that error.  This client error will supercede the original error.
+        # Unfortunately, the original error will not be the +#cause+ for the
+        # client error.  But it will be available on #process_error.
         def authenticate
           client.run_command(mechanism, initial_response) { process _1 }
+            .tap { raise process_error if process_error }
             .tap { raise AuthenticationIncomplete, _1 unless done? }
         rescue AuthenticationCanceled, *client.response_errors
           raise # but don't drop the connection
@@ -128,9 +138,12 @@ def initial_response
         end
 
         def process(challenge)
-          client.encode authenticator.process client.decode challenge
-        ensure
           @processed = true
+          return client.cancel_response if process_error
+          client.encode authenticator.process client.decode challenge
+        rescue => process_error
+          @process_error = process_error
+          client.cancel_response
         end
 
       end
diff --git a/test/net/imap/fake_server/command_router.rb b/test/net/imap/fake_server/command_router.rb
index 82430316..f170d8d9 100644
--- a/test/net/imap/fake_server/command_router.rb
+++ b/test/net/imap/fake_server/command_router.rb
@@ -97,8 +97,9 @@ def handler_for(command)
         response_b64 = resp.request_continuation("") || ""
         state.commands << {continuation: response_b64}
       end
-      response = Base64.decode64(response_b64)
-      response.empty?                   and return resp.fail_bad "canceled"
+      response_b64.strip == ?*          and return resp.fail_bad "canceled"
+      response = Base64.decode64(response_b64) rescue :decode64_failed
+      response == :decode64_failed      and return resp.fail_bad "invalid b64"
       # TODO: support mechanisms other than PLAIN.
       parts = response.split("\0")
       parts.length == 3                  or return resp.fail_bad "invalid"
diff --git a/test/net/imap/test_imap.rb b/test/net/imap/test_imap.rb
index 0ad7ca51..2e4ce84a 100644
--- a/test/net/imap/test_imap.rb
+++ b/test/net/imap/test_imap.rb
@@ -1014,6 +1014,25 @@ def test_id
     end
   end
 
+  test "#authenticate can be canceled with SASL::AuthenticationCanceled" do
+    with_fake_server(
+      preauth: false, cleartext_auth: true,
+      sasl_ir: false, sasl_mechanisms: %i[PLAIN]
+    ) do |server, imap|
+      registry = Net::IMAP::SASL::Authenticators.new(use_defaults: false)
+      registry.add_authenticator :plain, ->(*a, **kw, &b) {
+        ->(challenge) {
+          raise(Net::IMAP::SASL::AuthenticationCanceled,
+                "a: %p, kw: %p, b: %p" % [a, kw, b])
+        }
+      }
+      assert_raise_with_message(Net::IMAP::BadResponseError, "canceled") do
+        imap.authenticate(:plain, hello: :world, registry: registry)
+      end
+      refute imap.disconnected?
+    end
+  end
+
   def test_uidplus_uid_expunge
     with_fake_server(select: "INBOX",
                      extensions: %i[UIDPLUS]) do |server, imap|