1919
2020from twisted .internet import defer , protocol
2121from twisted .internet .error import ConnectError
22- from twisted .internet .interfaces import IStreamClientEndpoint
23- from twisted .internet .protocol import connectionDone
22+ from twisted .internet .interfaces import IReactorCore , IStreamClientEndpoint
23+ from twisted .internet .protocol import ClientFactory , Protocol , connectionDone
2424from twisted .web import http
25+ from twisted .web .http_headers import Headers
2526
2627logger = logging .getLogger (__name__ )
2728
@@ -43,23 +44,33 @@ class HTTPConnectProxyEndpoint:
4344
4445 Args:
4546 reactor: the Twisted reactor to use for the connection
46- proxy_endpoint (IStreamClientEndpoint) : the endpoint to use to connect to the
47- proxy
48- host (bytes): hostname that we want to CONNECT to
49- port (int): port that we want to connect to
47+ proxy_endpoint: the endpoint to use to connect to the proxy
48+ host: hostname that we want to CONNECT to
49+ port: port that we want to connect to
50+ headers: Extra HTTP headers to include in the CONNECT request
5051 """
5152
52- def __init__ (self , reactor , proxy_endpoint , host , port ):
53+ def __init__ (
54+ self ,
55+ reactor : IReactorCore ,
56+ proxy_endpoint : IStreamClientEndpoint ,
57+ host : bytes ,
58+ port : int ,
59+ headers : Headers ,
60+ ):
5361 self ._reactor = reactor
5462 self ._proxy_endpoint = proxy_endpoint
5563 self ._host = host
5664 self ._port = port
65+ self ._headers = headers
5766
5867 def __repr__ (self ):
5968 return "<HTTPConnectProxyEndpoint %s>" % (self ._proxy_endpoint ,)
6069
61- def connect (self , protocolFactory ):
62- f = HTTPProxiedClientFactory (self ._host , self ._port , protocolFactory )
70+ def connect (self , protocolFactory : ClientFactory ):
71+ f = HTTPProxiedClientFactory (
72+ self ._host , self ._port , protocolFactory , self ._headers
73+ )
6374 d = self ._proxy_endpoint .connect (f )
6475 # once the tcp socket connects successfully, we need to wait for the
6576 # CONNECT to complete.
@@ -74,15 +85,23 @@ class HTTPProxiedClientFactory(protocol.ClientFactory):
7485 HTTP Protocol object and run the rest of the connection.
7586
7687 Args:
77- dst_host (bytes): hostname that we want to CONNECT to
78- dst_port (int): port that we want to connect to
79- wrapped_factory (protocol.ClientFactory): The original Factory
88+ dst_host: hostname that we want to CONNECT to
89+ dst_port: port that we want to connect to
90+ wrapped_factory: The original Factory
91+ headers: Extra HTTP headers to include in the CONNECT request
8092 """
8193
82- def __init__ (self , dst_host , dst_port , wrapped_factory ):
94+ def __init__ (
95+ self ,
96+ dst_host : bytes ,
97+ dst_port : int ,
98+ wrapped_factory : ClientFactory ,
99+ headers : Headers ,
100+ ):
83101 self .dst_host = dst_host
84102 self .dst_port = dst_port
85103 self .wrapped_factory = wrapped_factory
104+ self .headers = headers
86105 self .on_connection = defer .Deferred ()
87106
88107 def startedConnecting (self , connector ):
@@ -92,7 +111,11 @@ def buildProtocol(self, addr):
92111 wrapped_protocol = self .wrapped_factory .buildProtocol (addr )
93112
94113 return HTTPConnectProtocol (
95- self .dst_host , self .dst_port , wrapped_protocol , self .on_connection
114+ self .dst_host ,
115+ self .dst_port ,
116+ wrapped_protocol ,
117+ self .on_connection ,
118+ self .headers ,
96119 )
97120
98121 def clientConnectionFailed (self , connector , reason ):
@@ -112,24 +135,37 @@ class HTTPConnectProtocol(protocol.Protocol):
112135 """Protocol that wraps an existing Protocol to do a CONNECT handshake at connect
113136
114137 Args:
115- host (bytes) : The original HTTP(s) hostname or IPv4 or IPv6 address literal
138+ host: The original HTTP(s) hostname or IPv4 or IPv6 address literal
116139 to put in the CONNECT request
117140
118- port (int) : The original HTTP(s) port to put in the CONNECT request
141+ port: The original HTTP(s) port to put in the CONNECT request
119142
120- wrapped_protocol (interfaces.IProtocol) : the original protocol (probably
121- HTTPChannel or TLSMemoryBIOProtocol, but could be anything really)
143+ wrapped_protocol: the original protocol (probably HTTPChannel or
144+ TLSMemoryBIOProtocol, but could be anything really)
122145
123- connected_deferred (Deferred) : a Deferred which will be callbacked with
146+ connected_deferred: a Deferred which will be callbacked with
124147 wrapped_protocol when the CONNECT completes
148+
149+ headers: Extra HTTP headers to include in the CONNECT request
125150 """
126151
127- def __init__ (self , host , port , wrapped_protocol , connected_deferred ):
152+ def __init__ (
153+ self ,
154+ host : bytes ,
155+ port : int ,
156+ wrapped_protocol : Protocol ,
157+ connected_deferred : defer .Deferred ,
158+ headers : Headers ,
159+ ):
128160 self .host = host
129161 self .port = port
130162 self .wrapped_protocol = wrapped_protocol
131163 self .connected_deferred = connected_deferred
132- self .http_setup_client = HTTPConnectSetupClient (self .host , self .port )
164+ self .headers = headers
165+
166+ self .http_setup_client = HTTPConnectSetupClient (
167+ self .host , self .port , self .headers
168+ )
133169 self .http_setup_client .on_connected .addCallback (self .proxyConnected )
134170
135171 def connectionMade (self ):
@@ -154,7 +190,7 @@ def proxyConnected(self, _):
154190 if buf :
155191 self .wrapped_protocol .dataReceived (buf )
156192
157- def dataReceived (self , data ):
193+ def dataReceived (self , data : bytes ):
158194 # if we've set up the HTTP protocol, we can send the data there
159195 if self .wrapped_protocol .connected :
160196 return self .wrapped_protocol .dataReceived (data )
@@ -168,21 +204,29 @@ class HTTPConnectSetupClient(http.HTTPClient):
168204 """HTTPClient protocol to send a CONNECT message for proxies and read the response.
169205
170206 Args:
171- host (bytes): The hostname to send in the CONNECT message
172- port (int): The port to send in the CONNECT message
207+ host: The hostname to send in the CONNECT message
208+ port: The port to send in the CONNECT message
209+ headers: Extra headers to send with the CONNECT message
173210 """
174211
175- def __init__ (self , host , port ):
212+ def __init__ (self , host : bytes , port : int , headers : Headers ):
176213 self .host = host
177214 self .port = port
215+ self .headers = headers
178216 self .on_connected = defer .Deferred ()
179217
180218 def connectionMade (self ):
181219 logger .debug ("Connected to proxy, sending CONNECT" )
182220 self .sendCommand (b"CONNECT" , b"%s:%d" % (self .host , self .port ))
221+
222+ # Send any additional specified headers
223+ for name , values in self .headers .getAllRawHeaders ():
224+ for value in values :
225+ self .sendHeader (name , value )
226+
183227 self .endHeaders ()
184228
185- def handleStatus (self , version , status , message ):
229+ def handleStatus (self , version : bytes , status : bytes , message : bytes ):
186230 logger .debug ("Got Status: %s %s %s" , status , message , version )
187231 if status != b"200" :
188232 raise ProxyConnectError ("Unexpected status on CONNECT: %s" % status )
0 commit comments