twisted/internet/test/test_newtls.py

   1 # Copyright (c) Twisted Matrix Laboratories.
   2 # See LICENSE for details.
   3
   4 """
   5 Tests for L{twisted.internet._newtls}.
   6 """
   7
   8 from twisted.trial import unittest
   9 from twisted.internet.test.reactormixins import ReactorBuilder, runProtocolsWithReactor
  10 from twisted.internet.test.reactormixins import ConnectableProtocol
  11 from twisted.internet.test.test_tls import SSLCreator, TLSMixin
  12 from twisted.internet.test.test_tls import StartTLSClientCreator
  13 from twisted.internet.test.test_tls import ContextGeneratingMixin
  14 from twisted.internet.test.test_tcp import TCPCreator
  15 try:
  16     from twisted.protocols import tls
  17     from twisted.internet import _newtls
  18 except ImportError:
  19     _newtls = None
  20
  21
  22 class BypassTLSTests(unittest.TestCase):
  23     """
  24     Tests for the L{_newtls._BypassTLS} class.
  25     """
  26
  27     if not _newtls:
  28         skip = "Couldn't import _newtls, perhaps pyOpenSSL is old or missing"
  29
  30     def test_loseConnectionPassThrough(self):
  31         """
  32         C{_BypassTLS.loseConnection} calls C{loseConnection} on the base
  33         class, while preserving any default argument in the base class'
  34         C{loseConnection} implementation.
  35         """
  36         default = object()
  37         result = []
  38
  39         class FakeTransport(object):
  40             def loseConnection(self, _connDone=default):
  41                 result.append(_connDone)
  42
  43         bypass = _newtls._BypassTLS(FakeTransport, FakeTransport())
  44
  45         # The default from FakeTransport is used:
  46         bypass.loseConnection()
  47         self.assertEqual(result, [default])
  48
  49         # And we can pass our own:
  50         notDefault = object()
  51         bypass.loseConnection(notDefault)
  52         self.assertEqual(result, [default, notDefault])
  53
  54
  55
  56 class FakeProducer(object):
  57     """
  58     A producer that does nothing.
  59     """
  60
  61     def pauseProducing(self):
  62         pass
  63
  64
  65     def resumeProducing(self):
  66         pass
  67
  68
  69     def stopProducing(self):
  70         pass
  71
  72
  73
  74 class ProducerProtocol(ConnectableProtocol):
  75     """
  76     Register a producer, unregister it, and verify the producer hooks up to
  77     innards of C{TLSMemoryBIOProtocol}.
  78     """
  79
  80     def __init__(self, producer, result):
  81         self.producer = producer
  82         self.result = result
  83
  84
  85     def connectionMade(self):
  86         if not isinstance(self.transport.protocol,
  87                           tls.TLSMemoryBIOProtocol):
  88             # Either the test or the code have a bug...
  89             raise RuntimeError("TLSMemoryBIOProtocol not hooked up.")
  90
  91         self.transport.registerProducer(self.producer, True)
  92         # The producer was registered with the TLSMemoryBIOProtocol:
  93         self.result.append(self.transport.protocol._producer._producer)
  94
  95         self.transport.unregisterProducer()
  96         # The producer was unregistered from the TLSMemoryBIOProtocol:
  97         self.result.append(self.transport.protocol._producer)
  98         self.transport.loseConnection()
  99
 100
 101
 102 class ProducerTestsMixin(ReactorBuilder, TLSMixin, ContextGeneratingMixin):
 103     """
 104     Test the new TLS code integrates C{TLSMemoryBIOProtocol} correctly.
 105     """
 106
 107     if not _newtls:
 108         skip = "Could not import twisted.internet._newtls"
 109
 110     def test_producerSSLFromStart(self):
 111         """
 112         C{registerProducer} and C{unregisterProducer} on TLS transports
 113         created as SSL from the get go are passed to the
 114         C{TLSMemoryBIOProtocol}, not the underlying transport directly.
 115         """
 116         result = []
 117         producer = FakeProducer()
 118
 119         runProtocolsWithReactor(self, ConnectableProtocol(),
 120                                 ProducerProtocol(producer, result),
 121                                 SSLCreator())
 122         self.assertEqual(result, [producer, None])
 123
 124
 125     def test_producerAfterStartTLS(self):
 126         """
 127         C{registerProducer} and C{unregisterProducer} on TLS transports
 128         created by C{startTLS} are passed to the C{TLSMemoryBIOProtocol}, not
 129         the underlying transport directly.
 130         """
 131         result = []
 132         producer = FakeProducer()
 133
 134         runProtocolsWithReactor(self, ConnectableProtocol(),
 135                                 ProducerProtocol(producer, result),
 136                                 StartTLSClientCreator())
 137         self.assertEqual(result, [producer, None])
 138
 139
 140     def startTLSAfterRegisterProducer(self, streaming):
 141         """
 142         When a producer is registered, and then startTLS is called,
 143         the producer is re-registered with the C{TLSMemoryBIOProtocol}.
 144         """
 145         clientContext = self.getClientContext()
 146         serverContext = self.getServerContext()
 147         result = []
 148         producer = FakeProducer()
 149
 150         class RegisterTLSProtocol(ConnectableProtocol):
 151             def connectionMade(self):
 152                 self.transport.registerProducer(producer, streaming)
 153                 self.transport.startTLS(serverContext)
 154                 # Store TLSMemoryBIOProtocol and underlying transport producer
 155                 # status:
 156                 if streaming:
 157                     # _ProducerMembrane -> producer:
 158                     result.append(self.transport.protocol._producer._producer)
 159                     result.append(self.transport.producer._producer)
 160                 else:
 161                     # _ProducerMembrane -> _PullToPush -> producer:
 162                     result.append(
 163                         self.transport.protocol._producer._producer._producer)
 164                     result.append(self.transport.producer._producer._producer)
 165                 self.transport.unregisterProducer()
 166                 self.transport.loseConnection()
 167
 168         class StartTLSProtocol(ConnectableProtocol):
 169             def connectionMade(self):
 170                 self.transport.startTLS(clientContext)
 171
 172         runProtocolsWithReactor(self, RegisterTLSProtocol(),
 173                                 StartTLSProtocol(), TCPCreator())
 174         self.assertEqual(result, [producer, producer])
 175
 176
 177     def test_startTLSAfterRegisterProducerStreaming(self):
 178         """
 179         When a streaming producer is registered, and then startTLS is called,
 180         the producer is re-registered with the C{TLSMemoryBIOProtocol}.
 181         """
 182         self.startTLSAfterRegisterProducer(True)
 183
 184
 185     def test_startTLSAfterRegisterProducerNonStreaming(self):
 186         """
 187         When a non-streaming producer is registered, and then startTLS is
 188         called, the producer is re-registered with the
 189         C{TLSMemoryBIOProtocol}.
 190         """
 191         self.startTLSAfterRegisterProducer(False)
 192
 193
 194 globals().update(ProducerTestsMixin.makeTestCaseClasses())