🚧 Handle cancellations more carefully...

Author: nevans

lib/net/imap/sasl.rb

@@ -103,7 +103,15 @@ module SASL
 
       # Indicates an authentication exchange that will be or has been canceled
       # by the client, not due to any error or failure during processing.
-      AuthenticationCanceled = Class.new(Error)
+      class AuthenticationCanceled < Error
+        # The error response from the server
+        attr_reader :response
+
+        def initialize(message = "authentication canceled", response: nil)
+          super(message)
+          @response = response
+        end
+      end
 
       # Indicates an error when processing a server challenge, e.g: an invalid
       # or unparsable challenge.  An underlying exception may be available as

lib/net/imap/sasl/authentication_exchange.rb

@@ -82,6 +82,9 @@ def self.build(client, mechanism, *args, sasl_ir: true, **kwargs, &block)
         # An exception that has been raised by <tt>authenticator.process</tt>.
         attr_reader :process_error
 
+        # An exception that represents an error response from the server.
+        attr_reader :response_error
+
         def initialize(client, mechanism, authenticator, sasl_ir: true)
           @client = client
           @mechanism = Authenticators.normalize_name(mechanism)
@@ -103,9 +106,11 @@ def initialize(client, mechanism, authenticator, sasl_ir: true)
         # 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? }
+          handle_cancellation do
+            client.run_command(mechanism, initial_response) { process _1 }
+              .tap { raise process_error if process_error }
+              .tap { raise AuthenticationIncomplete, _1 unless done? }
+          end
         rescue AuthenticationCanceled, *client.response_errors
           raise # but don't drop the connection
         rescue
@@ -141,11 +146,53 @@ def process(challenge)
           @processed = true
           return client.cancel_response if process_error
           client.encode authenticator.process client.decode challenge
-        rescue => process_error
-          @process_error = process_error
+        rescue AuthenticationCanceled => error
+          @process_error = error
+          client.cancel_response
+        rescue => error
+          @process_error = begin
+            raise AuthenticationError, "error while processing server challenge"
+          rescue
+            $!
+          end
           client.cancel_response
         end
 
+        # | process | response | => result                                |
+        # |---------|----------|------------------------------------------|
+        # | success | success  | success                                  |
+        # | success | error    | reraise response error                   |
+        # | error   | success  | raise incomplete error (cause = process) |
+        # | error   | error    | raise canceled error   (cause = process) |
+        def handle_cancellation
+          result = begin
+            yield
+          rescue *client.response_errors => error
+            @response_error = error
+            raise unless process_error
+          end
+          raise_mutual_cancellation!       if process_error &&  response_error
+          raise_incomplete_cancel!(result) if process_error && !response_error
+          result
+        end
+
+        def raise_mutual_cancellation!
+          raise process_error # sets the cause
+        rescue
+          raise AuthenticationCanceled.new(
+            "authentication canceled (see error #cause and #response)",
+            response: response_error
+          )
+        end
+
+        def raise_incomplete_cancellation!
+          raise process_error # sets the cause
+        rescue
+          raise AuthenticationIncomplete.new(
+            response_error, "server ignored canceled authentication"
+          )
+        end
+
       end
     end
   end

test/net/imap/test_imap_authenticate.rb

@@ -166,20 +166,32 @@ class IMAPAuthenticateTest < Net::IMAP::TestCase
   end
 
   test "#authenticate can be canceled with SASL::AuthenticationCanceled" do
+    authenticator = Class.new do
+      def initialize(*a, **kw, &b) @a, @kw, @b = a, kw, b end
+      def process(challenge)
+        raise(Net::IMAP::SASL::AuthenticationCanceled,
+              "a: %p, kw: %p, b: %p, c: %p" % [@a, @kw, @b, challenge])
+      end
+    end
+    registry = Net::IMAP::SASL::Authenticators.new(use_defaults: false)
+    registry.add_authenticator :plain, authenticator
+
     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)
+      error = nil
+      assert_raise_with_message(Net::IMAP::SASL::AuthenticationCanceled,
+                                /authentication canceled/i) do
+        imap.authenticate(:plain, foo: :bar, registry: registry)
+      rescue => error
+        raise # for assert_raise
       end
+      assert_kind_of Net::IMAP::SASL::AuthenticationCanceled,  error.cause
+      assert_equal("a: %p, kw: %p, b: %p, c: %p" % [[], {foo: :bar}, nil, ""],
+                   error.cause.to_s)
+      assert_kind_of Net::IMAP::BadResponseError, error.response
+      assert_equal   "canceled",                  error.response.to_s
       refute imap.disconnected?
     end
   end