Viewing file: tcp.py (19.71 KB) -rw-r--r-- Select action/file-type: (+) | (+) | (+) | Code (+) | Session (+) | (+) | SDB (+) | (+) | (+) | (+) | (+) | (+) |
# Copyright (c) 2008-2009 Twisted Matrix Laboratories. # See LICENSE for details.
""" TCP support for IOCP reactor """
import socket, operator, errno, struct
from zope.interface import implements, directlyProvides
from twisted.internet import interfaces, error, address, main, defer from twisted.internet.abstract import isIPAddress from twisted.internet.tcp import _SocketCloser, Connector as TCPConnector from twisted.persisted import styles from twisted.python import log, failure, reflect, util
from twisted.internet.iocpreactor import iocpsupport as _iocp, abstract from twisted.internet.iocpreactor.interfaces import IReadWriteHandle from twisted.internet.iocpreactor.const import ERROR_IO_PENDING from twisted.internet.iocpreactor.const import SO_UPDATE_CONNECT_CONTEXT from twisted.internet.iocpreactor.const import SO_UPDATE_ACCEPT_CONTEXT from twisted.internet.iocpreactor.const import ERROR_CONNECTION_REFUSED from twisted.internet.iocpreactor.const import ERROR_NETWORK_UNREACHABLE
try: from twisted.protocols.tls import TLSMemoryBIOFactory, TLSMemoryBIOProtocol except ImportError: TLSMemoryBIOProtocol = TLSMemoryBIOFactory = None _extraInterfaces = () else: _extraInterfaces = (interfaces.ITLSTransport,)
# ConnectEx returns these. XXX: find out what it does for timeout connectExErrors = { ERROR_CONNECTION_REFUSED: errno.WSAECONNREFUSED, ERROR_NETWORK_UNREACHABLE: errno.WSAENETUNREACH, }
class _BypassTLS(object): """ L{_BypassTLS} is used as the transport object for the TLS protocol object used to implement C{startTLS}. Its methods skip any TLS logic which C{startTLS} enables.
@ivar _connection: A L{Connection} which TLS has been started on which will be proxied to by this object. Any method which has its behavior altered after C{startTLS} will be skipped in favor of the base class's implementation. This allows the TLS protocol object to have direct access to the transport, necessary to actually implement TLS. """ def __init__(self, connection): self._connection = connection
def __getattr__(self, name): return getattr(self._connection, name)
def write(self, data): return abstract.FileHandle.write(self._connection, data)
def writeSequence(self, iovec): return abstract.FileHandle.writeSequence(self._connection, iovec)
def loseConnection(self, reason=None): return abstract.FileHandle.loseConnection(self._connection, reason)
class Connection(abstract.FileHandle, _SocketCloser): """ @ivar _tls: C{False} to indicate the connection is in normal TCP mode, C{True} to indicate that TLS has been started and that operations must be routed through the L{TLSMemoryBIOProtocol} instance.
@ivar _tlsClientDefault: A flag which must be set by a subclass. If set to C{True}, L{startTLS} will default to initiating SSL as a client. If set to C{False}, L{startTLS} will default to initiating SSL as a server. """ implements(IReadWriteHandle, interfaces.ITCPTransport, interfaces.ISystemHandle, *_extraInterfaces)
_tls = False
def __init__(self, sock, proto, reactor=None): abstract.FileHandle.__init__(self, reactor) self.socket = sock self.getFileHandle = sock.fileno self.protocol = proto
def getHandle(self): return self.socket
def dataReceived(self, rbuffer): # XXX: some day, we'll have protocols that can handle raw buffers self.protocol.dataReceived(str(rbuffer))
def readFromHandle(self, bufflist, evt): return _iocp.recv(self.getFileHandle(), bufflist, evt)
def writeToHandle(self, buff, evt): return _iocp.send(self.getFileHandle(), buff, evt)
def _closeWriteConnection(self): try: getattr(self.socket, self._socketShutdownMethod)(1) except socket.error: pass p = interfaces.IHalfCloseableProtocol(self.protocol, None) if p: try: p.writeConnectionLost() except: f = failure.Failure() log.err() self.connectionLost(f)
def readConnectionLost(self, reason): p = interfaces.IHalfCloseableProtocol(self.protocol, None) if p: try: p.readConnectionLost() except: log.err() self.connectionLost(failure.Failure()) else: self.connectionLost(reason)
def connectionLost(self, reason): abstract.FileHandle.connectionLost(self, reason) self._closeSocket() protocol = self.protocol del self.protocol del self.socket del self.getFileHandle protocol.connectionLost(reason)
def logPrefix(self): """ Return the prefix to log with when I own the logging thread. """ return self.logstr
def getTcpNoDelay(self): return operator.truth(self.socket.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY))
def setTcpNoDelay(self, enabled): self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, enabled)
def getTcpKeepAlive(self): return operator.truth(self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE))
def setTcpKeepAlive(self, enabled): self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, enabled)
if TLSMemoryBIOFactory is not None: def startTLS(self, contextFactory, normal=True): """ @see: L{ITLSTransport.startTLS} """ # Figure out which direction the SSL goes in. If normal is True, # we'll go in the direction indicated by the subclass. Otherwise, # we'll go the other way (client = not normal ^ _tlsClientDefault, # in other words). if normal: client = self._tlsClientDefault else: client = not self._tlsClientDefault
tlsFactory = TLSMemoryBIOFactory(contextFactory, client, None) tlsProtocol = TLSMemoryBIOProtocol(tlsFactory, self.protocol, False) self.protocol = tlsProtocol
self.getHandle = tlsProtocol.getHandle self.getPeerCertificate = tlsProtocol.getPeerCertificate
# Mark the transport as secure. directlyProvides(self, interfaces.ISSLTransport)
# Remember we did this so that write and writeSequence can send the # data to the right place. self._tls = True
# Hook it up self.protocol.makeConnection(_BypassTLS(self))
def write(self, data): """ Write some data, either directly to the underlying handle or, if TLS has been started, to the L{TLSMemoryBIOProtocol} for it to encrypt and send.
@see: L{ITCPTransport.write} """ if self._tls: self.protocol.write(data) else: abstract.FileHandle.write(self, data)
def writeSequence(self, iovec): """ Write some data, either directly to the underlying handle or, if TLS has been started, to the L{TLSMemoryBIOProtocol} for it to encrypt and send.
@see: L{ITCPTransport.writeSequence} """ if self._tls: self.protocol.writeSequence(iovec) else: abstract.FileHandle.writeSequence(self, iovec)
def loseConnection(self, reason=None): """ Close the underlying handle or, if TLS has been started, first shut it down.
@see: L{ITCPTransport.loseConnection} """ if self._tls: if self.connected and not self.disconnecting: self.protocol.loseConnection() else: abstract.FileHandle.loseConnection(self, reason)
class Client(Connection): addressFamily = socket.AF_INET socketType = socket.SOCK_STREAM
_tlsClientDefault = True
def __init__(self, host, port, bindAddress, connector, reactor): self.connector = connector self.addr = (host, port) self.reactor = reactor # ConnectEx documentation says socket _has_ to be bound if bindAddress is None: bindAddress = ('', 0)
try: try: skt = reactor.createSocket(self.addressFamily, self.socketType) except socket.error, se: raise error.ConnectBindError(se[0], se[1]) else: try: skt.bind(bindAddress) except socket.error, se: raise error.ConnectBindError(se[0], se[1]) self.socket = skt Connection.__init__(self, skt, None, reactor) reactor.callLater(0, self.resolveAddress) except error.ConnectBindError, err: reactor.callLater(0, self.failIfNotConnected, err)
def resolveAddress(self): if isIPAddress(self.addr[0]): self._setRealAddress(self.addr[0]) else: d = self.reactor.resolve(self.addr[0]) d.addCallbacks(self._setRealAddress, self.failIfNotConnected)
def _setRealAddress(self, address): self.realAddress = (address, self.addr[1]) self.doConnect()
def failIfNotConnected(self, err): if (self.connected or self.disconnected or not hasattr(self, "connector")): return
try: self._closeSocket() except AttributeError: pass else: del self.socket, self.getFileHandle self.reactor.removeActiveHandle(self)
self.connector.connectionFailed(failure.Failure(err)) del self.connector
def stopConnecting(self): """ Stop attempt to connect. """ self.failIfNotConnected(error.UserError())
def cbConnect(self, rc, bytes, evt): if rc: rc = connectExErrors.get(rc, rc) self.failIfNotConnected(error.getConnectError((rc, errno.errorcode.get(rc, 'Unknown error')))) else: self.socket.setsockopt(socket.SOL_SOCKET, SO_UPDATE_CONNECT_CONTEXT, struct.pack('I', self.socket.fileno())) self.protocol = self.connector.buildProtocol(self.getPeer()) self.connected = True self.logstr = self.protocol.__class__.__name__+",client" self.protocol.makeConnection(self) self.startReading()
def doConnect(self): if not hasattr(self, "connector"): # this happens if we connector.stopConnecting in # factory.startedConnecting return assert _iocp.have_connectex self.reactor.addActiveHandle(self) evt = _iocp.Event(self.cbConnect, self)
rc = _iocp.connect(self.socket.fileno(), self.realAddress, evt) if rc == ERROR_IO_PENDING: return else: evt.ignore = True self.cbConnect(rc, 0, 0, evt)
def getHost(self): """ Returns an IPv4Address.
This indicates the address from which I am connecting. """ return address.IPv4Address('TCP', *(self.socket.getsockname() + ('INET',)))
def getPeer(self): """ Returns an IPv4Address.
This indicates the address that I am connected to. """ return address.IPv4Address('TCP', *(self.realAddress + ('INET',)))
def __repr__(self): s = ('<%s to %s at %x>' % (self.__class__, self.addr, util.unsignedID(self))) return s
def connectionLost(self, reason): if not self.connected: self.failIfNotConnected(error.ConnectError(string=reason)) else: Connection.connectionLost(self, reason) self.connector.connectionLost(reason)
class Server(Connection): """ Serverside socket-stream connection class.
I am a serverside network connection transport; a socket which came from an accept() on a server. """
_tlsClientDefault = False
def __init__(self, sock, protocol, clientAddr, serverAddr, sessionno, reactor): """ Server(sock, protocol, client, server, sessionno)
Initialize me with a socket, a protocol, a descriptor for my peer (a tuple of host, port describing the other end of the connection), an instance of Port, and a session number. """ Connection.__init__(self, sock, protocol, reactor) self.serverAddr = serverAddr self.clientAddr = clientAddr self.sessionno = sessionno self.logstr = "%s,%s,%s" % (self.protocol.__class__.__name__, sessionno, self.clientAddr.host) self.repstr = "<%s #%s on %s>" % (self.protocol.__class__.__name__, self.sessionno, self.serverAddr.port) self.connected = True self.startReading()
def __repr__(self): """ A string representation of this connection. """ return self.repstr
def getHost(self): """ Returns an IPv4Address.
This indicates the server's address. """ return self.serverAddr
def getPeer(self): """ Returns an IPv4Address.
This indicates the client's address. """ return self.clientAddr
class Connector(TCPConnector): def _makeTransport(self): return Client(self.host, self.port, self.bindAddress, self, self.reactor)
class Port(styles.Ephemeral, _SocketCloser): implements(interfaces.IListeningPort)
connected = False disconnected = False disconnecting = False addressFamily = socket.AF_INET socketType = socket.SOCK_STREAM
sessionno = 0
maxAccepts = 100
# Actual port number being listened on, only set to a non-None # value when we are actually listening. _realPortNumber = None
def __init__(self, port, factory, backlog=50, interface='', reactor=None): self.port = port self.factory = factory self.backlog = backlog self.interface = interface self.reactor = reactor
def __repr__(self): if self._realPortNumber is not None: return "<%s of %s on %s>" % (self.__class__, self.factory.__class__, self._realPortNumber) else: return "<%s of %s (not listening)>" % (self.__class__, self.factory.__class__)
def startListening(self): try: skt = self.reactor.createSocket(self.addressFamily, self.socketType) # TODO: resolve self.interface if necessary skt.bind((self.interface, self.port)) except socket.error, le: raise error.CannotListenError, (self.interface, self.port, le)
self.addrLen = _iocp.maxAddrLen(skt.fileno())
# Make sure that if we listened on port 0, we update that to # reflect what the OS actually assigned us. self._realPortNumber = skt.getsockname()[1]
log.msg("%s starting on %s" % (self.factory.__class__, self._realPortNumber))
self.factory.doStart() skt.listen(self.backlog) self.connected = True self.disconnected = False self.reactor.addActiveHandle(self) self.socket = skt self.getFileHandle = self.socket.fileno self.doAccept()
def loseConnection(self, connDone=failure.Failure(main.CONNECTION_DONE)): """ Stop accepting connections on this port.
This will shut down my socket and call self.connectionLost(). It returns a deferred which will fire successfully when the port is actually closed. """ self.disconnecting = True if self.connected: self.deferred = defer.Deferred() self.reactor.callLater(0, self.connectionLost, connDone) return self.deferred
stopListening = loseConnection
def connectionLost(self, reason): """ Cleans up the socket. """ log.msg('(Port %s Closed)' % self._realPortNumber) self._realPortNumber = None d = None if hasattr(self, "deferred"): d = self.deferred del self.deferred
self.disconnected = True self.reactor.removeActiveHandle(self) self.connected = False self._closeSocket() del self.socket del self.getFileHandle
try: self.factory.doStop() except: self.disconnecting = False if d is not None: d.errback(failure.Failure()) else: raise else: self.disconnecting = False if d is not None: d.callback(None)
def logPrefix(self): """ Returns the name of my class, to prefix log entries with. """ return reflect.qual(self.factory.__class__)
def getHost(self): """ Returns an IPv4Address.
This indicates the server's address. """ return address.IPv4Address('TCP', *(self.socket.getsockname() + ('INET',)))
def cbAccept(self, rc, bytes, evt): self.handleAccept(rc, evt) if not (self.disconnecting or self.disconnected): self.doAccept()
def handleAccept(self, rc, evt): if self.disconnecting or self.disconnected: return False
# possible errors: # (WSAEMFILE, WSAENOBUFS, WSAENFILE, WSAENOMEM, WSAECONNABORTED) if rc: log.msg("Could not accept new connection -- %s (%s)" % (errno.errorcode.get(rc, 'unknown error'), rc)) return False else: evt.newskt.setsockopt(socket.SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT, struct.pack('I', self.socket.fileno())) family, lAddr, rAddr = _iocp.get_accept_addrs(evt.newskt.fileno(), evt.buff) assert family == self.addressFamily
protocol = self.factory.buildProtocol( address._ServerFactoryIPv4Address('TCP', rAddr[0], rAddr[1])) if protocol is None: evt.newskt.close() else: s = self.sessionno self.sessionno = s+1 transport = Server(evt.newskt, protocol, address.IPv4Address('TCP', rAddr[0], rAddr[1], 'INET'), address.IPv4Address('TCP', lAddr[0], lAddr[1], 'INET'), s, self.reactor) protocol.makeConnection(transport) return True
def doAccept(self): numAccepts = 0 while 1: evt = _iocp.Event(self.cbAccept, self)
# see AcceptEx documentation evt.buff = buff = _iocp.AllocateReadBuffer(2 * (self.addrLen + 16))
evt.newskt = newskt = self.reactor.createSocket(self.addressFamily, self.socketType) rc = _iocp.accept(self.socket.fileno(), newskt.fileno(), buff, evt)
if (rc == ERROR_IO_PENDING or (not rc and numAccepts >= self.maxAccepts)): break else: evt.ignore = True if not self.handleAccept(rc, evt): break numAccepts += 1
|