Improve robustness of proxy CONNECT, test coverage to 100%.

2 files changed, 30 insertions, 2 deletions

diff --git a/libpathod/pathoc.py b/libpathod/pathoc.py
index 650aa42a..1540d817 100644
--- a/libpathod/pathoc.py
+++ b/libpathod/pathoc.py
@@ -11,6 +11,9 @@ class Response:
         self.httpversion, self.status_code, self.msg = httpversion, status_code, msg
         self.headers, self.content = headers, content
 
+    def __repr__(self):
+        return "Response(%s - %s)"%(self.status_code, self.msg)
+
 
 class Pathoc(tcp.TCPClient):
     def __init__(self, host, port, ssl=None, sni=None, clientcert=None):
@@ -28,8 +31,13 @@ class Pathoc(tcp.TCPClient):
                     '\r\n'
                     )
         wfile.flush()
-        rfile.readline()
-        headers = http.read_headers(self.rfile)
+        l = rfile.readline()
+        if not l:
+            raise PathocError("Proxy CONNECT failed")
+        parsed = http.parse_response_line(l)
+        if not parsed[1] == 200:
+            raise PathocError("Proxy CONNECT failed: %s - %s"%(parsed[1], parsed[2]))
+        headers = http.read_headers(rfile)
 
     def connect(self, connect_to=None):
         """
diff --git a/test/test_pathoc.py b/test/test_pathoc.py
index 52a1b5ee..5391167f 100644
--- a/test/test_pathoc.py
+++ b/test/test_pathoc.py
@@ -2,6 +2,10 @@ import json, cStringIO
 from libpathod import pathoc, test, version
 import tutils
 
+def test_response():
+    r = pathoc.Response("1.1", 200, "Message", {}, None)
+    assert repr(r)
+
 
 class _TestDaemon:
     @classmethod
@@ -126,3 +130,19 @@ class TestDaemon(_TestDaemon):
         assert "foo" in self.tval(["+%s"%d], showreq=True)
         assert "File" in self.tval(["+/nonexistent"])
 
+    def test_connect_fail(self):
+        to = ("foobar", 80)
+        c = pathoc.Pathoc("127.0.0.1", self.d.port)
+        r, w = cStringIO.StringIO(), cStringIO.StringIO()
+        tutils.raises("connect failed", c.http_connect, to, w, r)
+        r = cStringIO.StringIO(
+            "HTTP/1.1 500 OK\r\n"
+        )
+        tutils.raises("connect failed", c.http_connect, to, w, r)
+        r = cStringIO.StringIO(
+            "HTTP/1.1 200 OK\r\n"
+        )
+        c.http_connect(to, w, r)
+
+
+