dns: trim DNS channel handlers after a response, or after a timeout.

This avoids memory/socket leaks.

2 files changed, 25 insertions, 11 deletions

diff --git a/client.py b/client.py
index 7872bae..0f596eb 100644
--- a/client.py
+++ b/client.py
@@ -1,4 +1,4 @@
-import struct, socket, select, errno, re, signal
+import struct, socket, select, errno, re, signal, time
 import compat.ssubprocess as ssubprocess
 import helpers, ssnet, ssh, ssyslog
 from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
@@ -293,22 +293,27 @@ def _main(listener, fw, ssh_cmd, remotename, python, latency_control,
         handlers.append(Proxy(SockWrapper(sock, sock), outwrap))
     handlers.append(Handler([listener], onaccept))
 
-    dnspeers = {}
+    dnsreqs = {}
     def dns_done(chan, data):
-        peer = dnspeers.get(chan)
-        debug1('dns_done: channel=%r peer=%r\n' % (chan, peer))
+        peer,timeout = dnsreqs.get(chan)
+        debug3('dns_done: channel=%r peer=%r\n' % (chan, peer))
         if peer:
-            del dnspeers[chan]
-            debug1('doing sendto %r\n' % (peer,))
+            del dnsreqs[chan]
+            debug3('doing sendto %r\n' % (peer,))
             dnslistener.sendto(data, peer)
     def ondns():
         pkt,peer = dnslistener.recvfrom(4096)
+        now = time.time()
         if pkt:
-            debug1('Got DNS request from %r: %d bytes\n' % (peer, len(pkt)))
+            debug1('DNS request from %r: %d bytes\n' % (peer, len(pkt)))
             chan = mux.next_channel()
-            dnspeers[chan] = peer
+            dnsreqs[chan] = peer,now+30
             mux.send(chan, ssnet.CMD_DNS_REQ, pkt)
             mux.channels[chan] = lambda cmd,data: dns_done(chan,data)
+        for chan,(peer,timeout) in dnsreqs.items():
+            if timeout < now:
+                del dnsreqs[chan]
+        debug3('Remaining DNS requests: %d\n' % len(dnsreqs))
     if dnslistener:
         handlers.append(Handler([dnslistener], ondns))
 
diff --git a/server.py b/server.py
index 60eaa42..45bc2fc 100644
--- a/server.py
+++ b/server.py
@@ -1,4 +1,4 @@
-import re, struct, socket, select, traceback
+import re, struct, socket, select, traceback, time
 if not globals().get('skip_imports'):
     import ssnet, helpers, hostwatch
     import compat.ssubprocess as ssubprocess
@@ -111,6 +111,7 @@ class DnsProxy(Handler):
         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         Handler.__init__(self, [sock])
         self.sock = sock
+        self.timeout = time.time()+30
         self.mux = mux
         self.chan = chan
         self.sock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42)
@@ -119,8 +120,9 @@ class DnsProxy(Handler):
 
     def callback(self):
         data = self.sock.recv(4096)
-        debug2('dns response: %d bytes\n' % len(data))
+        debug2('DNS response: %d bytes\n' % len(data))
         self.mux.send(self.chan, ssnet.CMD_DNS_RESPONSE, data)
+        self.ok = False
 
 
 def main():
@@ -184,7 +186,7 @@ def main():
 
     dnshandlers = {}
     def dns_req(channel, data):
-        debug1('got dns request!\n')
+        debug2('Incoming DNS request.\n')
         h = DnsProxy(mux, channel, data)
         handlers.append(h)
         dnshandlers[channel] = h
@@ -201,3 +203,10 @@ def main():
         if latency_control:
             mux.check_fullness()
         mux.callback()
+
+        if dnshandlers:
+            now = time.time()
+            for channel,h in dnshandlers.items():
+                if h.timeout < now or not h.ok:
+                    del dnshandlers[channel]
+                    h.ok = False