# HG changeset patch
# User cmlenz
# Date 1119966709 0
# Node ID d108f780b76c6ef9c1a83a21f4ae8f4c61759dfa
# Parent  47ab019508dd758448831aaf4842cba059b1846a
Fixes to the BEEP implementation:
 * Handle 'ERR' and 'NUL' messages as acknowledgements to sent messages (i.e. dealloc the message number).
 * When the peer disconnects without having closing any channels, close the channels automatically.

diff --git a/bitten/util/beep.py b/bitten/util/beep.py
--- a/bitten/util/beep.py
+++ b/bitten/util/beep.py
@@ -180,6 +180,14 @@
             logging.info('Session terminated')
         asynchat.async_chat.close(self)
 
+    def handle_close(self):
+        logging.warning('Peer %s:%s closed connection' % self.addr)
+        channels = self.channels.keys()
+        channels.reverse()
+        for channelno in channels:
+            self.channels[channelno].close()
+        asynchat.async_chat.handle_close(self)
+
     def handle_error(self):
         """Called by asyncore when an exception is raised."""
         cls, value = sys.exc_info()[:2]
@@ -435,7 +443,7 @@
             # Recombine queued messages
             payload = ''.join(self.inqueue[msgno]) + payload
             del self.inqueue[msgno]
-        if cmd == 'RPY' and msgno in self.msgnos:
+        if cmd in ('ERR', 'RPY', 'NUL') and msgno in self.msgnos:
             # Final reply using this message number, so dealloc
             self.msgnos.remove(msgno)
         message = None
diff --git a/bitten/util/tests/beep.py b/bitten/util/tests/beep.py
--- a/bitten/util/tests/beep.py
+++ b/bitten/util/tests/beep.py
@@ -41,11 +41,22 @@
 
     def handle_msg(self, msgno, message):
         text = message.as_string().strip()
-        self.handled_messages.append(('MSG', msgno, text))
+        self.handled_messages.append(('MSG', msgno, text, None))
 
     def handle_rpy(self, msgno, message):
         text = message.as_string().strip()
-        self.handled_messages.append(('RPY', msgno, text))
+        self.handled_messages.append(('RPY', msgno, text, None))
+
+    def handle_err(self, msgno, message):
+        text = message.as_string().strip()
+        self.handled_messages.append(('ERR', msgno, text, None))
+
+    def handle_ans(self, msgno, ansno, message):
+        text = message.as_string().strip()
+        self.handled_messages.append(('ANS', msgno, text, ansno))
+
+    def handle_nul(self, msgno):
+        self.handled_messages.append(('NUL', msgno, '', None))
 
 
 class ChannelTestCase(unittest.TestCase):
@@ -60,7 +71,7 @@
         """
         channel = beep.Channel(self.session, 0, MockProfileHandler)
         channel.handle_data_frame('MSG', 0, False, 0, None, 'foo bar')
-        self.assertEqual(('MSG', 0, 'foo bar'),
+        self.assertEqual(('MSG', 0, 'foo bar', None),
                          channel.profile.handled_messages[0])
 
     def test_handle_segmented_msg_frames(self):
@@ -71,7 +82,7 @@
         channel = beep.Channel(self.session, 0, MockProfileHandler)
         channel.handle_data_frame('MSG', 0, True, 0, None, 'foo ')
         channel.handle_data_frame('MSG', 0, False, 4, None, 'bar')
-        self.assertEqual(('MSG', 0, 'foo bar'),
+        self.assertEqual(('MSG', 0, 'foo bar', None),
                          channel.profile.handled_messages[0])
 
     def test_handle_out_of_sync_frame(self):
@@ -137,8 +148,8 @@
 
     def test_message_and_reply(self):
         """
-        Verify that a message number is deallocated after a final reply has been
-        received.
+        Verify that a message number is deallocated after a final "RPY" reply
+        has been received.
         """
         channel = beep.Channel(self.session, 0, MockProfileHandler)
         msgno = channel.send_msg(beep.MIMEMessage('foo bar', None))
@@ -146,10 +157,44 @@
                          self.session.sent_messages[0])
         assert msgno in channel.msgnos
         channel.handle_data_frame('RPY', msgno, False, 0, None, '42')
-        self.assertEqual(('RPY', msgno, '42'),
+        self.assertEqual(('RPY', msgno, '42', None),
                          channel.profile.handled_messages[0])
         assert msgno not in channel.msgnos
 
+    def test_message_and_error(self):
+        """
+        Verify that a message number is deallocated after a final "ERR" reply
+        has been received.
+        """
+        channel = beep.Channel(self.session, 0, MockProfileHandler)
+        msgno = channel.send_msg(beep.MIMEMessage('foo bar', None))
+        self.assertEqual(('MSG', 0, msgno, False, 0L, None, 'foo bar'),
+                         self.session.sent_messages[0])
+        assert msgno in channel.msgnos
+        channel.handle_data_frame('ERR', msgno, False, 0, None, '42')
+        self.assertEqual(('ERR', msgno, '42', None),
+                         channel.profile.handled_messages[0])
+        assert msgno not in channel.msgnos
+
+    def test_message_and_ans_nul(self):
+        """
+        Verify that a message number is deallocated after a final "NUL" reply
+        has been received.
+        """
+        channel = beep.Channel(self.session, 0, MockProfileHandler)
+        msgno = channel.send_msg(beep.MIMEMessage('foo bar', None))
+        self.assertEqual(('MSG', 0, msgno, False, 0L, None, 'foo bar'),
+                         self.session.sent_messages[0])
+        assert msgno in channel.msgnos
+        channel.handle_data_frame('ANS', msgno, False, 0, 0, '42')
+        self.assertEqual(('ANS', msgno, '42', 0),
+                         channel.profile.handled_messages[0])
+        assert msgno in channel.msgnos
+        channel.handle_data_frame('NUL', msgno, False, 2, None, '42')
+        self.assertEqual(('NUL', msgno, '', None),
+                         channel.profile.handled_messages[1])
+        assert msgno not in channel.msgnos
+
     def test_send_error(self):
         """
         Verify that sending an ERR message is processed correctly.