1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
|
# -*- test-case-name: twisted.test.test_protocols.PortforwardingTests -*-
# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.
"""
A simple port forwarder.
"""
from __future__ import annotations
from twisted.internet import protocol
from twisted.internet.interfaces import (
IAddress,
IConsumer,
IPushProducer,
IStreamServerEndpoint,
ITransport,
)
class Proxy(protocol.Protocol):
noisy = True
peer: Proxy | None = None
factory: protocol.Factory[Proxy]
def setPeer(self, peer: Proxy | None) -> None:
self.peer = peer
def connectionLost(self, reason):
if self.peer is not None:
self.peer.transport.loseConnection()
self.peer = None
def dataReceived(self, data):
self.peer.transport.write(data)
class ProxyClient(Proxy):
factory: protocol.Factory[ProxyClient] # type:ignore
def connectionMade(self) -> None:
assert self.peer is not None
self.peer.setPeer(self)
# Wire this and the peer transport together to enable
# flow control (this stops connections from filling
# this proxy memory when one side produces data at a
# higher rate than the other can consume).
self.transport.registerProducer(self.peer.transport, True) # type:ignore
self.peer.transport.registerProducer(self.transport, True) # type:ignore
# We're connected, everybody can read to their hearts content.
self.peer.transport.resumeProducing() # type:ignore
class ProxyClientFactory(protocol.ClientFactory[ProxyClient]):
protocol = ProxyClient
def setServer(self, server):
self.server = server
def buildProtocol(self, addr: IAddress | None) -> ProxyClient:
prot = super().buildProtocol(addr)
assert prot is not None, "peer must build protocol"
prot.setPeer(self.server)
return prot
def clientConnectionFailed(self, connector, reason):
self.server.transport.loseConnection()
class _MakeTypesHappy(IPushProducer, IConsumer, ITransport):
"""
L{ProxyServer}'s transport is implicitly assumed to provide several
interfaces so include them all here.
"""
class ProxyServer(Proxy):
clientProtocolFactory = ProxyClientFactory
reactor = None
transport: _MakeTypesHappy
factory: ProxyFactory # type:ignore[assignment]
endpoint: IStreamServerEndpoint | None
def connectionMade(self) -> None:
# Don't read anything from the connecting client until we have
# somewhere to send it to.
self.transport.pauseProducing()
client = self.clientProtocolFactory()
client.setServer(self)
if self.reactor is None:
from twisted.internet import reactor
self.reactor = reactor
self.reactor.connectTCP(self.factory.host, self.factory.port, client)
class ProxyFactory(protocol.Factory[ProxyServer]):
"""
Factory for port forwarder.
"""
protocol = ProxyServer
def __init__(self, host: str, port: int) -> None:
self.host = host
self.port = port
|
