Commit b5707788 authored by Avery Pennarun's avatar Avery Pennarun

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

This avoids memory/socket leaks.
parent 4c5185dc
import struct, socket, select, errno, re, signal import struct, socket, select, errno, re, signal, time
import compat.ssubprocess as ssubprocess import compat.ssubprocess as ssubprocess
import helpers, ssnet, ssh, ssyslog import helpers, ssnet, ssh, ssyslog
from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
...@@ -293,22 +293,27 @@ def _main(listener, fw, ssh_cmd, remotename, python, latency_control, ...@@ -293,22 +293,27 @@ def _main(listener, fw, ssh_cmd, remotename, python, latency_control,
handlers.append(Proxy(SockWrapper(sock, sock), outwrap)) handlers.append(Proxy(SockWrapper(sock, sock), outwrap))
handlers.append(Handler([listener], onaccept)) handlers.append(Handler([listener], onaccept))
dnspeers = {} dnsreqs = {}
def dns_done(chan, data): def dns_done(chan, data):
peer = dnspeers.get(chan) peer,timeout = dnsreqs.get(chan)
debug1('dns_done: channel=%r peer=%r\n' % (chan, peer)) debug3('dns_done: channel=%r peer=%r\n' % (chan, peer))
if peer: if peer:
del dnspeers[chan] del dnsreqs[chan]
debug1('doing sendto %r\n' % (peer,)) debug3('doing sendto %r\n' % (peer,))
dnslistener.sendto(data, peer) dnslistener.sendto(data, peer)
def ondns(): def ondns():
pkt,peer = dnslistener.recvfrom(4096) pkt,peer = dnslistener.recvfrom(4096)
now = time.time()
if pkt: 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() chan = mux.next_channel()
dnspeers[chan] = peer dnsreqs[chan] = peer,now+30
mux.send(chan, ssnet.CMD_DNS_REQ, pkt) mux.send(chan, ssnet.CMD_DNS_REQ, pkt)
mux.channels[chan] = lambda cmd,data: dns_done(chan,data) 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: if dnslistener:
handlers.append(Handler([dnslistener], ondns)) handlers.append(Handler([dnslistener], ondns))
......
import re, struct, socket, select, traceback import re, struct, socket, select, traceback, time
if not globals().get('skip_imports'): if not globals().get('skip_imports'):
import ssnet, helpers, hostwatch import ssnet, helpers, hostwatch
import compat.ssubprocess as ssubprocess import compat.ssubprocess as ssubprocess
...@@ -111,6 +111,7 @@ class DnsProxy(Handler): ...@@ -111,6 +111,7 @@ class DnsProxy(Handler):
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
Handler.__init__(self, [sock]) Handler.__init__(self, [sock])
self.sock = sock self.sock = sock
self.timeout = time.time()+30
self.mux = mux self.mux = mux
self.chan = chan self.chan = chan
self.sock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42) self.sock.setsockopt(socket.SOL_IP, socket.IP_TTL, 42)
...@@ -119,8 +120,9 @@ class DnsProxy(Handler): ...@@ -119,8 +120,9 @@ class DnsProxy(Handler):
def callback(self): def callback(self):
data = self.sock.recv(4096) 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.mux.send(self.chan, ssnet.CMD_DNS_RESPONSE, data)
self.ok = False
def main(): def main():
...@@ -184,7 +186,7 @@ def main(): ...@@ -184,7 +186,7 @@ def main():
dnshandlers = {} dnshandlers = {}
def dns_req(channel, data): def dns_req(channel, data):
debug1('got dns request!\n') debug2('Incoming DNS request.\n')
h = DnsProxy(mux, channel, data) h = DnsProxy(mux, channel, data)
handlers.append(h) handlers.append(h)
dnshandlers[channel] = h dnshandlers[channel] = h
...@@ -201,3 +203,10 @@ def main(): ...@@ -201,3 +203,10 @@ def main():
if latency_control: if latency_control:
mux.check_fullness() mux.check_fullness()
mux.callback() 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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment