🥅 Guard against tagged responses sent too soon

lib/net/imap.rb

@@ -1479,16 +1479,19 @@ def starttls(**options)
       rescue Exception => error
         raise # note that the error backtrace is in the receiver_thread
       end
-      if error
-        disconnect
-        reraise error
-      end
-      unless handled
+      synchronize do
+        if error
+          reraise error
+        elsif !handled
+          raise InvalidResponseError,
+                "STARTTLS handler was bypassed, although server responded %p" % [
+                  ok.raw_data.chomp
+                ]
+        end
+      rescue Exception => error
+        @client_disconnect_error ||= error
         disconnect
-        raise InvalidResponseError,
-              "STARTTLS handler was bypassed, although server responded %p" % [
-                ok.raw_data.chomp
-              ]
+        raise
       end
       ok
     end
@@ -3249,8 +3252,8 @@ def idle(timeout = nil, &response_handler)
       response = nil
 
       synchronize do
-        tag = Thread.current[:net_imap_tag] = generate_tag
-        command = Command[tag:, name: "IDLE"]
+        command = start_command(name: "IDLE")
+        tag = Thread.current[:net_imap_tag] = command.tag # TODO: remove this
         put_string("#{tag} IDLE#{CRLF}")
         finish_sending_command(command)
 
@@ -3269,7 +3272,8 @@ def idle(timeout = nil, &response_handler)
             response = get_tagged_response(tag, "IDLE", idle_response_timeout)
           end
         end
-      rescue InvalidResponseError
+      rescue InvalidResponseError => error
+        @client_disconnect_error ||= error
         disconnect
         raise
       end
@@ -3683,8 +3687,8 @@ def send_command_with_continuations(cmd, *args)
     def send_command(cmd, *args, &block)
       args.each do validate_data _1 end
       synchronize do
-        tag = generate_tag
-        command = Command[tag:, name: cmd]
+        command = start_command(name: cmd)
+        tag = command.tag
         put_string(tag + " " + cmd)
         args.each do |i|
           put_string(" ")
@@ -3699,12 +3703,22 @@ def send_command(cmd, *args, &block)
         ensure
           remove_response_handler(block) if block
         end
-      rescue InvalidResponseError
+      rescue InvalidResponseError => error
+        @client_disconnect_error ||= error
         disconnect
         raise
       end
     end
 
+    def start_command(**)
+      raise IOError, "closed stream" if disconnected?
+      command = Command.new(**, tag: generate_tag)
+      if (response = @tagged_responses[command.tag])
+        raise InvalidTaggedResponseError.new(:unstarted, command:, response:)
+      end
+      command
+    end
+
     # NOTE: This must be synchronized with sending the command's final CRLF and
     # adding any command-related response handlers.
     def finish_sending_command(command)
@@ -3736,7 +3750,6 @@ def get_tagged_response(tag, cmd, timeout = nil)
       when /\A(?:BAD)\z/ni
         raise BadResponseError, resp
       else
-        disconnect
         raise InvalidResponseError, "invalid tagged resp: %p" % [resp.raw_data.chomp]
       end
     end

test/net/imap/test_imap.rb

@@ -220,13 +220,16 @@ def test_starttls_stripping_ok_sent_before_response
     end
   end
 
-  # Similar to STARTTLS stripping test, but checks other commands too
+  # Similar to STARTTLS stripping test, but checks other commands and statuses
   data(
-    "IDLE"   => ->imap do imap.idle(1) do end end,
-    "NOOP"   => ->imap do imap.noop end,
-    "SELECT" => ->imap do imap.select("inbox") end,
+    "OK for IDLE"    => ["OK",  ->imap do imap.idle(1) do end end],
+    "OK for NOOP"    => ["OK",  ->imap do imap.noop end],
+    "NO for IDLE"    => ["NO",  ->imap do imap.idle(1) do end end],
+    "NO for EXAMINE" => ["NO",  ->imap do imap.examine("inbox") end],
+    "BAD for IDLE"   => ["BAD", ->imap do imap.idle(1) do end end],
+    "BAD for SELECT" => ["BAD", ->imap do imap.select("inbox") end],
   )
-  test "premature tagged OK response" do |cmd|
+  test "premature tagged response" do |(status, cmd)|
     timeout = 5
     timeout *= EnvUtil.timeout_scale || 1 if defined?(EnvUtil.timeout_scale)
     Timeout.timeout(timeout) do
@@ -240,9 +243,9 @@ def test_starttls_stripping_ok_sent_before_response
         begin
           sock.print("* OK test server\r\n")
           assert_equal :send_malicious_responses, client_to_server.pop
-          sock.print("RUBY0001 OK invalid\r\n")
-          sock.print("RUBY0002 OK false\r\n")
-          sock.print("RUBY0003 OK tricky\r\n")
+          sock.print("RUBY0001 #{status} invalid\r\n")
+          sock.print("RUBY0002 #{status} false\r\n")
+          sock.print("RUBY0003 #{status} tricky\r\n")
           server_to_client << :sent_malicious_responses
           sock.gets
         ensure
@@ -260,7 +263,9 @@ def test_starttls_stripping_ok_sent_before_response
         assert_equal :sent_malicious_responses, server_to_client.pop
         assert_equal [1, 2, 3], 3.times.map { rcvr_to_client.pop }
         # should respond this way for _any_ command
-        assert_local_raise(Net::IMAP::InvalidTaggedResponseError) do
+        assert_local_raise(
+          Net::IMAP::InvalidTaggedResponseError, / unstarted .*tag=RUBY0001/
+        ) do
           cmd.(imap)
         end
         assert imap.disconnected?