From 10d6a23eb2b9e8193bf844a780c7c26e101567c5 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 12:31:17 -0500
Subject: [PATCH 1/7] Fix remaining mypy issues in tests.

---
 tests/replication/_base.py | 40 +++++++++++++-------------------------
 tests/server.py            |  2 ++
 2 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/tests/replication/_base.py b/tests/replication/_base.py
index 20940c81071e..d95ea7b1ecf7 100644
--- a/tests/replication/_base.py
+++ b/tests/replication/_base.py
@@ -13,9 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-import attr
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 from twisted.internet.interfaces import IConsumer, IPullProducer, IReactorTime
 from twisted.internet.protocol import Protocol
@@ -158,10 +156,8 @@ def handle_http_replication_attempt(self) -> SynapseRequest:
         # Set up client side protocol
         client_protocol = client_factory.buildProtocol(None)
 
-        request_factory = OneShotRequestFactory()
-
         # Set up the server side protocol
-        channel = _PushHTTPChannel(self.reactor, request_factory, self.site)
+        channel = _PushHTTPChannel(self.reactor, SynapseRequest, self.site)
 
         # Connect client to server and vice versa.
         client_to_server_transport = FakeTransport(
@@ -183,7 +179,7 @@ def handle_http_replication_attempt(self) -> SynapseRequest:
         server_to_client_transport.loseConnection()
         client_to_server_transport.loseConnection()
 
-        return request_factory.request
+        return channel.request
 
     def assert_request_is_get_repl_stream_updates(
         self, request: SynapseRequest, stream_name: str
@@ -392,10 +388,8 @@ def _handle_http_replication_attempt(self, hs, repl_port):
         # Set up client side protocol
         client_protocol = client_factory.buildProtocol(None)
 
-        request_factory = OneShotRequestFactory()
-
         # Set up the server side protocol
-        channel = _PushHTTPChannel(self.reactor, request_factory, self._hs_to_site[hs])
+        channel = _PushHTTPChannel(self.reactor, SynapseRequest, self._hs_to_site[hs])
 
         # Connect client to server and vice versa.
         client_to_server_transport = FakeTransport(
@@ -453,21 +447,6 @@ async def on_rdata(self, stream_name, instance_name, token, rows):
             self.received_rdata_rows.append((stream_name, token, r))
 
 
-@attr.s()
-class OneShotRequestFactory:
-    """A simple request factory that generates a single `SynapseRequest` and
-    stores it for future use. Can only be used once.
-    """
-
-    request = attr.ib(default=None)
-
-    def __call__(self, *args, **kwargs):
-        assert self.request is None
-
-        self.request = SynapseRequest(*args, **kwargs)
-        return self.request
-
-
 class _PushHTTPChannel(HTTPChannel):
     """A HTTPChannel that wraps pull producers to push producers.
 
@@ -479,7 +458,7 @@ class _PushHTTPChannel(HTTPChannel):
     """
 
     def __init__(
-        self, reactor: IReactorTime, request_factory: Callable[..., Request], site: Site
+        self, reactor: IReactorTime, request_factory: Type[Request], site: Site
     ):
         super().__init__()
         self.reactor = reactor
@@ -510,6 +489,11 @@ def checkPersistence(self, request, version):
         request.responseHeaders.setRawHeaders(b"connection", [b"close"])
         return False
 
+    def requestDone(self, request):
+        # Store the request for inspection.
+        self.request = request
+        super().requestDone(request)
+
 
 class _PullToPushProducer:
     """A push producer that wraps a pull producer."""
@@ -597,6 +581,8 @@ def buildProtocol(self, addr):
 class FakeRedisPubSubProtocol(Protocol):
     """A connection from a client talking to the fake Redis server."""
 
+    transport = None  # type: Optional[FakeTransport]
+
     def __init__(self, server: FakeRedisPubSubServer):
         self._server = server
         self._reader = hiredis.Reader()
@@ -641,6 +627,8 @@ def handle_command(self, command, *args):
 
     def send(self, msg):
         """Send a message back to the client."""
+        assert self.transport is not None
+
         raw = self.encode(msg).encode("utf-8")
 
         self.transport.write(raw)
diff --git a/tests/server.py b/tests/server.py
index 863f6da7385c..2287d200763b 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -16,6 +16,7 @@
     IReactorPluggableNameResolver,
     IReactorTCP,
     IResolverSimple,
+    ITransport,
 )
 from twisted.python.failure import Failure
 from twisted.test.proto_helpers import AccumulatingProtocol, MemoryReactorClock
@@ -467,6 +468,7 @@ def get_clock():
     return clock, hs_clock
 
 
+@implementer(ITransport)
 @attr.s(cmp=False)
 class FakeTransport:
     """

From 292d53d8feddfaec99a3bb398143964807c492a8 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:20:58 -0500
Subject: [PATCH 2/7] Fix issues with a transport being optional / the wrong
 type.

---
 synapse/http/client.py              | 7 +++++++
 synapse/replication/tcp/protocol.py | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/synapse/http/client.py b/synapse/http/client.py
index d4ab3a273261..cb5f9948c01a 100644
--- a/synapse/http/client.py
+++ b/synapse/http/client.py
@@ -45,6 +45,7 @@
     IHostResolution,
     IReactorPluggableNameResolver,
     IResolutionReceiver,
+    ITCPTransport,
 )
 from twisted.internet.task import Cooperator
 from twisted.python.failure import Failure
@@ -760,6 +761,8 @@ class BodyExceededMaxSize(Exception):
 class _DiscardBodyWithMaxSizeProtocol(protocol.Protocol):
     """A protocol which immediately errors upon receiving data."""
 
+    transport = None  # type: Optional[ITCPTransport]
+
     def __init__(self, deferred: defer.Deferred):
         self.deferred = deferred
 
@@ -771,6 +774,7 @@ def _maybe_fail(self):
             self.deferred.errback(BodyExceededMaxSize())
             # Close the connection (forcefully) since all the data will get
             # discarded anyway.
+            assert self.transport is not None
             self.transport.abortConnection()
 
     def dataReceived(self, data: bytes) -> None:
@@ -783,6 +787,8 @@ def connectionLost(self, reason: Failure) -> None:
 class _ReadBodyWithMaxSizeProtocol(protocol.Protocol):
     """A protocol which reads body to a stream, erroring if the body exceeds a maximum size."""
 
+    transport = None  # type: Optional[ITCPTransport]
+
     def __init__(
         self, stream: BinaryIO, deferred: defer.Deferred, max_size: Optional[int]
     ):
@@ -805,6 +811,7 @@ def dataReceived(self, data: bytes) -> None:
             self.deferred.errback(BodyExceededMaxSize())
             # Close the connection (forcefully) since all the data will get
             # discarded anyway.
+            assert self.transport is not None
             self.transport.abortConnection()
 
     def connectionLost(self, reason: Failure) -> None:
diff --git a/synapse/replication/tcp/protocol.py b/synapse/replication/tcp/protocol.py
index 8e4734b59cca..408c60af0a20 100644
--- a/synapse/replication/tcp/protocol.py
+++ b/synapse/replication/tcp/protocol.py
@@ -56,6 +56,7 @@
 from zope.interface import Interface, implementer
 
 from twisted.internet import task
+from twisted.internet.tcp import Connection
 from twisted.protocols.basic import LineOnlyReceiver
 from twisted.python.failure import Failure
 
@@ -145,6 +146,10 @@ class BaseReplicationStreamProtocol(LineOnlyReceiver):
     (if they send a `PING` command)
     """
 
+    # The transport is going to be an ITCPTransport, but that doesn't have the
+    # (un)registerProducer methods, those are only on the implementation.
+    transport = None  # type: Connection
+
     delimiter = b"\n"
 
     # Valid commands we expect to receive
@@ -189,6 +194,7 @@ def connectionMade(self):
 
         connected_connections.append(self)  # Register connection for metrics
 
+        assert self.transport is not None
         self.transport.registerProducer(self, True)  # For the *Producing callbacks
 
         self._send_pending_commands()
@@ -213,6 +219,7 @@ def send_ping(self):
                 logger.info(
                     "[%s] Failed to close connection gracefully, aborting", self.id()
                 )
+                assert self.transport is not None
                 self.transport.abortConnection()
         else:
             if now - self.last_sent_command >= PING_TIME:
@@ -302,6 +309,7 @@ def handle_command(self, cmd: Command) -> None:
     def close(self):
         logger.warning("[%s] Closing connection", self.id())
         self.time_we_closed = self.clock.time_msec()
+        assert self.transport is not None
         self.transport.loseConnection()
         self.on_connection_closed()
 

From 027f9fec3d6088f0a895b0aacf328f7d1354ad03 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:22:30 -0500
Subject: [PATCH 3/7] Fix Redis stubs to inherit from Protocol.

---
 stubs/txredisapi.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stubs/txredisapi.pyi b/stubs/txredisapi.pyi
index 34787e0b1ecb..080ca40287d2 100644
--- a/stubs/txredisapi.pyi
+++ b/stubs/txredisapi.pyi
@@ -19,7 +19,7 @@ from typing import Any, List, Optional, Type, Union
 
 from twisted.internet import protocol
 
-class RedisProtocol:
+class RedisProtocol(protocol.Protocol):
     def publish(self, channel: str, message: bytes): ...
     async def ping(self) -> None: ...
     async def set(

From 91fe968b62c050e975ae9655deb2fa3f9b0a0944 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:22:57 -0500
Subject: [PATCH 4/7] The host is expected to be bytes.

---
 synapse/replication/tcp/handler.py | 4 ++--
 synapse/replication/tcp/redis.py   | 2 +-
 tests/replication/_base.py         | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/synapse/replication/tcp/handler.py b/synapse/replication/tcp/handler.py
index ee909f3fc5da..a8894beadfd1 100644
--- a/synapse/replication/tcp/handler.py
+++ b/synapse/replication/tcp/handler.py
@@ -302,7 +302,7 @@ def start_replication(self, hs):
                 hs, outbound_redis_connection
             )
             hs.get_reactor().connectTCP(
-                hs.config.redis.redis_host,
+                hs.config.redis.redis_host.encode(),
                 hs.config.redis.redis_port,
                 self._factory,
             )
@@ -311,7 +311,7 @@ def start_replication(self, hs):
             self._factory = DirectTcpReplicationClientFactory(hs, client_name, self)
             host = hs.config.worker_replication_host
             port = hs.config.worker_replication_port
-            hs.get_reactor().connectTCP(host, port, self._factory)
+            hs.get_reactor().connectTCP(host.encode(), port, self._factory)
 
     def get_streams(self) -> Dict[str, Stream]:
         """Get a map from stream name to all streams."""
diff --git a/synapse/replication/tcp/redis.py b/synapse/replication/tcp/redis.py
index 7cccde097de7..2f4d407f9482 100644
--- a/synapse/replication/tcp/redis.py
+++ b/synapse/replication/tcp/redis.py
@@ -365,6 +365,6 @@ def lazyConnection(
     factory.continueTrying = reconnect
 
     reactor = hs.get_reactor()
-    reactor.connectTCP(host, port, factory, timeout=30, bindAddress=None)
+    reactor.connectTCP(host.encode(), port, factory, timeout=30, bindAddress=None)
 
     return factory.handler
diff --git a/tests/replication/_base.py b/tests/replication/_base.py
index d95ea7b1ecf7..67b7913666fc 100644
--- a/tests/replication/_base.py
+++ b/tests/replication/_base.py
@@ -233,7 +233,7 @@ def setUp(self):
         if self.hs.config.redis.redis_enabled:
             # Handle attempts to connect to fake redis server.
             self.reactor.add_tcp_client_callback(
-                "localhost",
+                b"localhost",
                 6379,
                 self.connect_any_redis_attempts,
             )
@@ -415,7 +415,7 @@ def connect_any_redis_attempts(self):
         clients = self.reactor.tcpClients
         while clients:
             (host, port, client_factory, _timeout, _bindAddress) = clients.pop(0)
-            self.assertEqual(host, "localhost")
+            self.assertEqual(host, b"localhost")
             self.assertEqual(port, 6379)
 
             client_protocol = client_factory.buildProtocol(None)

From 63e5e33d0abc6135f0a939a1403606d1c90b7169 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:23:15 -0500
Subject: [PATCH 5/7] A Failure's type is always set.

---
 synapse/replication/tcp/protocol.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/synapse/replication/tcp/protocol.py b/synapse/replication/tcp/protocol.py
index 408c60af0a20..825900f64c1d 100644
--- a/synapse/replication/tcp/protocol.py
+++ b/synapse/replication/tcp/protocol.py
@@ -407,6 +407,7 @@ def stopProducing(self):
     def connectionLost(self, reason):
         logger.info("[%s] Replication connection closed: %r", self.id(), reason)
         if isinstance(reason, Failure):
+            assert reason.type is not None
             connection_close_counter.labels(reason.type.__name__).inc()
         else:
             connection_close_counter.labels(reason.__class__.__name__).inc()

From 06d09af4df36b5b09e7b61cf58d4d5a71c4212a0 Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:23:26 -0500
Subject: [PATCH 6/7] Match the parent signature for connectionLost.

---
 synapse/http/client.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/synapse/http/client.py b/synapse/http/client.py
index cb5f9948c01a..1e01e0a9f2b5 100644
--- a/synapse/http/client.py
+++ b/synapse/http/client.py
@@ -47,6 +47,7 @@
     IResolutionReceiver,
     ITCPTransport,
 )
+from twisted.internet.protocol import connectionDone
 from twisted.internet.task import Cooperator
 from twisted.python.failure import Failure
 from twisted.web._newclient import ResponseDone
@@ -780,7 +781,7 @@ def _maybe_fail(self):
     def dataReceived(self, data: bytes) -> None:
         self._maybe_fail()
 
-    def connectionLost(self, reason: Failure) -> None:
+    def connectionLost(self, reason: Failure = connectionDone) -> None:
         self._maybe_fail()
 
 
@@ -814,7 +815,7 @@ def dataReceived(self, data: bytes) -> None:
             assert self.transport is not None
             self.transport.abortConnection()
 
-    def connectionLost(self, reason: Failure) -> None:
+    def connectionLost(self, reason: Failure = connectionDone) -> None:
         # If the maximum size was already exceeded, there's nothing to do.
         if self.deferred.called:
             return

From f9b682899b89d81cdd5131ec0130538b084a3f8a Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Fri, 12 Mar 2021 13:24:21 -0500
Subject: [PATCH 7/7] Newsfragment

---
 changelog.d/9608.misc | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/9608.misc

diff --git a/changelog.d/9608.misc b/changelog.d/9608.misc
new file mode 100644
index 000000000000..14c7b78dd97e
--- /dev/null
+++ b/changelog.d/9608.misc
@@ -0,0 +1 @@
+Fix incorrect type hints.