🥅 Fix premature tagged response guard for IDLE

lib/net/imap.rb

@@ -3202,8 +3202,8 @@ def idle(timeout = nil, &response_handler)
       synchronize do
         tag = Thread.current[:net_imap_tag] = generate_tag
         command = Command[tag:, name: "IDLE"]
-        finish_sending_command(command)
         put_string("#{tag} IDLE#{CRLF}")
+        finish_sending_command(command)
 
         begin
           add_response_handler(&response_handler)
@@ -3220,6 +3220,9 @@ def idle(timeout = nil, &response_handler)
             response = get_tagged_response(tag, "IDLE", idle_response_timeout)
           end
         end
+      rescue InvalidResponseError
+        disconnect
+        raise
       end
 
       return response

test/lib/helper.rb

@@ -184,6 +184,12 @@ def assert_pattern
     end
   end
 
+  def assert_stream_closed_error
+    assert_raise_with_message(IOError, /\A(?:stream closed|closed stream)\z/) do
+      yield
+    end
+  end
+
   def pend_if(condition, *args, &block)
     if condition
       pend(*args, &block)

test/net/imap/test_imap.rb

@@ -210,6 +210,57 @@ def test_starttls_stripping_ok_sent_before_response
     end
   end
 
+  # Similar to STARTTLS stripping test, but checks other commands too
+  data(
+    "IDLE"   => ->imap do imap.idle(1) do end end,
+    "NOOP"   => ->imap do imap.noop end,
+    "SELECT" => ->imap do imap.select("inbox") end,
+  )
+  test "premature tagged OK response" do |cmd|
+    timeout = 5
+    timeout *= EnvUtil.timeout_scale || 1 if defined?(EnvUtil.timeout_scale)
+    Timeout.timeout(timeout) do
+      server_to_client = Queue.new
+      client_to_server = Queue.new
+      rcvr_to_client   = Queue.new
+      server = create_tcp_server
+      port = server.addr[1]
+      start_server do
+        sock = server.accept
+        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")
+          server_to_client << :sent_malicious_responses
+          sock.gets
+        ensure
+          sock.close
+          server.close
+        end
+      end
+      begin
+        imap = Net::IMAP.new(server_addr, port:)
+        i = 0
+        imap.add_response_handler do |resp|
+          rcvr_to_client << (i += 1)
+        end
+        client_to_server << :send_malicious_responses
+        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 imap.disconnected?
+        assert_stream_closed_error do cmd.(imap) end
+        assert_stream_closed_error do cmd.(imap) end
+        assert_stream_closed_error do cmd.(imap) end
+      ensure
+        imap.disconnect if imap
+      end
+    end
+  end
+
   def start_server
     th = Thread.new do
       yield