1 files changed, 31 insertions, 0 deletions
diff --git a/src/git_tftpd/protocol.py b/src/git_tftpd/protocol.py
new file mode 100644
index 0000000..ce270c7
--- /dev/null
+++ b/src/git_tftpd/protocol.py
@@ -0,0 +1,31 @@
+from tftp.protocol import TFTP
+from twisted.internet.defer import inlineCallbacks, returnValue
+
+
+class TFTPProtocol(TFTP):
+    """Simple hack to stop sending packets, if the connection was closed."""
+
+    def __init__(self, backend, _clock=None):
+        TFTP.__init__(self, backend, _clock)
+        self._real_session_sendData = None
+
+    def _session_sendData(self, session, bytes):
+        if session.transport.connected:
+            # Send data, if "connection" is established
+            self._real_session_sendData(bytes)
+        elif session.timeout_watchdog is not None and \
+             session.timeout_watchdog.active():
+            # Kill the timeout watchdog, so that we do not get:
+            # "Timed out after a successful transfer"
+            session.timeout_watchdog.cancel()
+    
+    @inlineCallbacks
+    def _startSession(self, datagram, addr, mode):
+        session = yield TFTP._startSession(self, datagram, addr, mode)
+
+        # monkey patch the sendData method
+        self._real_session_sendData = session.session.sendData
+        session.session.sendData = \
+            lambda bytes: self._session_sendData(session.session, bytes)
+
+        returnValue(session)