123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594 |
- # -*- test-case-name: twisted.internet.test.test_tcp -*-
- # Copyright (c) Twisted Matrix Laboratories.
- # See LICENSE for details.
-
- """
- Various helpers for tests for connection-oriented transports.
- """
-
-
- import socket
- from gc import collect
- from typing import Optional
- from weakref import ref
-
- from zope.interface.verify import verifyObject
-
- from twisted.internet.defer import Deferred, gatherResults
- from twisted.internet.interfaces import IConnector, IReactorFDSet
- from twisted.internet.protocol import ClientFactory, Protocol, ServerFactory
- from twisted.internet.test.reactormixins import needsRunningReactor
- from twisted.python import context, log
- from twisted.python.failure import Failure
- from twisted.python.log import ILogContext, err, msg
- from twisted.python.runtime import platform
- from twisted.test.test_tcp import ClosingProtocol
- from twisted.trial.unittest import SkipTest
-
-
- def findFreePort(interface="127.0.0.1", family=socket.AF_INET, type=socket.SOCK_STREAM):
- """
- Ask the platform to allocate a free port on the specified interface, then
- release the socket and return the address which was allocated.
-
- @param interface: The local address to try to bind the port on.
- @type interface: C{str}
-
- @param type: The socket type which will use the resulting port.
-
- @return: A two-tuple of address and port, like that returned by
- L{socket.getsockname}.
- """
- addr = socket.getaddrinfo(interface, 0)[0][4]
- probe = socket.socket(family, type)
- try:
- probe.bind(addr)
- if family == socket.AF_INET6:
- sockname = probe.getsockname()
- hostname = socket.getnameinfo(
- sockname, socket.NI_NUMERICHOST | socket.NI_NUMERICSERV
- )[0]
- return (hostname, sockname[1])
- else:
- return probe.getsockname()
- finally:
- probe.close()
-
-
- class ConnectableProtocol(Protocol):
- """
- A protocol to be used with L{runProtocolsWithReactor}.
-
- The protocol and its pair should eventually disconnect from each other.
-
- @ivar reactor: The reactor used in this test.
-
- @ivar disconnectReason: The L{Failure} passed to C{connectionLost}.
-
- @ivar _done: A L{Deferred} which will be fired when the connection is
- lost.
- """
-
- disconnectReason = None
-
- def _setAttributes(self, reactor, done):
- """
- Set attributes on the protocol that are known only externally; this
- will be called by L{runProtocolsWithReactor} when this protocol is
- instantiated.
-
- @param reactor: The reactor used in this test.
-
- @param done: A L{Deferred} which will be fired when the connection is
- lost.
- """
- self.reactor = reactor
- self._done = done
-
- def connectionLost(self, reason):
- self.disconnectReason = reason
- self._done.callback(None)
- del self._done
-
-
- class EndpointCreator:
- """
- Create client and server endpoints that know how to connect to each other.
- """
-
- def server(self, reactor):
- """
- Return an object providing C{IStreamServerEndpoint} for use in creating
- a server to use to establish the connection type to be tested.
- """
- raise NotImplementedError()
-
- def client(self, reactor, serverAddress):
- """
- Return an object providing C{IStreamClientEndpoint} for use in creating
- a client to use to establish the connection type to be tested.
- """
- raise NotImplementedError()
-
-
- class _SingleProtocolFactory(ClientFactory):
- """
- Factory to be used by L{runProtocolsWithReactor}.
-
- It always returns the same protocol (i.e. is intended for only a single
- connection).
- """
-
- def __init__(self, protocol):
- self._protocol = protocol
-
- def buildProtocol(self, addr):
- return self._protocol
-
-
- def runProtocolsWithReactor(
- reactorBuilder, serverProtocol, clientProtocol, endpointCreator
- ):
- """
- Connect two protocols using endpoints and a new reactor instance.
-
- A new reactor will be created and run, with the client and server protocol
- instances connected to each other using the given endpoint creator. The
- protocols should run through some set of tests, then disconnect; when both
- have disconnected the reactor will be stopped and the function will
- return.
-
- @param reactorBuilder: A L{ReactorBuilder} instance.
-
- @param serverProtocol: A L{ConnectableProtocol} that will be the server.
-
- @param clientProtocol: A L{ConnectableProtocol} that will be the client.
-
- @param endpointCreator: An instance of L{EndpointCreator}.
-
- @return: The reactor run by this test.
- """
- reactor = reactorBuilder.buildReactor()
- serverProtocol._setAttributes(reactor, Deferred())
- clientProtocol._setAttributes(reactor, Deferred())
- serverFactory = _SingleProtocolFactory(serverProtocol)
- clientFactory = _SingleProtocolFactory(clientProtocol)
-
- # Listen on a port:
- serverEndpoint = endpointCreator.server(reactor)
- d = serverEndpoint.listen(serverFactory)
-
- # Connect to the port:
- def gotPort(p):
- clientEndpoint = endpointCreator.client(reactor, p.getHost())
- return clientEndpoint.connect(clientFactory)
-
- d.addCallback(gotPort)
-
- # Stop reactor when both connections are lost:
- def failed(result):
- log.err(result, "Connection setup failed.")
-
- disconnected = gatherResults([serverProtocol._done, clientProtocol._done])
- d.addCallback(lambda _: disconnected)
- d.addErrback(failed)
- d.addCallback(lambda _: needsRunningReactor(reactor, reactor.stop))
-
- reactorBuilder.runReactor(reactor)
- return reactor
-
-
- def _getWriters(reactor):
- """
- Like L{IReactorFDSet.getWriters}, but with support for IOCP reactor as
- well.
- """
- if IReactorFDSet.providedBy(reactor):
- return reactor.getWriters()
- elif "IOCP" in reactor.__class__.__name__:
- return reactor.handles
- else:
- # Cannot tell what is going on.
- raise Exception(f"Cannot find writers on {reactor!r}")
-
-
- class _AcceptOneClient(ServerFactory):
- """
- This factory fires a L{Deferred} with a protocol instance shortly after it
- is constructed (hopefully long enough afterwards so that it has been
- connected to a transport).
-
- @ivar reactor: The reactor used to schedule the I{shortly}.
-
- @ivar result: A L{Deferred} which will be fired with the protocol instance.
- """
-
- def __init__(self, reactor, result):
- self.reactor = reactor
- self.result = result
-
- def buildProtocol(self, addr):
- protocol = ServerFactory.buildProtocol(self, addr)
- self.reactor.callLater(0, self.result.callback, protocol)
- return protocol
-
-
- class _SimplePullProducer:
- """
- A pull producer which writes one byte whenever it is resumed. For use by
- C{test_unregisterProducerAfterDisconnect}.
- """
-
- def __init__(self, consumer):
- self.consumer = consumer
-
- def stopProducing(self):
- pass
-
- def resumeProducing(self):
- log.msg("Producer.resumeProducing")
- self.consumer.write(b"x")
-
-
- class Stop(ClientFactory):
- """
- A client factory which stops a reactor when a connection attempt fails.
- """
-
- failReason = None
-
- def __init__(self, reactor):
- self.reactor = reactor
-
- def clientConnectionFailed(self, connector, reason):
- self.failReason = reason
- msg(f"Stop(CF) cCFailed: {reason.getErrorMessage()}")
- self.reactor.stop()
-
-
- class ClosingLaterProtocol(ConnectableProtocol):
- """
- ClosingLaterProtocol exchanges one byte with its peer and then disconnects
- itself. This is mostly a work-around for the fact that connectionMade is
- called before the SSL handshake has completed.
- """
-
- def __init__(self, onConnectionLost):
- self.lostConnectionReason = None
- self.onConnectionLost = onConnectionLost
-
- def connectionMade(self):
- msg("ClosingLaterProtocol.connectionMade")
-
- def dataReceived(self, bytes):
- msg(f"ClosingLaterProtocol.dataReceived {bytes!r}")
- self.transport.loseConnection()
-
- def connectionLost(self, reason):
- msg("ClosingLaterProtocol.connectionLost")
- self.lostConnectionReason = reason
- self.onConnectionLost.callback(self)
-
-
- class ConnectionTestsMixin:
- """
- This mixin defines test methods which should apply to most L{ITransport}
- implementations.
- """
-
- endpoints: Optional[EndpointCreator] = None
-
- def test_logPrefix(self):
- """
- Client and server transports implement L{ILoggingContext.logPrefix} to
- return a message reflecting the protocol they are running.
- """
-
- class CustomLogPrefixProtocol(ConnectableProtocol):
- def __init__(self, prefix):
- self._prefix = prefix
- self.system = None
-
- def connectionMade(self):
- self.transport.write(b"a")
-
- def logPrefix(self):
- return self._prefix
-
- def dataReceived(self, bytes):
- self.system = context.get(ILogContext)["system"]
- self.transport.write(b"b")
- # Only close connection if both sides have received data, so
- # that both sides have system set.
- if b"b" in bytes:
- self.transport.loseConnection()
-
- client = CustomLogPrefixProtocol("Custom Client")
- server = CustomLogPrefixProtocol("Custom Server")
- runProtocolsWithReactor(self, server, client, self.endpoints)
- self.assertIn("Custom Client", client.system)
- self.assertIn("Custom Server", server.system)
-
- def test_writeAfterDisconnect(self):
- """
- After a connection is disconnected, L{ITransport.write} and
- L{ITransport.writeSequence} are no-ops.
- """
- reactor = self.buildReactor()
-
- finished = []
-
- serverConnectionLostDeferred = Deferred()
- protocol = lambda: ClosingLaterProtocol(serverConnectionLostDeferred)
- portDeferred = self.endpoints.server(reactor).listen(
- ServerFactory.forProtocol(protocol)
- )
-
- def listening(port):
- msg(f"Listening on {port.getHost()!r}")
- endpoint = self.endpoints.client(reactor, port.getHost())
-
- lostConnectionDeferred = Deferred()
- protocol = lambda: ClosingLaterProtocol(lostConnectionDeferred)
- client = endpoint.connect(ClientFactory.forProtocol(protocol))
-
- def write(proto):
- msg(f"About to write to {proto!r}")
- proto.transport.write(b"x")
-
- client.addCallbacks(write, lostConnectionDeferred.errback)
-
- def disconnected(proto):
- msg(f"{proto!r} disconnected")
- proto.transport.write(b"some bytes to get lost")
- proto.transport.writeSequence([b"some", b"more"])
- finished.append(True)
-
- lostConnectionDeferred.addCallback(disconnected)
- serverConnectionLostDeferred.addCallback(disconnected)
- return gatherResults([lostConnectionDeferred, serverConnectionLostDeferred])
-
- def onListen():
- portDeferred.addCallback(listening)
- portDeferred.addErrback(err)
- portDeferred.addCallback(lambda ignored: reactor.stop())
-
- needsRunningReactor(reactor, onListen)
-
- self.runReactor(reactor)
- self.assertEqual(finished, [True, True])
-
- def test_protocolGarbageAfterLostConnection(self):
- """
- After the connection a protocol is being used for is closed, the
- reactor discards all of its references to the protocol.
- """
- lostConnectionDeferred = Deferred()
- clientProtocol = ClosingLaterProtocol(lostConnectionDeferred)
- clientRef = ref(clientProtocol)
-
- reactor = self.buildReactor()
- portDeferred = self.endpoints.server(reactor).listen(
- ServerFactory.forProtocol(Protocol)
- )
-
- def listening(port):
- msg(f"Listening on {port.getHost()!r}")
- endpoint = self.endpoints.client(reactor, port.getHost())
-
- client = endpoint.connect(ClientFactory.forProtocol(lambda: clientProtocol))
-
- def disconnect(proto):
- msg(f"About to disconnect {proto!r}")
- proto.transport.loseConnection()
-
- client.addCallback(disconnect)
- client.addErrback(lostConnectionDeferred.errback)
- return lostConnectionDeferred
-
- def onListening():
- portDeferred.addCallback(listening)
- portDeferred.addErrback(err)
- portDeferred.addBoth(lambda ignored: reactor.stop())
-
- needsRunningReactor(reactor, onListening)
-
- self.runReactor(reactor)
-
- # Drop the reference and get the garbage collector to tell us if there
- # are no references to the protocol instance left in the reactor.
- clientProtocol = None
- collect()
- self.assertIsNone(clientRef())
-
-
- class LogObserverMixin:
- """
- Mixin for L{TestCase} subclasses which want to observe log events.
- """
-
- def observe(self):
- loggedMessages = []
- log.addObserver(loggedMessages.append)
- self.addCleanup(log.removeObserver, loggedMessages.append)
- return loggedMessages
-
-
- class BrokenContextFactory:
- """
- A context factory with a broken C{getContext} method, for exercising the
- error handling for such a case.
- """
-
- message = "Some path was wrong maybe"
-
- def getContext(self):
- raise ValueError(self.message)
-
-
- class StreamClientTestsMixin:
- """
- This mixin defines tests applicable to SOCK_STREAM client implementations.
-
- This must be mixed in to a L{ReactorBuilder
- <twisted.internet.test.reactormixins.ReactorBuilder>} subclass, as it
- depends on several of its methods.
-
- Then the methods C{connect} and C{listen} must defined, defining a client
- and a server communicating with each other.
- """
-
- def test_interface(self):
- """
- The C{connect} method returns an object providing L{IConnector}.
- """
- reactor = self.buildReactor()
- connector = self.connect(reactor, ClientFactory())
- self.assertTrue(verifyObject(IConnector, connector))
-
- def test_clientConnectionFailedStopsReactor(self):
- """
- The reactor can be stopped by a client factory's
- C{clientConnectionFailed} method.
- """
- reactor = self.buildReactor()
- needsRunningReactor(reactor, lambda: self.connect(reactor, Stop(reactor)))
- self.runReactor(reactor)
-
- def test_connectEvent(self):
- """
- This test checks that we correctly get notifications event for a
- client. This ought to prevent a regression under Windows using the
- GTK2 reactor. See #3925.
- """
- reactor = self.buildReactor()
-
- self.listen(reactor, ServerFactory.forProtocol(Protocol))
- connected = []
-
- class CheckConnection(Protocol):
- def connectionMade(self):
- connected.append(self)
- reactor.stop()
-
- clientFactory = Stop(reactor)
- clientFactory.protocol = CheckConnection
-
- needsRunningReactor(reactor, lambda: self.connect(reactor, clientFactory))
-
- reactor.run()
-
- self.assertTrue(connected)
-
- def test_unregisterProducerAfterDisconnect(self):
- """
- If a producer is unregistered from a transport after the transport has
- been disconnected (by the peer) and after C{loseConnection} has been
- called, the transport is not re-added to the reactor as a writer as
- would be necessary if the transport were still connected.
- """
- reactor = self.buildReactor()
- self.listen(reactor, ServerFactory.forProtocol(ClosingProtocol))
-
- finished = Deferred()
- finished.addErrback(log.err)
- finished.addCallback(lambda ign: reactor.stop())
-
- writing = []
-
- class ClientProtocol(Protocol):
- """
- Protocol to connect, register a producer, try to lose the
- connection, wait for the server to disconnect from us, and then
- unregister the producer.
- """
-
- def connectionMade(self):
- log.msg("ClientProtocol.connectionMade")
- self.transport.registerProducer(
- _SimplePullProducer(self.transport), False
- )
- self.transport.loseConnection()
-
- def connectionLost(self, reason):
- log.msg("ClientProtocol.connectionLost")
- self.unregister()
- writing.append(self.transport in _getWriters(reactor))
- finished.callback(None)
-
- def unregister(self):
- log.msg("ClientProtocol unregister")
- self.transport.unregisterProducer()
-
- clientFactory = ClientFactory()
- clientFactory.protocol = ClientProtocol
- self.connect(reactor, clientFactory)
- self.runReactor(reactor)
- self.assertFalse(writing[0], "Transport was writing after unregisterProducer.")
-
- def test_disconnectWhileProducing(self):
- """
- If C{loseConnection} is called while a producer is registered with the
- transport, the connection is closed after the producer is unregistered.
- """
- reactor = self.buildReactor()
-
- # For some reason, pyobject/pygtk will not deliver the close
- # notification that should happen after the unregisterProducer call in
- # this test. The selectable is in the write notification set, but no
- # notification ever arrives. Probably for the same reason #5233 led
- # win32eventreactor to be broken.
- skippedReactors = ["Glib2Reactor", "Gtk2Reactor"]
- reactorClassName = reactor.__class__.__name__
- if reactorClassName in skippedReactors and platform.isWindows():
- raise SkipTest(
- "A pygobject/pygtk bug disables this functionality " "on Windows."
- )
-
- class Producer:
- def resumeProducing(self):
- log.msg("Producer.resumeProducing")
-
- self.listen(reactor, ServerFactory.forProtocol(Protocol))
-
- finished = Deferred()
- finished.addErrback(log.err)
- finished.addCallback(lambda ign: reactor.stop())
-
- class ClientProtocol(Protocol):
- """
- Protocol to connect, register a producer, try to lose the
- connection, unregister the producer, and wait for the connection to
- actually be lost.
- """
-
- def connectionMade(self):
- log.msg("ClientProtocol.connectionMade")
- self.transport.registerProducer(Producer(), False)
- self.transport.loseConnection()
- # Let the reactor tick over, in case synchronously calling
- # loseConnection and then unregisterProducer is the same as
- # synchronously calling unregisterProducer and then
- # loseConnection (as it is in several reactors).
- reactor.callLater(0, reactor.callLater, 0, self.unregister)
-
- def unregister(self):
- log.msg("ClientProtocol unregister")
- self.transport.unregisterProducer()
- # This should all be pretty quick. Fail the test
- # if we don't get a connectionLost event really
- # soon.
- reactor.callLater(
- 1.0, finished.errback, Failure(Exception("Connection was not lost"))
- )
-
- def connectionLost(self, reason):
- log.msg("ClientProtocol.connectionLost")
- finished.callback(None)
-
- clientFactory = ClientFactory()
- clientFactory.protocol = ClientProtocol
- self.connect(reactor, clientFactory)
- self.runReactor(reactor)
- # If the test failed, we logged an error already and trial
- # will catch it.
|