🥅 Guard against tagged responses sent too soon

lib/net/imap.rb

@@ -3201,8 +3201,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)
 
@@ -3665,8 +3665,8 @@ def send_command(cmd, *args, &block)
         args.each do |i|
           validate_data(i)
         end
-        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(" ")
@@ -3687,6 +3687,16 @@ def send_command(cmd, *args, &block)
       raise
     end
 
+    def start_command(**)
+      raise IOError, "closed stream" if disconnected?
+      tag = generate_tag
+      command = Command.new(**, 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)

test/net/imap/test_imap.rb

@@ -210,13 +210,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
@@ -230,9 +233,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
@@ -250,7 +253,11 @@ 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_raise(Net::IMAP::InvalidTaggedResponseError) do cmd.(imap) end
+        assert_raise_with_message(
+          Net::IMAP::InvalidTaggedResponseError, / unstarted .*tag=RUBY0001/
+        ) do
+          cmd.(imap)
+        end
         assert imap.disconnected?
         assert_stream_closed_error do cmd.(imap) end
         assert_stream_closed_error do cmd.(imap) end