diff --git a/src/wormhole/blocking/transcribe.py b/src/wormhole/blocking/transcribe.py index 82be96c..23ad5b0 100644 --- a/src/wormhole/blocking/transcribe.py +++ b/src/wormhole/blocking/transcribe.py @@ -10,6 +10,7 @@ from .. import __version__ from .. import codes from ..errors import ServerError, Timeout, WrongPasswordError, UsageError from ..util.hkdf import HKDF +from ..channel_monitor import monitor SECOND = 1 MINUTE = 60*SECOND @@ -141,7 +142,6 @@ class Wormhole: self.verifier = None self._sent_data = False self._got_data = False - self._closed = False def handle_welcome(self, welcome): if ("motd" in welcome and @@ -194,6 +194,7 @@ class Wormhole: self.code = code channelid = int(mo.group(1)) self.channel = self._channel_manager.connect(channelid) + monitor.add(self.channel) def _start(self): # allocate the rest now too, so it can be serialized @@ -263,9 +264,5 @@ class Wormhole: raise WrongPasswordError def close(self): + monitor.close(self.channel) self.channel.deallocate() - self._closed = True - - def __del__(self): - if not self._closed: - print("Error: a Wormhole instance was not closed", file=sys.stderr) diff --git a/src/wormhole/channel_monitor.py b/src/wormhole/channel_monitor.py new file mode 100644 index 0000000..964b2a5 --- /dev/null +++ b/src/wormhole/channel_monitor.py @@ -0,0 +1,16 @@ +from __future__ import print_function +import sys +from weakref import ref + +class ChannelMonitor: + def __init__(self): + self._open_channels = set() + def add(self, w): + wr = ref(w, self._lost) + self._open_channels.add(wr) + def _lost(self, wr): + print("Error: a Wormhole instance was not closed", file=sys.stderr) + def close(self, w): + self._open_channels.discard(ref(w)) + +monitor = ChannelMonitor() # singleton diff --git a/src/wormhole/twisted/transcribe.py b/src/wormhole/twisted/transcribe.py index 742999d..89ea7ca 100644 --- a/src/wormhole/twisted/transcribe.py +++ b/src/wormhole/twisted/transcribe.py @@ -15,6 +15,7 @@ from .. import __version__ from .. import codes from ..errors import ServerError, WrongPasswordError, UsageError from ..util.hkdf import HKDF +from ..channel_monitor import monitor @implementer(IBodyProducer) class DataProducer: @@ -152,7 +153,6 @@ class Wormhole: self._started_get_code = False self._sent_data = False self._got_data = False - self._closed = False def _set_side(self, side): self._side = side @@ -209,6 +209,7 @@ class Wormhole: self.code = code channelid = int(mo.group(1)) self.channel = self._channel_manager.connect(channelid) + monitor.add(self.channel) def _start(self): # allocate the rest now too, so it can be serialized @@ -325,13 +326,6 @@ class Wormhole: return d def close(self, res=None): + monitor.close(self.channel) d = self.channel.deallocate() - def _closed(_): - self._closed = True - return res - d.addCallback(_closed) return d - - def __del__(self): - if not self._closed: - print("Error: a Wormhole instance was not closed", file=sys.stderr)