Commit b01d891

@@ -9,7 +9,7 @@ module Minbox
       @socket = socket
     end
 
-    def mail_message
+    def mail_message(&block)
       socket.puts "220"
       while socket && (line = socket.gets)
         case line
@@ -27,7 +27,7 @@ module Minbox
           socket.puts('502 Invalid/unsupported command')
         end
       end
-      Mail.new(@body.join)
+      block.call(Mail.new(@body.join))
     end
 
     private

@@ -8,21 +8,21 @@ module Minbox
       @logger = logger
     end
 
-    def listen!
+    def listen!(&block)
       logger.debug("Starting server on port #{port}...")
       @server = TCPServer.new(port.to_i)
       logger.debug("Server started!")
 
       loop do
-        yield handle(@server.accept)
+        handle(@server.accept, &block)
       rescue StandardError => error
         logger.error(error)
       end
     end
 
-    def handle(socket)
+    def handle(socket, &block)
       logger.debug("client connected: #{socket.inspect}")
-      Client.new(host, socket, logger).mail_message
+      Client.new(host, socket, logger).mail_message(&block)
     end
 
     def shutdown!