[Shtoom] STUN retransmit delayed calls not cleaned up

[Antoine Pitrou]

Hi,

(is shtoom still active? "svn up" has not given changes for a long time)

I've found a potential problem with the STUN implementation.
The delayed calls which handle message restranmission are not cleaned up
at the end. This means if you contact several servers at once, and one
of them does not answer readily (or the DNS resolution takes some time),
the STUN protocol handler will still try to contact it even after it has
finished with another server.

It seems to cause problems with Twisted 2.0 in that the "self.transport"
in a DatagramProtocol is reset to None when we stopListening(). It leads
to errors like "NoneType object has no attribute 'write'" when the
_StunBase tries to send a request.

Here is the correction I've done to stun.py (modifications inside
<AP> ... </AP>):


class  _StunBase(object):

    def sendRequest(self, server, tid=None, avpairs=()):
        # <AP>
        if not self.transport:
            print "No transport defined, cannot send STUN request"
            return
        # </AP>
        if tid is None:
            tid = getRandomTID()
        mt = 0x1 # binding request
        avstr = ''
        # add any attributes
        if not avpairs:
            avpairs = ('CHANGE-REQUEST', CHANGE_NONE),
        for a,v in avpairs:
            avstr = avstr + struct.pack('!hh', StunTypes[a], len(v)) + v
        pktlen = len(avstr)
        if pktlen > 65535:
            raise ValueError, "stun request too big (%d bytes)"%pktlen
        pkt = struct.pack('!hh16s', mt, pktlen, tid) + avstr
        if STUNVERBOSE:
            print "sending request %r with %d avpairs to %r (in state %s)"%(
                            hexify(tid), len(avpairs), server, self._stunState)
        self.transport.write(pkt, server)

class StunDiscoveryProtocol(DatagramProtocol, _StunBase):

    stunDiscoveryRetries = 0

    def __init__(self, servers=DefaultServers, *args, **kwargs):
        # Potential STUN servers
        self._potentialStuns = {}
        # See flowchart ascii art at bottom of file.
        self._stunState = '1'
        self._finished = False
        self._altStunAddress = None
        self.externalAddress = None
        self.localAddress = None
        self.expectedTID = None
        self.oldTIDs = sets.Set()
        self.natType = None
        # <AP>
