Refactor duplicate code

riemann · Jul 2, 2024 · 3a62c71 · 3a62c71
1 parent 0a2641b
commit 3a62c71
Showing 1 changed file with 14 additions and 27 deletions.
diff --git a/lib/riemann/tools/tls_check.rb b/lib/riemann/tools/tls_check.rb
@@ -494,19 +494,9 @@ def postgres_tls_socket(uri, address)
 
       def smtp_tls_socket(uri, address)
         socket = tcp_socket(address, uri.port)
-        while (line = socket.gets)
-          next if line =~ /\A220-/
-          break if line =~ /\A220 /
-
-          raise UnexpectedMessage, line
-        end
+        read_socket_lines_until_prefix_matched(socket, '220 ', also_accept_prefixes: ['220-'])
         socket.puts("EHLO #{my_hostname}")
-        while (line = socket.gets)
-          next if line =~ /\A250-/
-          break if line =~ /\A250 /
-
-          raise UnexpectedMessage, line
-        end
+        read_socket_lines_until_prefix_matched(socket, '250 ', also_accept_prefixes: ['250-'])
         socket.puts('STARTTLS')
         socket.gets
 
@@ -519,26 +509,23 @@ def my_hostname
         Socket.gethostname
       end
 
-      def imap_tls_socket(uri, address)
-        socket = tcp_socket(address, uri.port)
-        while (line = socket.gets) # rubocop:disable Lint/UnreachableLoop
-          break if line =~ /\A\* OK/
+      def read_socket_lines_until_prefix_matched(socket, prefix, also_accept_prefixes: [])
+        loop do
+          line = socket.gets
+          break if line.start_with?(prefix)
+          next if also_accept_prefixes.map { |accepted_prefix| line.start_with?(accepted_prefix) }.any?
 
           raise UnexpectedMessage, line
         end
-        socket.puts('. CAPABILITY')
-        while (line = socket.gets)
-          next if line =~ /\A\* CAPABILITY/
-          break if line =~ /\A\. OK/
+      end
 
-          raise UnexpectedMessage, line
-        end
+      def imap_tls_socket(uri, address)
+        socket = tcp_socket(address, uri.port)
+        read_socket_lines_until_prefix_matched(socket, '* OK')
+        socket.puts('. CAPABILITY')
+        read_socket_lines_until_prefix_matched(socket, '. OK', also_accept_prefixes: ['* CAPABILITY'])
         socket.puts('. STARTTLS')
-        while (line = socket.gets) # rubocop:disable Lint/UnreachableLoop
-          break if line =~ /\A\. OK/
-
-          raise UnexpectedMessage, line
-        end
+        read_socket_lines_until_prefix_matched(socket, '. OK')
 
         tls_handshake(socket, uri.host)
       end