← Back to branch summary

~ntt-pf-lab/nova/monkey_patch_notification

« back to all changes in this revision

Viewing changes to vendor/Twisted-10.0.0/twisted/test/iosim.py

  • Committer: Jesse Andrews
  • Date: 2010-05-28 06:05:26 UTC
  • Revision ID: git-v1:bf6e6e718cdc7488e2da87b21e258ccc065fe499
initial commit

Show diffs side-by-side

added added

removed removed

Lines of Context:
vendor/Twisted-10.0.0/twisted/test/iosim.py
 
1
# -*- test-case-name: twisted.test.test_amp.TLSTest -*-
 
2
# Copyright (c) 2008-2009 Twisted Matrix Laboratories.
 
3
# See LICENSE for details.
 
4
 
 
5
"""
 
6
Utilities and helpers for simulating a network
 
7
"""
 
8
 
 
9
import itertools
 
10
 
 
11
try:
 
12
    from OpenSSL.SSL import Error as NativeOpenSSLError
 
13
except ImportError:
 
14
    pass
 
15
 
 
16
from zope.interface import implements, directlyProvides
 
17
 
 
18
from twisted.python.failure import Failure
 
19
from twisted.internet import error
 
20
from twisted.internet import interfaces
 
21
 
 
22
class TLSNegotiation:
 
23
    def __init__(self, obj, connectState):
 
24
        self.obj = obj
 
25
        self.connectState = connectState
 
26
        self.sent = False
 
27
        self.readyToSend = connectState
 
28
 
 
29
    def __repr__(self):
 
30
        return 'TLSNegotiation(%r)' % (self.obj,)
 
31
 
 
32
    def pretendToVerify(self, other, tpt):
 
33
        # Set the transport problems list here?  disconnections?
 
34
        # hmmmmm... need some negative path tests.
 
35
 
 
36
        if not self.obj.iosimVerify(other.obj):
 
37
            tpt.disconnectReason = NativeOpenSSLError()
 
38
            tpt.loseConnection()
 
39
 
 
40
 
 
41
class FakeTransport:
 
42
    """
 
43
    A wrapper around a file-like object to make it behave as a Transport.
 
44
 
 
45
    This doesn't actually stream the file to the attached protocol,
 
46
    and is thus useful mainly as a utility for debugging protocols.
 
47
    """
 
48
    implements(interfaces.ITransport,
 
49
               interfaces.ITLSTransport) # ha ha not really
 
50
 
 
51
    _nextserial = itertools.count().next
 
52
    closed = 0
 
53
    disconnecting = 0
 
54
    disconnected = 0
 
55
    disconnectReason = error.ConnectionDone("Connection done")
 
56
    producer = None
 
57
    streamingProducer = 0
 
58
    tls = None
 
59
 
 
60
    def __init__(self):
 
61
        self.stream = []
 
62
        self.serial = self._nextserial()
 
63
 
 
64
    def __repr__(self):
 
65
        return 'FakeTransport<%s,%s,%s>' % (
 
66
            self.isServer and 'S' or 'C', self.serial,
 
67
            self.protocol.__class__.__name__)
 
68
 
 
69
    def write(self, data):
 
70
        if self.tls is not None:
 
71
            self.tlsbuf.append(data)
 
72
        else:
 
73
            self.stream.append(data)
 
74
 
 
75
    def _checkProducer(self):
 
76
        # Cheating; this is called at "idle" times to allow producers to be
 
77
        # found and dealt with
 
78
        if self.producer:
 
79
            self.producer.resumeProducing()
 
80
 
 
81
    def registerProducer(self, producer, streaming):
 
82
        """From abstract.FileDescriptor
 
83
        """
 
84
        self.producer = producer
 
85
        self.streamingProducer = streaming
 
86
        if not streaming:
 
87
            producer.resumeProducing()
 
88
 
 
89
    def unregisterProducer(self):
 
90
        self.producer = None
 
91
 
 
92
    def stopConsuming(self):
 
93
        self.unregisterProducer()
 
94
        self.loseConnection()
 
95
 
 
96
    def writeSequence(self, iovec):
 
97
        self.write("".join(iovec))
 
98
 
 
99
    def loseConnection(self):
 
100
        self.disconnecting = True
 
101
 
 
102
    def reportDisconnect(self):
 
103
        if self.tls is not None:
 
104
            # We were in the middle of negotiating!  Must have been a TLS problem.
 
105
            err = NativeOpenSSLError()
 
106
        else:
 
107
            err = self.disconnectReason
 
108
        self.protocol.connectionLost(Failure(err))
 
109
 
 
110
    def getPeer(self):
 
111
        # XXX: According to ITransport, this should return an IAddress!
 
112
        return 'file', 'file'
 
113
 
 
114
    def getHost(self):
 
115
        # XXX: According to ITransport, this should return an IAddress!
 
116
        return 'file'
 
117
 
 
118
    def resumeProducing(self):
 
119
        # Never sends data anyways
 
120
        pass
 
121
 
 
122
    def pauseProducing(self):
 
123
        # Never sends data anyways
 
124
        pass
 
125
 
 
126
    def stopProducing(self):
 
127
        self.loseConnection()
 
128
 
 
129
    def startTLS(self, contextFactory, beNormal=True):
 
130
        # Nothing's using this feature yet, but startTLS has an undocumented
 
