LDAPClient.unbind returns a Deferred that fires on actual disconnect (fixes #225)

docs/source/NEWS.rst

@@ -82,6 +82,11 @@ Fixes
   in-memory back end grew a ``deleteOldRDN`` kwarg on ``move`` and the
   client-side ``ldapsyntax.LDAPEntry.move`` grew the same flag so
   callers can request either behaviour (#89).
+- ``LDAPClient.unbind`` now returns a ``Deferred`` that fires with the
+  connection-lost reason once the transport has actually been closed,
+  which matters for TLS teardown. A new
+  ``LDAPClient.notifyOnDisconnect()`` helper lets callers register
+  additional disconnect Deferreds without initiating an unbind (#225).
 
 
 21.2.0 (2021-02-28)

ldaptor/protocols/ldap/ldapclient.py

@@ -39,6 +39,7 @@ def __init__(self):
         self.onwire = {}
         self.buffer = b""
         self.connected = None
+        self._disconnect_deferreds = []
 
     berdecoder = pureldap.LDAPBERDecoderContext_TopLevel(
         inherit=pureldap.LDAPBERDecoderContext_LDAPMessage(
@@ -75,6 +76,21 @@ def connectionLost(self, reason=protocol.connectionDone):
             k, v = self.onwire.popitem()
             d, _, _, _, _ = v
             d.errback(reason)
+        # fire any notifyOnDisconnect / unbind Deferreds
+        pending = self._disconnect_deferreds
+        self._disconnect_deferreds = []
+        for d in pending:
+            d.callback(reason)
+
+    def notifyOnDisconnect(self):
+        """
+        Return a Deferred that fires with the disconnect reason once the
+        transport has actually been closed. Useful for waiting on a clean
+        TLS shutdown after unbind() (#225).
+        """
+        d = defer.Deferred()
+        self._disconnect_deferreds.append(d)
+        return d
 
     def _send(self, op, controls=None):
         if not self.connected:
@@ -235,11 +251,21 @@ def _handle_bind_msg(self, msg):
         return (msg.matchedDN, msg.serverSaslCreds)
 
     def unbind(self):
+        """
+        Send an LDAP unbind and close the connection.
+
+        Returns a Deferred that fires with the connection-lost reason
+        once the transport has actually closed. Callers that don't need
+        to wait can ignore the return value; callers that do (e.g.
+        cleanup after a TLS shutdown) can chain on it.
+        """
         if not self.connected:
             raise Exception("Not connected (TODO)")  # TODO make this a real object
         r = pureldap.LDAPUnbindRequest()
         self.send_noResponse(r)
+        d = self.notifyOnDisconnect()
         self.transport.loseConnection()
+        return d
 
     def _cbStartTLS(self, msg, ctx):
         assert isinstance(msg, pureldap.LDAPExtendedResponse)

ldaptor/test/test_ldapclient.py

@@ -130,6 +130,34 @@ def test_unbind_not_connected(self):
         client = ldapclient.LDAPClient()
         self.assertRaises(Exception, client.unbind)
 
+    def test_unbind_returnsDeferredFiringOnDisconnect(self):
+        """
+        unbind() returns a Deferred that fires once the transport is
+        actually closed. Previously the caller had no way to await the
+        (potentially slow, e.g. TLS) shutdown (#225).
+        """
+        client, transport = self.create_test_client()
+        d = client.unbind()
+        fired = []
+        d.addBoth(fired.append)
+        self.assertEqual(fired, [], "unbind Deferred fired before disconnect")
+        client.connectionLost(SillyError("boom"))
+        self.assertEqual(len(fired), 1)
+
+    def test_notifyOnDisconnect_firesEachTime(self):
+        """
+        Multiple notifyOnDisconnect() Deferreds all fire on the same
+        connectionLost event.
+        """
+        client, transport = self.create_test_client()
+        d1 = client.notifyOnDisconnect()
+        d2 = client.notifyOnDisconnect()
+        seen = []
+        d1.addBoth(lambda r: seen.append(("d1", r)))
+        d2.addBoth(lambda r: seen.append(("d2", r)))
+        client.connectionLost(SillyError("bye"))
+        self.assertEqual(len(seen), 2)
+
     def test_TLS_failure(self):
         clock = Clock()
         ldapclient.reactor = clock