Add a timeout.

This commit is contained in:
Itamar Turner-Trauring 2022-06-23 12:47:33 -04:00
parent 0c99a9f7b0
commit eb1e48bcc3

View File

@ -2,13 +2,14 @@
Support for listening with both HTTP and Foolscap on the same port. Support for listening with both HTTP and Foolscap on the same port.
""" """
from typing import Optional, Tuple from typing import Optional
from twisted.internet.protocol import Protocol from twisted.internet.protocol import Protocol
from twisted.internet.interfaces import ITransport from twisted.internet.interfaces import IDelayedCall
from twisted.internet.ssl import CertificateOptions, PrivateCertificate from twisted.internet.ssl import CertificateOptions, PrivateCertificate
from twisted.web.server import Site from twisted.web.server import Site
from twisted.protocols.tls import TLSMemoryBIOFactory from twisted.protocols.tls import TLSMemoryBIOFactory
from twisted.internet import reactor
from foolscap.negotiate import Negotiation from foolscap.negotiate import Negotiation
@ -36,6 +37,8 @@ class FoolscapOrHttp(Protocol, metaclass=PretendToBeNegotiation):
certificate: PrivateCertificate certificate: PrivateCertificate
storage_server: StorageServer storage_server: StorageServer
_timeout: IDelayedCall
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
self._foolscap: Negotiation = Negotiation(*args, **kwargs) self._foolscap: Negotiation = Negotiation(*args, **kwargs)
self._buffer: bytes = b"" self._buffer: bytes = b""
@ -69,6 +72,9 @@ class FoolscapOrHttp(Protocol, metaclass=PretendToBeNegotiation):
self._convert_to_negotiation() self._convert_to_negotiation()
return self.initClient(*args, **kwargs) return self.initClient(*args, **kwargs)
def connectionMade(self):
self._timeout = reactor.callLater(30, self.transport.abortConnection)
def dataReceived(self, data: bytes) -> None: def dataReceived(self, data: bytes) -> None:
"""Handle incoming data. """Handle incoming data.
@ -80,7 +86,8 @@ class FoolscapOrHttp(Protocol, metaclass=PretendToBeNegotiation):
return return
# Check if it looks like a Foolscap request. If so, it can handle this # Check if it looks like a Foolscap request. If so, it can handle this
# and later data: # and later data, otherwise assume HTTPS.
self._timeout.cancel()
if self._buffer.startswith(b"GET /id/"): if self._buffer.startswith(b"GET /id/"):
transport = self.transport transport = self.transport
buf = self._buffer buf = self._buffer