131
        # second argument which defaults to true; if set to False, servers will
 
132
        # behave like clients and clients will behave like servers.
 
133
        connectState = self.isServer ^ beNormal
 
134
        self.tls = TLSNegotiation(contextFactory, connectState)
 
135
        self.tlsbuf = []
 
136
 
 
137
    def getOutBuffer(self):
 
138
        S = self.stream
 
139
        if S:
 
140
            self.stream = []
 
141
            return ''.join(S)
 
142
        elif self.tls is not None:
 
143
            if self.tls.readyToSend:
 
144
                # Only _send_ the TLS negotiation "packet" if I'm ready to.
 
145
                self.tls.sent = True
 
146
                return self.tls
 
147
            else:
 
148
                return None
 
149
        else:
 
150
            return None
 
151
 
 
152
    def bufferReceived(self, buf):
 
153
        if isinstance(buf, TLSNegotiation):
 
154
            assert self.tls is not None # By the time you're receiving a
 
155
                                        # negotiation, you have to have called
 
156
                                        # startTLS already.
 
157
            if self.tls.sent:
 
158
                self.tls.pretendToVerify(buf, self)
 
159
                self.tls = None # we're done with the handshake if we've gotten
 
160
                                # this far... although maybe it failed...?
 
161
                # TLS started!  Unbuffer...
 
162
                b, self.tlsbuf = self.tlsbuf, None
 
163
                self.writeSequence(b)
 
164
                directlyProvides(self, interfaces.ISSLTransport)
 
165
            else:
 
166
                # We haven't sent our own TLS negotiation: time to do that!
 
167
                self.tls.readyToSend = True
 
168
        else:
 
169
            self.protocol.dataReceived(buf)
 
170
 
 
171
 
 
172
 
 
173
def makeFakeClient(c):
 
174
    ft = FakeTransport()
 
175
    ft.isServer = False
 
176
    ft.protocol = c
 
177
    return ft
 
178
 
 
179
def makeFakeServer(s):
 
180
    ft = FakeTransport()
 
181
    ft.isServer = True
 
182
    ft.protocol = s
 
183
    return ft
 
184
 
 
185
class IOPump:
 
186
    """Utility to pump data between clients and servers for protocol testing.
 
187
 
 
188
    Perhaps this is a utility worthy of being in protocol.py?
 
189
    """
 
190
    def __init__(self, client, server, clientIO, serverIO, debug):
 
191
        self.client = client
 
192
        self.server = server
 
193
        self.clientIO = clientIO
 
194
        self.serverIO = serverIO
 
195
        self.debug = debug
 
196
 
 
197
    def flush(self, debug=False):
 
198
        """Pump until there is no more input or output.
 
199
 
 
200
        Returns whether any data was moved.
 
201
        """
 
202
        result = False
 
203
        for x in range(1000):
 
204
            if self.pump(debug):
 
205
                result = True
 
206
            else:
 
207
                break
 
208
        else:
 
209
            assert 0, "Too long"
 
210
        return result
 
211
 
 
212
 
 
213
    def pump(self, debug=False):
 
214
        """Move data back and forth.
 
215
 
 
216
        Returns whether any data was moved.
 
217
        """
 
218
        if self.debug or debug:
 
219
            print '-- GLUG --'
 
220
        sData = self.serverIO.getOutBuffer()
 
221
        cData = self.clientIO.getOutBuffer()
 
222
        self.clientIO._checkProducer()
 
223
        self.serverIO._checkProducer()
 
224
        if self.debug or debug:
 
225
            print '.'
 
226
            # XXX slightly buggy in the face of incremental output
 
227
            if cData:
 
228
                print 'C: '+repr(cData)
 
229
            if sData:
 
230
                print 'S: '+repr(sData)
 
231
        if cData:
 
232
            self.serverIO.bufferReceived(cData)
 
233
        if sData:
 
234
            self.clientIO.bufferReceived(sData)
 
235
        if cData or sData:
 
236
            return True
 
237
        if (self.serverIO.disconnecting and
 
238
            not self.serverIO.disconnected):
 
239
            if self.debug or debug:
 
240
                print '* C'
 
241
            self.serverIO.disconnected = True
 
242
            self.clientIO.disconnecting = True
 
243
            self.clientIO.reportDisconnect()
 
244
            return True
 
245
        if self.clientIO.disconnecting and not self.clientIO.disconnected:
 
246
            if self.debug or debug:
 
247
                print '* S'
 
248
            self.clientIO.disconnected = True
 
249
            self.serverIO.disconnecting = True
 
250
            self.serverIO.reportDisconnect()
 
251
            return True
 
252
        return False
 
253
 
 
254
 
 
255
def connectedServerAndClient(ServerClass, ClientClass,
 
256
                             clientTransportFactory=makeFakeClient,
 
257
                             serverTransportFactory=makeFakeServer,
 
258
                             debug=False):
 
259
    """Returns a 3-tuple: (client, server, pump)
 
260
    """
 
261
    c = ClientClass()
 
262
    s = ServerClass()
 
263
    cio = clientTransportFactory(c)
 
264
    sio = serverTransportFactory(s)
 
265
    c.makeConnection(cio)
 
266
    s.makeConnection(sio)
 
267
    pump = IOPump(c, s, cio, sio, debug)
 
268
    # kick off server greeting, etc
 
269
    pump.flush()
 
270
    return c, s, pump