#         self.servers = [(socket.gethostbyname(host), port)
#                                             for host, port in servers]
        self.servers = [(host, port) for host, port in servers]
        # </AP>
        super(StunDiscoveryProtocol, self).__init__(*args, **kwargs)

    def initialStunRequest(self, address):
        # <AP>
        if self._finished:
            return
        # </AP>
        tid = getRandomTID()
        delayed = reactor.callLater(INITIAL_TIMEOUT,
                                    self.retransmitInitial, address, tid)
        self._potentialStuns[tid] = delayed
        self.oldTIDs.add(tid)
        self.sendRequest(address, tid=tid)

    def retransmitInitial(self, address, tid, count=1):
        # <AP>
        if self._finished:
            return
        # </AP>
        if count <= MAX_RETRANSMIT:
            t = BACKOFF_TIME * 2**min(count, MAX_BACKOFF)
            delayed = reactor.callLater(t, self.retransmitInitial,
                                                address, tid, count+1)
            self._potentialStuns[tid] = delayed
            self.sendRequest(address, tid=tid)
        else:
            if STUNVERBOSE:
                print "giving up on %r"%(address,)
            del self._potentialStuns[tid]
            if not self._potentialStuns:
                if STUNVERBOSE:
                    print "stun state 1 timeout - no internet UDP possible"
                self.natType = NatTypeUDPBlocked
                self._finishedStun()

    def datagramReceived(self, dgram, address):
        if self._finished:
            return
        mt, pktlen, tid = struct.unpack('!hh16s', dgram[:20])
        # Check tid is one we sent and haven't had a reply to yet
        if tid in self._potentialStuns:
            delayed = self._potentialStuns.get(tid)
            if delayed is not None:
                delayed.cancel()
            del self._potentialStuns[tid]
            if self._stunState == '1':
                # We got a (potentially) working STUN server!
                # Cancel the retransmit timers for the other ones
                for k in self._potentialStuns.keys():
                    self._potentialStuns[k].cancel()
                    self._potentialStuns[k] = None
                resdict = _parseStunResponse(dgram, address, self.expectedTID,
                                                self.oldTIDs)
                if not resdict:
                    return
                self.handleStunState1(resdict, address)
            else:
                # We already have a working STUN server to play with.
                pass
            return
        resdict = _parseStunResponse(dgram, address, self.expectedTID,
                                                self.oldTIDs)
        if not resdict:
            return
        if STUNVERBOSE: print 'calling handleStunState%s'%(self._stunState)
        getattr(self, 'handleStunState%s'%(self._stunState))(resdict, address)

    def handleStunState1(self, resdict, address):
        self.__dict__.update(resdict)

        if self.externalAddress and self._altStunAddress:
            if self.localAddress == self.externalAddress[0]:
                self._stunState = '2a'
            else:
                self._stunState = '2b'
            self.expectedTID = tid = getRandomTID()
            self.oldTIDs.add(tid)
            self.state2DelayedCall = reactor.callLater(INITIAL_TIMEOUT,
                                                self.retransmitStunState2,
                                                address, tid)
            self.sendRequest(address, tid, avpairs=(
                                    ('CHANGE-REQUEST', CHANGE_BOTH),))

    def handleStunState2a(self, resdict, address):
        self.state2DelayedCall.cancel()
        del self.state2DelayedCall
        if STUNVERBOSE:
            print "2a", resdict
        self.natType = NatTypeNone
        self._finishedStun()

    def handleStunState2b(self, resdict, address):
        self.state2DelayedCall.cancel()
        del self.state2DelayedCall
        if STUNVERBOSE:
            print "2b", resdict
        self.natType = NatTypeFullCone
        self._finishedStun()

    def retransmitStunState2(self, address, tid, count=1):
        # <AP>
        if self._finished:
            return
        # </AP>
        if count <= MAX_RETRANSMIT:
            t = BACKOFF_TIME * 2**min(count, MAX_BACKOFF)
            self.state2DelayedCall = reactor.callLater(t,
                                                    self.retransmitStunState2,
                                                    address, tid, count+1)
            self.sendRequest(address, tid, avpairs=(
                                    ('CHANGE-REQUEST', CHANGE_BOTH),))
        elif self._stunState == '2a':
            self.natType = NatTypeSymUDP
            self._finishedStun()
        else: # 2b
            # Off to state 3 we go!
            self._stunState = '3'
            self.state3DelayedCall = reactor.callLater(INITIAL_TIMEOUT,
                                                    self.retransmitStunState3,
                                                    address, tid)
            self.expectedTID = tid = getRandomTID()
            self.oldTIDs.add(tid)
            self.sendRequest(self._altStunAddress, tid)

    def handleStunState3(self, resdict, address):
        self.state3DelayedCall.cancel()
        del self.state3DelayedCall
        if STUNVERBOSE:
            print "3", resdict
        if self.externalAddress == resdict['externalAddress']:
            # State 4! wheee!
            self._stunState = '4'
            self.expectedTID = tid = getRandomTID()
            self.oldTIDs.add(tid)
            self.state4DelayedCall = reactor.callLater(INITIAL_TIMEOUT,
                                                self.retransmitStunState4,
                                                address, tid)
            self.expectedTID = tid = getRandomTID()
            self.oldTIDs.add(tid)
            self.sendRequest(address, tid, avpairs=(
                                    ('CHANGE-REQUEST', CHANGE_PORT),))
        else:
            self.natType = NatTypeSymmetric
            self._finishedStun()

    def retransmitStunState3(self, address, tid, count=1):
        # <AP>
        if self._finished:
            return
        # </AP>
        if count <= (2 * MAX_RETRANSMIT):
            t = BACKOFF_TIME * 2**min(count, MAX_BACKOFF)
            self.state3DelayedCall = reactor.callLater(t,
                                                    self.retransmitStunState3,
                                                    address, tid, count+1)
            self.sendRequest(self._altStunAddress, tid)
        else:
            log.err("STUN Failed in state 3, retrying")
            # We should do _something_ here. a new type BrokenNAT?
            self.stunDiscoveryRetries = self.stunDiscoveryRetries + 1
            if self.stunDiscoveryRetries < 5:
                reactor.callLater(0.2, self.startDiscovery)

    def handleStunState4(self, resdict, address):
        self.state4DelayedCall.cancel()
        del self.state4DelayedCall
        self.natType = NatTypeRestrictedCone
        self._finishedStun()

    def retransmitStunState4(self, address, tid, count = 1):
        # <AP>
        if self._finished:
            return
        # </AP>
        if count < MAX_RETRANSMIT:
            t = BACKOFF_TIME * 2**min(count, MAX_BACKOFF)
            self.state4DelayedCall = reactor.callLater(t,
                                                    self.retransmitStunState4,
                                                    address, tid, count+1)
            self.sendRequest(address, tid, avpairs=(
                                    ('CHANGE-REQUEST', CHANGE_PORT),))
        else:
            self.natType = NatTypePortRestricted
            self._finishedStun()


    def _finishedStun(self):
        self._finished = True
        self.finishedStun()

    def finishedStun(self):
        # Override in a subclass
        if STUNVERBOSE:
            print "firewall type is", self.natType

    def startDiscovery(self):
        from shtoom.nat import isBogusAddress, getLocalIPAddress
        if _ForceStunType is not None:
            self.natType = _ForceStunType
            reactor.callLater(0, self._finishedStun)
            return
        localAddress = self.transport.getHost().host
        if isBogusAddress(localAddress):
            d = getLocalIPAddress()
            d.addCallback(self._resolveStunServers)
        else:
            self._resolveStunServers(localAddress)


    def _resolveStunServers(self, localAddress):
        self.localAddress = localAddress
        # reactor.resolve the hosts!
        for host, port in self.servers:
            d = reactor.resolve(host)
            d.addCallback(lambda x,p=port: self.initialStunRequest((x, p)))