lib/smtp_protocol.rb in rumbster-1.0.0 vs lib/smtp_protocol.rb in rumbster-1.0.1

@@ -17,21 +17,21 @@
     
     SmtpProtocol.new(initial_state, states)
   end
 
   def initialize(initial_state, states)
-    states.each_value { |state| state.protocol = self }
+    states[:read_mail].protocol = self unless states[:read_mail].nil?
 
     @states = states
-    @state = @states[initial_state]
   end
   
-  def state=(new_state)
-    @state = @states[new_state]
-  end
-  
   def serve(io)
-    @state.serve(io)
+    next_state = :init
+    count = 0
+    until next_state == :done or count > 10 do 
+      next_state = @states[next_state].serve(io)
+      count += 1
+    end
   end
   
   def new_message_received(message)
     changed
     notify_observers(message)