habedi · September 10, 2015 09:02 · Oct 8, 2012 · Mar 18, 2012 · Mar 18, 2012
diff --git a/server.py b/server.py
@@ -81,6 +81,7 @@ def connectionMade(self):
         self.factory.request.setResponseCode(200, "CONNECT OK")
         self.factory.request.setHeader('X-Connected-IP',
             self.transport.realAddress[0])
+        self.factory.request.setHeader('Content-Length', '0')
         self.factory.request.finish()
 
     def connectionLost(self, reason):
@@ -124,4 +125,4 @@ def clientConnectionFailed(self, connector, reason):
 
     import twisted.internet
     c = twisted.internet.reactor.listenTCP(ns.port, factory)
-    reactor.run()
+    twisted.internet.reactor.run()
diff --git a/server.py b/server.py
@@ -42,11 +42,6 @@ def processConnectRequest(self):
                       "Unable to parse port from URI: %s" % self.uri)
             return
 
-        # Copy headers from upstream reques
-        headers = self.getAllHeaders().copy()
-        if 'host' not in headers:            
-            headers['host'] = host           
-
         clientFactory = ConnectProxyClientFactory(host, port, self)
 
         # TODO provide an API to set proxy connect timeouts

diff --git a/server.py b/server.py
@@ -0,0 +1,132 @@
+from twisted.web.proxy import Proxy, ProxyRequest
+from twisted.internet.protocol import Protocol, ClientFactory
+import urlparse
+from twisted.python import log
+
+
+class ConnectProxyRequest(ProxyRequest):
+    """HTTP ProxyRequest handler (factory) that supports CONNECT"""
+
+    connectedProtocol = None
+
+    def process(self):
+        if self.method == 'CONNECT':
+            self.processConnectRequest()
+        else:
+            ProxyRequest.process(self)
+
+    def fail(self, message, body):
+        self.setResponseCode(501, message)
+        self.responseHeaders.addRawHeader("Content-Type", "text/html")
+        self.write(body)
+        self.finish()                                
+
+    def splitHostPort(self, hostport, default_port):
+        port = default_port
+        parts = hostport.split(':', 1)
+        if len(parts) == 2:
+            try:
+                port = int(parts[1])
+            except ValueError:
+                pass
+        return parts[0], port
+
+    def processConnectRequest(self):
+        parsed = urlparse.urlparse(self.uri)
+        default_port = self.ports.get(parsed.scheme)
+
+        host, port = self.splitHostPort(parsed.netloc or parsed.path,
+                                        default_port)
+        if port is None:
+            self.fail("Bad CONNECT Request",
+                      "Unable to parse port from URI: %s" % self.uri)
+            return
+
+        # Copy headers from upstream reques
+        headers = self.getAllHeaders().copy()
+        if 'host' not in headers:            
+            headers['host'] = host           
+
+        clientFactory = ConnectProxyClientFactory(host, port, self)
+
+        # TODO provide an API to set proxy connect timeouts
+        self.reactor.connectTCP(host, port, clientFactory)
+
+
+class ConnectProxy(Proxy):
+    """HTTP Server Protocol that supports CONNECT"""
+    requestFactory = ConnectProxyRequest
+    connectedRemote = None
+
+    def requestDone(self, request):
+        if request.method == 'CONNECT' and self.connectedRemote is not None:
+            self.connectedRemote.connectedClient = self
+        else:
+            Proxy.requestDone(self, request)
+
+    def connectionLost(self, reason):
+        if self.connectedRemote is not None:
+            self.connectedRemote.transport.loseConnection(reason)
+        Proxy.connectionLost(self, reason)
+
+    def dataReceived(self, data):
+        if self.connectedRemote is None:
+            Proxy.dataReceived(self, data)
+        else:
+            # Once proxy is connected, forward all bytes received
+            # from the original client to the remote server.
+            self.connectedRemote.transport.write(data)
+
+
+class ConnectProxyClient(Protocol):
+    connectedClient = None
+
+    def connectionMade(self):
+        self.factory.request.channel.connectedRemote = self
+        self.factory.request.setResponseCode(200, "CONNECT OK")
+        self.factory.request.setHeader('X-Connected-IP',
+            self.transport.realAddress[0])
+        self.factory.request.finish()
+
+    def connectionLost(self, reason):
+        if self.connectedClient is not None:
+            self.connectedClient.transport.loseConnection(reason)
+
+    def dataReceived(self, data):
+        if self.connectedClient is not None:
+            # Forward all bytes from the remote server back to the
+            # original connected client
+            self.connectedClient.transport.write(data)
+        else:
+            log.msg("UNEXPECTED DATA RECEIVED:", data)
+
+
+class ConnectProxyClientFactory(ClientFactory):
+    protocol = ConnectProxyClient
+
+    def __init__(self, host, port, request):
+        self.request = request
+        self.host = host
+        self.port = port
+
+    def clientConnectionFailed(self, connector, reason):
+        self.request.fail("Gateway Error", str(reason))
+
+
+
+if __name__ == '__main__':
+    import sys
+    log.startLogging(sys.stderr)
+
+    import argparse
+    ap = argparse.ArgumentParser()
+    ap.add_argument('port', default=0, nargs='?', type=int)
+    ns = ap.parse_args()
+
+    import twisted.web.http
+    factory = twisted.web.http.HTTPFactory()
+    factory.protocol = ConnectProxy
+
+    import twisted.internet
+    c = twisted.internet.reactor.listenTCP(ns.port, factory)
+    reactor.run()
No results found