← Back to branch summary

~certify-web-dev/twisted/certify-trunk

« back to all changes in this revision

Viewing changes to twisted/names/srvconnect.py

  • Committer: Bazaar Package Importer
  • Author(s): Matthias Klose
  • Date: 2007-01-17 14:52:35 UTC
  • mfrom: (1.1.5 upstream) (2.1.2 etch)
  • Revision ID: james.westby@ubuntu.com-20070117145235-btmig6qfmqfen0om
Tags: 2.5.0-0ubuntu1
New upstream version, compatible with python2.5.

Show diffs side-by-side

added added

removed removed

Lines of Context:
twisted/names/srvconnect.py
 
1
# Copyright (c) 2001-2006 Twisted Matrix Laboratories.
 
2
# See LICENSE for details.
 
3
 
 
4
import random
 
5
 
 
6
from zope.interface import implements
 
7
 
 
8
from twisted.internet import error, interfaces
 
9
 
 
10
from twisted.names import client
 
11
 
 
12
class _SRVConnector_ClientFactoryWrapper:
 
13
    def __init__(self, connector, wrappedFactory):
 
14
        self.__connector = connector
 
15
        self.__wrappedFactory = wrappedFactory
 
16
 
 
17
    def startedConnecting(self, connector):
 
18
        self.__wrappedFactory.startedConnecting(self.__connector)
 
19
 
 
20
    def clientConnectionFailed(self, connector, reason):
 
21
        self.__connector.connectionFailed(reason)
 
22
 
 
23
    def clientConnectionLost(self, connector, reason):
 
24
        self.__connector.connectionLost(reason)
 
25
 
 
26
    def __getattr__(self, key):
 
27
        return getattr(self.__wrappedFactory, key)
 
28
 
 
29
class SRVConnector:
 
30
    """A connector that looks up DNS SRV records. See RFC2782."""
 
31
 
 
32
    implements(interfaces.IConnector)
 
33
 
 
34
    stopAfterDNS=0
 
35
 
 
36
    def __init__(self, reactor, service, domain, factory,
 
37
                 protocol='tcp', connectFuncName='connectTCP',
 
38
                 connectFuncArgs=(),
 
39
                 connectFuncKwArgs={},
 
40
                 ):
 
41
        self.reactor = reactor
 
42
        self.service = service
 
43
        self.domain = domain
 
44
        self.factory = factory
 
45
 
 
46
        self.protocol = protocol
 
47
        self.connectFuncName = connectFuncName
 
48
        self.connectFuncArgs = connectFuncArgs
 
49
        self.connectFuncKwArgs = connectFuncKwArgs
 
50
 
 
51
        self.connector = None
 
52
        self.servers = None
 
53
        self.orderedServers = None # list of servers already used in this round
 
54
 
 
55
    def connect(self):
 
56
        """Start connection to remote server."""
 
57
        self.factory.doStart()
 
58
        self.factory.startedConnecting(self)
 
59
 
 
60
        if not self.servers:
 
61
            if self.domain is None:
 
62
                self.connectionFailed(error.DNSLookupError("Domain is not defined."))
 
63
                return
 
64
            d = client.lookupService('_%s._%s.%s' % (self.service,
 
65
                                                     self.protocol,
 
66
                                                     self.domain))
 
67
            d.addCallback(self._cbGotServers)
 
68
            d.addCallback(lambda x, self=self: self._reallyConnect())
 
69
            d.addErrback(self.connectionFailed)
 
70
        elif self.connector is None:
 
71
            self._reallyConnect()
 
72
        else:
 
73
            self.connector.connect()
 
74
 
 
75
    def _cbGotServers(self, (answers, auth, add)):
 
76
        if len(answers)==1 and answers[0].payload.target=='.':
 
77
            # decidedly not available
 
78
            raise error.DNSLookupError("Service %s not available for domain %s."
 
79
                                       % (repr(self.service), repr(self.domain)))
 
80
 
 
81
        self.servers = []
 
82
        self.orderedServers = []
 
83
        for a in answers:
 
84
            self.orderedServers.append((a.payload.priority, a.payload.weight,
 
85
                                        str(a.payload.target), a.payload.port))
 
86
 
 
87
    def _serverCmp(self, a, b):
 
88
        if a[0]!=b[0]:
 
89
            return cmp(a[0], b[0])
 
90
        else:
 
91
            return cmp(a[1], b[1])
 
92
 
 
93
    def pickServer(self):
 
94
        assert self.servers is not None
 
95
        assert self.orderedServers is not None
 
96
 
 
97
        if not self.servers and not self.orderedServers:
 
98
            # no SRV record, fall back..
 
99
            return self.domain, self.service
 
100
 
 
101
        if not self.servers and self.orderedServers:
 
102
            # start new round
 
103
            self.servers = self.orderedServers
 
104
            self.orderedServers = []
 
105
 
 
106
        assert self.servers
 
107
 
 
108
        self.servers.sort(self._serverCmp)
 
109
        minPriority=self.servers[0][0]
 
110
 
 
111
        weightIndex = zip(xrange(len(self.servers)), [x[1] for x in self.servers
 
112
                                                      if x[0]==minPriority])
 
113
        weightSum = reduce(lambda x, y: (None, x[1]+y[1]), weightIndex, (None, 0))[1]
 
114
        rand = random.randint(0, weightSum)
 
115
 
 
116
        for index, weight in weightIndex:
 
117
            weightSum -= weight
 
118
            if weightSum <= 0:
 
119
                chosen = self.servers[index]
 
120
                del self.servers[index]
 
121
                self.orderedServers.append(chosen)
 
122
 
 
123
                p, w, host, port = chosen
 
124
                return host, port
 
125
 
 
126
        raise RuntimeError, 'Impossible %s pickServer result.' % self.__class__.__name__
 
127
 
 
128
    def _reallyConnect(self):
 
129
        if self.stopAfterDNS:
 
130
            self.stopAfterDNS=0
 
131
            return
 
132
 
 
133
        self.host, self.port = self.pickServer()
 
134
        assert self.host is not None, 'Must have a host to connect to.'
 
135
        assert self.port is not None, 'Must have a port to connect to.'
 
136
 
 
137
        connectFunc = getattr(self.reactor, self.connectFuncName)
 
138
        self.connector=connectFunc(
 
139
            self.host, self.port,
 
140
            _SRVConnector_ClientFactoryWrapper(self, self.factory),
 
141
            *self.connectFuncArgs, **self.connectFuncKwArgs)
 
142
 
 
143
    def stopConnecting(self):
 
144
        """Stop attempting to connect."""
 
145
        if self.connector:
 
146
            self.connector.stopConnecting()
 
147
        else:
 
148
            self.stopAfterDNS=1
 
149
 
 
150
    def disconnect(self):
 
151
        """Disconnect whatever our are state is."""
 
152
        if self.connector is not None:
 
153
            self.connector.disconnect()
 
154
        else:
 
155
            self.stopConnecting()
 
156
 
 
157
    def getDestination(self):
 
158
        assert self.connector
 
159
        return self.connector.getDestination()
 
160
 
 
161
    def connectionFailed(self, reason):
 
162
        self.factory.clientConnectionFailed(self, reason)
 
163
        self.factory.doStop()
 
164
 
 
165
    def connectionLost(self, reason):
 
166
        self.factory.clientConnectionLost(self, reason)
 
167
        self.factory.doStop()
 
168