Bug 280661 - SOCKS proxy server connection timeout hard-coded; r=bzbarsky
authorChristopher Davis <chrisd@torproject.org>
Thu, 24 Feb 2011 15:10:08 +0200
changeset 63796 739bfb41cd012f739ebd361418af7125e593e784
parent 63795 7c42f37e0284f3e63c42d3734000aa73c09e6a6a
child 63797 2def74c0f2772aefab84802dbfc15fa71fea6d7c
push idunknown
push userunknown
push dateunknown
reviewersbzbarsky
bugs280661
milestone2.2a1pre
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 280661 - SOCKS proxy server connection timeout hard-coded; r=bzbarsky
netwerk/base/src/nsSocketTransport2.cpp
netwerk/socket/nsSOCKSIOLayer.cpp
--- a/netwerk/base/src/nsSocketTransport2.cpp
+++ b/netwerk/base/src/nsSocketTransport2.cpp
@@ -1222,16 +1222,26 @@ nsSocketTransport::InitiateSocket()
                 // connection... wouldn't we need to call ProxyStartSSL after a call
                 // to PR_ConnectContinue indicates that we are connected?
                 //
                 // XXX this appears to be what the old socket transport did.  why
                 // isn't this broken?
             }
         }
         //
+        // A SOCKS request was rejected; get the actual error code from
+        // the OS error
+        //
+        else if (PR_UNKNOWN_ERROR == code &&
+                 mProxyTransparent &&
+                 !mProxyHost.IsEmpty()) {
+            code = PR_GetOSError();
+            rv = ErrorAccordingToNSPR(code);
+        }
+        //
         // The connection was refused...
         //
         else {
             rv = ErrorAccordingToNSPR(code);
             if ((rv == NS_ERROR_CONNECTION_REFUSED) && !mProxyHost.IsEmpty())
                 rv = NS_ERROR_PROXY_CONNECTION_REFUSED;
         }
     }
@@ -1544,17 +1554,26 @@ nsSocketTransport::OnSocketReady(PRFileD
             //
             // If the connect is still not ready, then continue polling...
             //
             if ((PR_WOULD_BLOCK_ERROR == code) || (PR_IN_PROGRESS_ERROR == code)) {
                 // Set up the select flags for connect...
                 mPollFlags = (PR_POLL_EXCEPT | PR_POLL_WRITE);
                 // Update poll timeout in case it was changed
                 mPollTimeout = mTimeouts[TIMEOUT_CONNECT];
-            } 
+            }
+            //
+            // The SOCKS proxy rejected our request. Find out why.
+            //
+            else if (PR_UNKNOWN_ERROR == code &&
+                     mProxyTransparent &&
+                     !mProxyHost.IsEmpty()) {
+                code = PR_GetOSError();
+                mCondition = ErrorAccordingToNSPR(code);
+            }
             else {
                 //
                 // else, the connection failed...
                 //
                 mCondition = ErrorAccordingToNSPR(code);
                 if ((mCondition == NS_ERROR_CONNECTION_REFUSED) && !mProxyHost.IsEmpty())
                     mCondition = NS_ERROR_PROXY_CONNECTION_REFUSED;
                 SOCKET_LOG(("  connection failed! [reason=%x]\n", mCondition));
--- a/netwerk/socket/nsSOCKSIOLayer.cpp
+++ b/netwerk/socket/nsSOCKSIOLayer.cpp
@@ -20,16 +20,17 @@
  * Portions created by the Initial Developer are Copyright (C) 1998
  * the Initial Developer. All Rights Reserved.
  *
  * Contributor(s):
  *   Justin Bradford <jab@atdot.org>
  *   Bradley Baetz <bbaetz@acm.org>
  *   Darin Fisher <darin@meer.net>
  *   Malcolm Smith <malsmith@cs.rmit.edu.au>
+ *   Christopher Davis <chrisd@torproject.org>
  *
  * Alternatively, the contents of this file may be used under the terms of
  * either the GNU General Public License Version 2 or later (the "GPL"), or
  * the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
  * in which case the provisions of the GPL or the LGPL are applicable instead
  * of those above. If you wish to allow use of your version of this file only
  * under the terms of either the GPL or the LGPL, and not to allow others to
  * use your version of this file under the terms of the MPL, indicate your
@@ -63,51 +64,115 @@ static PRLogModuleInfo *gSOCKSLog;
 
 #else
 #define LOGDEBUG(args)
 #define LOGERROR(args)
 #endif
 
 class nsSOCKSSocketInfo : public nsISOCKSSocketInfo
 {
+    enum State {
+        SOCKS_INITIAL,
+        SOCKS_CONNECTING_TO_PROXY,
+        SOCKS4_WRITE_CONNECT_REQUEST,
+        SOCKS4_READ_CONNECT_RESPONSE,
+        SOCKS5_WRITE_AUTH_REQUEST,
+        SOCKS5_READ_AUTH_RESPONSE,
+        SOCKS5_WRITE_CONNECT_REQUEST,
+        SOCKS5_READ_CONNECT_RESPONSE_TOP,
+        SOCKS5_READ_CONNECT_RESPONSE_BOTTOM,
+        SOCKS_CONNECTED,
+        SOCKS_FAILED
+    };
+
+    // A buffer of 262 bytes should be enough for any request and response
+    // in case of SOCKS4 as well as SOCKS5
+    static const PRUint32 BUFFER_SIZE = 262;
+    static const PRUint32 MAX_HOSTNAME_LEN = 255;
+
 public:
     nsSOCKSSocketInfo();
-    virtual ~nsSOCKSSocketInfo() {}
+    virtual ~nsSOCKSSocketInfo() { HandshakeFinished(); }
 
     NS_DECL_ISUPPORTS
     NS_DECL_NSISOCKSSOCKETINFO
 
     void Init(PRInt32 version,
               const char *proxyHost,
               PRInt32 proxyPort,
               const char *destinationHost,
               PRUint32 flags);
 
-    const nsCString &DestinationHost() { return mDestinationHost; }
-    const nsCString &ProxyHost()       { return mProxyHost; }
-    PRInt32          ProxyPort()       { return mProxyPort; }
-    PRInt32          Version()         { return mVersion; }
-    PRUint32         Flags()           { return mFlags; }
+    void SetConnectTimeout(PRIntervalTime to);
+    PRStatus DoHandshake(PRFileDesc *fd, PRInt16 oflags = -1);
+    PRInt16 GetPollFlags() const;
+    bool IsConnected() const { return mState == SOCKS_CONNECTED; }
 
 private:
+    void HandshakeFinished(PRErrorCode err = 0);
+    PRStatus ConnectToProxy(PRFileDesc *fd);
+    PRStatus ContinueConnectingToProxy(PRFileDesc *fd, PRInt16 oflags);
+    PRStatus WriteV4ConnectRequest();
+    PRStatus ReadV4ConnectResponse();
+    PRStatus WriteV5AuthRequest();
+    PRStatus ReadV5AuthResponse();
+    PRStatus WriteV5ConnectRequest();
+    PRStatus ReadV5AddrTypeAndLength(PRUint8 *type, PRUint32 *len);
+    PRStatus ReadV5ConnectResponseTop();
+    PRStatus ReadV5ConnectResponseBottom();
+
+    void WriteUint8(PRUint8 d);
+    void WriteUint16(PRUint16 d);
+    void WriteUint32(PRUint32 d);
+    void WriteNetAddr(const PRNetAddr *addr);
+    void WriteNetPort(const PRNetAddr *addr);
+    void WriteString(const nsACString &str);
+
+    PRUint8 ReadUint8();
+    PRUint16 ReadUint16();
+    PRUint32 ReadUint32();
+    void ReadNetAddr(PRNetAddr *addr, PRUint16 fam);
+    void ReadNetPort(PRNetAddr *addr);
+
+    void WantRead(PRUint32 sz);
+    PRStatus ReadFromSocket(PRFileDesc *fd);
+    PRStatus WriteToSocket(PRFileDesc *fd);
+
+private:
+    State     mState;
+    PRUint8 * mData;
+    PRUint8 * mDataIoPtr;
+    PRUint32  mDataLength;
+    PRUint32  mReadOffset;
+    PRUint32  mAmountToRead;
+    nsCOMPtr<nsIDNSRecord> mDnsRec;
+
     nsCString mDestinationHost;
     nsCString mProxyHost;
     PRInt32   mProxyPort;
     PRInt32   mVersion;   // SOCKS version 4 or 5
     PRUint32  mFlags;
     PRNetAddr mInternalProxyAddr;
     PRNetAddr mExternalProxyAddr;
     PRNetAddr mDestinationAddr;
+    PRIntervalTime mTimeout;
 };
 
 nsSOCKSSocketInfo::nsSOCKSSocketInfo()
-    : mProxyPort(-1)
+    : mState(SOCKS_INITIAL)
+    , mDataIoPtr(nsnull)
+    , mDataLength(0)
+    , mReadOffset(0)
+    , mAmountToRead(0)
+    , mProxyPort(-1)
     , mVersion(-1)
     , mFlags(0)
+    , mTimeout(PR_INTERVAL_NO_TIMEOUT)
 {
+    mData = new PRUint8[BUFFER_SIZE];
     PR_InitializeNetAddr(PR_IpAddrAny, 0, &mInternalProxyAddr);
     PR_InitializeNetAddr(PR_IpAddrAny, 0, &mExternalProxyAddr);
     PR_InitializeNetAddr(PR_IpAddrAny, 0, &mDestinationAddr);
 }
 
 void
 nsSOCKSSocketInfo::Init(PRInt32 version, const char *proxyHost, PRInt32 proxyPort, const char *host, PRUint32 flags)
 {
@@ -157,647 +222,817 @@ nsSOCKSSocketInfo::GetInternalProxyAddr(
 
 NS_IMETHODIMP 
 nsSOCKSSocketInfo::SetInternalProxyAddr(PRNetAddr *aInternalProxyAddr)
 {
     memcpy(&mInternalProxyAddr, aInternalProxyAddr, sizeof(PRNetAddr));
     return NS_OK;
 }
 
-static PRInt32
-pr_RecvAll(PRFileDesc *fd, unsigned char *buf, PRInt32 amount, PRIntn flags, 
-           PRIntervalTime *timeout)
-{
-    PRInt32 bytesRead = 0;
-    PRInt32 offset = 0;
-
-    while (offset < amount) {
-        PRIntervalTime start_time = PR_IntervalNow();
-        bytesRead = PR_Recv(fd, buf + offset, amount - offset, flags, *timeout);
-        PRIntervalTime elapsed = PR_IntervalNow() - start_time;
-
-        if (elapsed > *timeout) {
-            *timeout = 0;
-        } else {
-            *timeout -= elapsed;
-        }
-
-        if (bytesRead > 0) {
-            offset += bytesRead;
-        } else if (bytesRead == 0 || offset != 0) {
-            return offset;
-        } else {
-            return bytesRead;
-        }
-
-        if (*timeout == 0) {
-            LOGERROR(("PR_Recv() timed out. amount = %d. offset = %d.",
-                     amount, offset));
-            return offset;
-        }
-    }
-    return offset;
-}
-
-static PRInt32
-pr_Send(PRFileDesc *fd, const void *buf, PRInt32 amount, PRIntn flags,
-        PRIntervalTime *timeout)
-{
-    PRIntervalTime start_time = PR_IntervalNow();
-    PRInt32 retval = PR_Send(fd, buf, amount, flags, *timeout);
-    PRIntervalTime elapsed = PR_IntervalNow() - start_time;
-
-    if (elapsed > *timeout) {
-        *timeout = 0;
-        LOGERROR(("PR_Send() timed out. amount = %d. retval = %d.",
-                 amount, retval));
-        return retval;
-    } else {
-        *timeout -= elapsed;
-    }
-
-    if (retval <= 0) {
-        LOGERROR(("PR_Send() failed. amount = %d. retval = %d.",
-                 amount, retval));
-    }
-
-    return retval;
-}
-
-// Negotiate a SOCKS 5 connection. Assumes the TCP connection to the socks 
-// server port has been established.
-static nsresult
-ConnectSOCKS5(PRFileDesc *fd, const PRNetAddr *addr, PRNetAddr *extAddr, PRIntervalTime timeout)
+// There needs to be a means of distinguishing between connection errors
+// that the SOCKS server reports when it rejects a connection request, and
+// connection errors that happen while attempting to connect to the SOCKS
+// server. Otherwise, Firefox will report incorrectly that the proxy server
+// is refusing connections when a SOCKS request is rejected by the proxy.
+// When a SOCKS handshake failure occurs, the PR error is set to
+// PR_UNKNOWN_ERROR, and the real error code is returned via the OS error.
+void
+nsSOCKSSocketInfo::HandshakeFinished(PRErrorCode err)
 {
-    int request_len = 0;
-    int response_len = 0;
-    int desired_len = 0;
-    unsigned char request[22];
-    unsigned char response[262];
-
-    NS_ENSURE_TRUE(fd, NS_ERROR_NOT_INITIALIZED);
-    NS_ENSURE_TRUE(addr, NS_ERROR_NOT_INITIALIZED);
-    NS_ENSURE_TRUE(extAddr, NS_ERROR_NOT_INITIALIZED);
-
-    request[0] = 0x05; // SOCKS version 5
-    request[1] = 0x01; // number of auth procotols we recognize
-    // auth protocols
-    request[2] = 0x00; // no authentication required
-    // compliant implementations MUST implement GSSAPI
-    // and SHOULD implement username/password and MAY
-    // implement CHAP
-    // TODO: we don't implement these
-    //request[3] = 0x01; // GSSAPI
-    //request[4] = 0x02; // username/password
-    //request[5] = 0x03; // CHAP
-
-    request_len = 2 + request[1];
-    int write_len = pr_Send(fd, request, request_len, 0, &timeout);
-    if (write_len != request_len) {
-        return NS_ERROR_FAILURE;
-    }
-
-    // get the server's response. 
-    desired_len = 2;
-    response_len = pr_RecvAll(fd, response, desired_len, 0, &timeout);
-
-    if (response_len < desired_len) {
-        LOGERROR(("pr_RecvAll() failed. response_len = %d.", response_len));
-        return NS_ERROR_FAILURE;
-    }
-
-    if (response[0] != 0x05) {
-        // it's a either not SOCKS or not our version
-        LOGERROR(("Not a SOCKS 5 reply. Expected: 5; received: %x", response[0]));
-        return NS_ERROR_FAILURE;
-    }
-    switch (response[1]) {
-        case 0x00:
-            // no auth
-            break;
-        case 0x01:
-            // GSSAPI
-            // TODO: implement
-            LOGERROR(("Server want to use GSSAPI to authenticate, but we don't support it."));
-            return NS_ERROR_FAILURE;
-        case 0x02:
-            // username/password
-            // TODO: implement
-            LOGERROR(("Server want to use username/password to authenticate, but we don't support it."));
-            return NS_ERROR_FAILURE;
-        case 0x03:
-            // CHAP
-            // TODO: implement?
-            LOGERROR(("Server want to use CHAP to authenticate, but we don't support it."));
-            return NS_ERROR_FAILURE;
-        default:
-            // unrecognized auth method
-            LOGERROR(("Uncrecognized authentication method received: %x", response[1]));
-            return NS_ERROR_FAILURE;
+    if (err == 0) {
+        mState = SOCKS_CONNECTED;
+    } else {
+        mState = SOCKS_FAILED;
+        PR_SetError(PR_UNKNOWN_ERROR, err);
     }
 
-    // we are now authenticated, so lets tell
-    // the server where to connect to
-
-    request_len = 0;
-
-    request[0] = 0x05; // SOCKS version 5
-    request[1] = 0x01; // CONNECT command
-    request[2] = 0x00; // obligatory reserved field (perfect for MS tampering!)
-
-    // get destination port
-    PRInt32 destPort = PR_ntohs(PR_NetAddrInetPort(addr));
-    nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
-
-    if (info->Flags() & nsISocketProvider::PROXY_RESOLVES_HOST) {
-
-        LOGDEBUG(("using server to resolve hostnames rather than resolving it first\n"));
-
-        // if the PROXY_RESOLVES_HOST flag is set, we assume
-        // that the transport wants us to pass the SOCKS server the 
-        // hostname and port and let it do the name resolution.
-
-        // the real destination hostname and port was stored
-        // in our info object earlier when this layer was created.
-
-        const nsCString& destHost = info->DestinationHost();
-
-        LOGDEBUG(("host:port -> %s:%li", destHost.get(), destPort));
-
-        request[3] = 0x03; // encoding of destination address (3 == hostname)
-
-        int host_len = destHost.Length();
-        if (host_len > 255) {
-            // SOCKS5 transmits the length of the hostname in a single char.
-            // This gives us an absolute limit of 255 chars in a hostname, and
-            // there's nothing we can do to extend it.  I don't think many
-            // hostnames will ever be bigger than this, so hopefully it's an
-            // uneventful abort condition.
-            LOGERROR (("Hostname too big for SOCKS5."));
-            return NS_ERROR_INVALID_ARG;
-        }
-        request[4] = (char) host_len;
-        request_len = 5;
-
-        // Send the initial header first...
-        write_len = pr_Send(fd, request, request_len, 0, &timeout);
-        if (write_len != request_len) {
-            // bad write
-            return NS_ERROR_FAILURE;
-        }
-
-        // Now send the hostname...
-        write_len = pr_Send(fd, destHost.get(), host_len, 0, &timeout);
-        if (write_len != host_len) {
-            // bad write
-            return NS_ERROR_FAILURE;
-        }
-
-        // There's no data left because we just sent it.
-        request_len = 0;
-
-    } else if (PR_NetAddrFamily(addr) == PR_AF_INET) {
-
-        request[3] = 0x01; // encoding of destination address (1 == IPv4)
-        request_len = 8;   // 4 for address, 4 SOCKS headers
-
-        char * ip = (char*)(&addr->inet.ip);
-        request[4] = *ip++;
-        request[5] = *ip++;
-        request[6] = *ip++;
-        request[7] = *ip++;
-
-    } else if (PR_NetAddrFamily(addr) == PR_AF_INET6) {
-
-        request[3] = 0x04; // encoding of destination address (4 == IPv6)
-        request_len = 20;  // 16 for address, 4 SOCKS headers
-
-        char * ip = (char*)(&addr->ipv6.ip.pr_s6_addr);
-        request[4] = *ip++; request[5] = *ip++; 
-        request[6] = *ip++; request[7] = *ip++;
-        request[8] = *ip++; request[9] = *ip++; 
-        request[10] = *ip++; request[11] = *ip++;
-        request[12] = *ip++; request[13] = *ip++; 
-        request[14] = *ip++; request[15] = *ip++;
-        request[16] = *ip++; request[17] = *ip++; 
-        request[18] = *ip++; request[19] = *ip++;
-
-        // we're going to test to see if this address can
-        // be mapped back into IPv4 without loss. if so,
-        // we'll use IPv4 instead, as reliable SOCKS server 
-        // support for IPv6 is probably questionable.
-
-        if (PR_IsNetAddrType(addr, PR_IpAddrV4Mapped)) {
-            request[3] = 0x01; // ipv4 encoding
-            request[4] = request[16];
-            request[5] = request[17];
-            request[6] = request[18];
-            request[7] = request[19];
-            request_len -= 12;
-        }
-    } else {
-        // Unknown address type
-        LOGERROR(("Don't know what kind of IP address this is."));
-        return NS_ERROR_FAILURE;
-    }
-
-    // add the destination port to the request
-    request[request_len] = (unsigned char)(destPort >> 8);
-    request[request_len+1] = (unsigned char)destPort;
-    request_len += 2;
-
-    write_len = pr_Send(fd, request, request_len, 0, &timeout);
-    if (write_len != request_len) {
-        // bad write
-        return NS_ERROR_FAILURE;
-    }
-
-    desired_len = 5;
-    response_len = pr_RecvAll(fd, response, desired_len, 0, &timeout);
-    if (response_len < desired_len) { // bad read
-        LOGERROR(("pr_RecvAll() failed getting connect command reply. response_len = %d.", response_len));
-        return NS_ERROR_FAILURE;
-    }
-
-    if (response[0] != 0x05) {
-        // bad response
-        LOGERROR(("Not a SOCKS 5 reply. Expected: 5; received: %x", response[0]));
-        return NS_ERROR_FAILURE;
-    }
-
-    switch(response[1]) {
-        case 0x00:  break;      // success
-        case 0x01:  LOGERROR(("SOCKS 5 server rejected connect request: 01, General SOCKS server failure."));
-                    return NS_ERROR_FAILURE;
-        case 0x02:  LOGERROR(("SOCKS 5 server rejected connect request: 02, Connection not allowed by ruleset."));
-                    return NS_ERROR_FAILURE;
-        case 0x03:  LOGERROR(("SOCKS 5 server rejected connect request: 03, Network unreachable."));
-                    return NS_ERROR_FAILURE;
-        case 0x04:  LOGERROR(("SOCKS 5 server rejected connect request: 04, Host unreachable."));
-                    return NS_ERROR_FAILURE;
-        case 0x05:  LOGERROR(("SOCKS 5 server rejected connect request: 05, Connection refused."));
-                    return NS_ERROR_FAILURE;
-        case 0x06:  LOGERROR(("SOCKS 5 server rejected connect request: 06, TTL expired."));
-                    return NS_ERROR_FAILURE;
-        case 0x07:  LOGERROR(("SOCKS 5 server rejected connect request: 07, Command not supported."));
-                    return NS_ERROR_FAILURE;
-        case 0x08:  LOGERROR(("SOCKS 5 server rejected connect request: 08, Address type not supported."));
-                    return NS_ERROR_FAILURE;
-        default:    LOGERROR(("SOCKS 5 server rejected connect request: %x.", response[1]));
-                    return NS_ERROR_FAILURE;
-
-
-    }
-
-    switch (response[3]) {
-        case 0x01: // IPv4
-	    desired_len = 4 + 2 - 1;
-            break;
-        case 0x03: // FQDN 
-	    desired_len = response[4] + 2;
-            break;
-        case 0x04: // IPv6
-	    desired_len = 16 + 2 - 1;
-            break;
-        default: // unknown format
-            return NS_ERROR_FAILURE;
-            break;
-    }
-    response_len = pr_RecvAll(fd, response + 5, desired_len, 0, &timeout);
-    if (response_len < desired_len) { // bad read
-        LOGERROR(("pr_RecvAll() failed getting connect command reply. response_len = %d.", response_len));
-        return NS_ERROR_FAILURE;
-    }
-    response_len += 5;
-
-    // get external bound address (this is what 
-    // the outside world sees as "us")
-    char *ip = nsnull;
-    PRUint16 extPort = 0;
-
-    switch (response[3]) {
-        case 0x01: // IPv4
-
-            extPort = (response[8] << 8) | response[9];
-
-            PR_SetNetAddr(PR_IpAddrAny, PR_AF_INET, extPort, extAddr);
-
-            ip = (char*)(&extAddr->inet.ip);
-            *ip++ = response[4];
-            *ip++ = response[5];
-            *ip++ = response[6];
-            *ip++ = response[7];
-
-            break;
-        case 0x04: // IPv6
-
-            extPort = (response[20] << 8) | response[21];
-
-            PR_SetNetAddr(PR_IpAddrAny, PR_AF_INET6, extPort, extAddr);
-
-            ip = (char*)(&extAddr->ipv6.ip.pr_s6_addr);
-            *ip++ = response[4]; *ip++ = response[5]; 
-            *ip++ = response[6]; *ip++ = response[7];
-            *ip++ = response[8]; *ip++ = response[9]; 
-            *ip++ = response[10]; *ip++ = response[11];
-            *ip++ = response[12]; *ip++ = response[13]; 
-            *ip++ = response[14]; *ip++ = response[15];
-            *ip++ = response[16]; *ip++ = response[17]; 
-            *ip++ = response[18]; *ip++ = response[19];
-
-            break;
-        case 0x03: // FQDN 
-            // if we get here, we don't know our external address.
-            // however, as that's possibly not critical to the user,
-            // we let it slide.
-            extPort = (response[response_len - 2] << 8) | 
-                       response[response_len - 1];
-            PR_InitializeNetAddr(PR_IpAddrNull, extPort, extAddr);
-            break;
-    }
-    return NS_OK;
+    // We don't need the buffer any longer, so free it.
+    delete [] mData;
+    mData = nsnull;
+    mDataIoPtr = nsnull;
+    mDataLength = 0;
+    mReadOffset = 0;
+    mAmountToRead = 0;
 }
 
-// Negotiate a SOCKS 4 connection. Assumes the TCP connection to the socks 
-// server port has been established.
-static nsresult
-ConnectSOCKS4(PRFileDesc *fd, const PRNetAddr *addr, PRIntervalTime timeout)
+PRStatus
+nsSOCKSSocketInfo::ConnectToProxy(PRFileDesc *fd)
 {
-    int request_len = 0;
-    int write_len;
-    int response_len = 0;
-    int desired_len = 0;
-    char *ip = nsnull;
-    unsigned char request[12];
-    unsigned char response[10];
-
-    NS_ENSURE_TRUE(fd, NS_ERROR_NOT_INITIALIZED);
-    NS_ENSURE_TRUE(addr, NS_ERROR_NOT_INITIALIZED);
-
-    request[0] = 0x04; // SOCKS version 4
-    request[1] = 0x01; // CD command code -- 1 for connect
-
-    // destination port
-    PRInt32 destPort = PR_ntohs(PR_NetAddrInetPort(addr));
-
-    // store the port
-    request[2] = (unsigned char)(destPort >> 8);
-    request[3] = (unsigned char)destPort;
+    PRStatus status;
+    nsresult rv;
 
-    // username
-    request[8] = 'M';
-    request[9] = 'O';
-    request[10] = 'Z';
-
-    request[11] = 0x00;
-
-    request_len = 12;
-
-    nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
-
-    if (info->Flags() & nsISocketProvider::PROXY_RESOLVES_HOST) {
-
-        LOGDEBUG(("using server to resolve hostnames rather than resolving it first\n"));
-
-        // if the PROXY_RESOLVES_HOST flag is set, we assume that the
-        // transport wants us to pass the SOCKS server the hostname
-        // and port and let it do the name resolution.
-
-        // an extension to SOCKS 4, called 4a, specifies a way
-        // to do this, so we'll try that and hope the
-        // server supports it.
-
-        // the real destination hostname and port was stored
-        // in our info object earlier when this layer was created.
-
-        const nsCString& destHost = info->DestinationHost();
-
-        LOGDEBUG(("host:port -> %s:%li\n", destHost.get(), destPort));
+    NS_ABORT_IF_FALSE(mState == SOCKS_INITIAL,
+                      "Must be in initial state to make connection!");
 
-        // the IP portion of the query is set to this special address.
-        request[4] = 0;
-        request[5] = 0;
-        request[6] = 0;
-        request[7] = 1;
-
-        write_len = pr_Send(fd, request, request_len, 0, &timeout);
-        if (write_len != request_len) {
-            return NS_ERROR_FAILURE;
-        }
-
-        // Remember the NULL.
-        int host_len = destHost.Length() + 1;
-
-        write_len = pr_Send(fd, destHost.get(), host_len, 0, &timeout);
-        if (write_len != host_len) {
-            return NS_ERROR_FAILURE;
-        }
-
-        // No data to send, just sent it.
-        request_len = 0;
-
-    } else if (PR_NetAddrFamily(addr) == PR_AF_INET) { // IPv4
+    // If we haven't performed the DNS lookup, do that now.
+    if (!mDnsRec) {
+        nsCOMPtr<nsIDNSService> dns = do_GetService(NS_DNSSERVICE_CONTRACTID);
+        if (!dns)
+            return PR_FAILURE;
 
-        // store the ip
-        ip = (char*)(&addr->inet.ip);
-        request[4] = *ip++;
-        request[5] = *ip++;
-        request[6] = *ip++;
-        request[7] = *ip++;
-
-    } else if (PR_NetAddrFamily(addr) == PR_AF_INET6) { // IPv6
-
-        // IPv4 address encoded in an IPv6 address
-        if (PR_IsNetAddrType(addr, PR_IpAddrV4Mapped)) {
-            // store the ip
-            ip = (char*)(&addr->ipv6.ip.pr_s6_addr[12]);
-            request[4] = *ip++;
-            request[5] = *ip++;
-            request[6] = *ip++;
-            request[7] = *ip++;
-        } else {
-            LOGERROR(("IPv6 is not supported in SOCKS 4."));
-            return NS_ERROR_FAILURE;	// SOCKS 4 can't do IPv6
-        }
-
-    } else {
-        LOGERROR(("Don't know what kind of IP address this is."));
-        return NS_ERROR_FAILURE;		// don't recognize this type
-    }
-
-    if (request_len > 0) {
-        write_len = pr_Send(fd, request, request_len, 0, &timeout);
-        if (write_len != request_len) {
-            return NS_ERROR_FAILURE;
+        rv = dns->Resolve(mProxyHost, 0, getter_AddRefs(mDnsRec));
+        if (NS_FAILED(rv)) {
+            LOGERROR(("socks: DNS lookup for SOCKS proxy %s failed",
+                     mProxyHost.get()));
+            return PR_FAILURE;
         }
     }
 
-    // get the server's response
-    desired_len = 8;	// size of the response
-    response_len = pr_RecvAll(fd, response, desired_len, 0, &timeout);
-    if (response_len < desired_len) {
-        LOGERROR(("pr_RecvAll() failed. response_len = %d.", response_len));
-        return NS_ERROR_FAILURE;
+    do {
+        rv = mDnsRec->GetNextAddr(mProxyPort, &mInternalProxyAddr);
+        // No more addresses to try? If so, we'll need to bail
+        if (NS_FAILED(rv)) {
+            LOGERROR(("socks: unable to connect to SOCKS proxy, %s",
+                     mProxyHost.get()));
+            return PR_FAILURE;
+        }
+
+#if defined(PR_LOGGING)
+        char buf[64];
+        PR_NetAddrToString(&mInternalProxyAddr, buf, sizeof(buf));
+        LOGDEBUG(("socks: trying proxy server, %s:%hu",
+                 buf, PR_ntohs(PR_NetAddrInetPort(&mInternalProxyAddr))));
+#endif
+        status = fd->lower->methods->connect(fd->lower,
+                        &mInternalProxyAddr, mTimeout);
+        if (status != PR_SUCCESS) {
+            PRErrorCode c = PR_GetError();
+            // If EINPROGRESS, return now and check back later after polling
+            if (c == PR_WOULD_BLOCK_ERROR || c == PR_IN_PROGRESS_ERROR) {
+                mState = SOCKS_CONNECTING_TO_PROXY;
+                return status;
+            }
+        }
+    } while (status != PR_SUCCESS);
+
+    // Connected now, start SOCKS
+    if (mVersion == 4)
+        return WriteV4ConnectRequest();
+    return WriteV5AuthRequest();
+}
+
+PRStatus
+nsSOCKSSocketInfo::ContinueConnectingToProxy(PRFileDesc *fd, PRInt16 oflags)
+{
+    PRStatus status;
+
+    NS_ABORT_IF_FALSE(mState == SOCKS_CONNECTING_TO_PROXY,
+                      "Continuing connection in wrong state!");
+
+    LOGDEBUG(("socks: continuing connection to proxy"));
+
+    status = fd->lower->methods->connectcontinue(fd->lower, oflags);
+    if (status != PR_SUCCESS) {
+        PRErrorCode c = PR_GetError();
+        if (c != PR_WOULD_BLOCK_ERROR && c != PR_IN_PROGRESS_ERROR) {
+            // A connection failure occured, try another address
+            mState = SOCKS_INITIAL;
+            return ConnectToProxy(fd);
+        }
+
+        // We're still connecting
+        return PR_FAILURE;
+    }
+
+    // Connected now, start SOCKS
+    if (mVersion == 4)
+        return WriteV4ConnectRequest();
+    return WriteV5AuthRequest();
+}
+
+PRStatus
+nsSOCKSSocketInfo::WriteV4ConnectRequest()
+{
+    PRNetAddr *addr = &mDestinationAddr;
+    PRInt32 proxy_resolve;
+
+    NS_ABORT_IF_FALSE(mState == SOCKS_CONNECTING_TO_PROXY,
+                      "Invalid state!");
+    
+    proxy_resolve = mFlags & nsISocketProvider::PROXY_RESOLVES_HOST;
+
+    mDataLength = 0;
+    mState = SOCKS4_WRITE_CONNECT_REQUEST;
+
+    LOGDEBUG(("socks4: sending connection request (socks4a resolve? %s)",
+             proxy_resolve? "yes" : "no"));
+
+    // Send a SOCKS 4 connect request.
+    WriteUint8(0x04); // version -- 4
+    WriteUint8(0x01); // command -- connect
+    WriteNetPort(addr);
+    if (proxy_resolve) {
+        // Add the full name, null-terminated, to the request
+        // according to SOCKS 4a. A fake IP address, with the first
+        // four bytes set to 0 and the last byte set to something other
+        // than 0, is used to notify the proxy that this is a SOCKS 4a
+        // request. This request type works for Tor and perhaps others.
+        WriteUint32(PR_htonl(0x00000001)); // Fake IP
+        WriteUint8(0x00); // Send an emtpy username
+        if (mDestinationHost.Length() > MAX_HOSTNAME_LEN) {
+            LOGERROR(("socks4: destination host name is too long!"));
+            HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+            return PR_FAILURE;
+        }
+        WriteString(mDestinationHost); // Hostname
+        WriteUint8(0x00);
+    } else if (PR_NetAddrFamily(addr) == PR_AF_INET) {
+        WriteNetAddr(addr); // Add the IPv4 address
+        WriteUint8(0x00); // Send an emtpy username
+    } else if (PR_NetAddrFamily(addr) == PR_AF_INET6) {
+        LOGERROR(("socks: SOCKS 4 can't handle IPv6 addresses!"));
+        HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+        return PR_FAILURE;
+    }
+
+    return PR_SUCCESS;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadV4ConnectResponse()
+{
+    NS_ABORT_IF_FALSE(mState == SOCKS4_READ_CONNECT_RESPONSE,
+                      "Handling SOCKS 4 connection reply in wrong state!");
+    NS_ABORT_IF_FALSE(mDataLength == 8,
+                      "SOCKS 4 connection reply must be 8 bytes!");
+
+    LOGDEBUG(("socks4: checking connection reply"));
+
+    if (ReadUint8() != 0x00) {
+        LOGERROR(("socks4: wrong connection reply"));
+        HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+        return PR_FAILURE;
+    }
+
+    // See if our connection request was granted
+    if (ReadUint8() == 90) {
+        LOGDEBUG(("socks4: connection successful!"));
+        HandshakeFinished();
+        return PR_SUCCESS;
+    }
+
+    LOGERROR(("socks4: unable to connect"));
+    HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+    return PR_FAILURE;
+}
+
+PRStatus
+nsSOCKSSocketInfo::WriteV5AuthRequest()
+{
+    NS_ABORT_IF_FALSE(mVersion == 5, "SOCKS version must be 5!");
+
+    mState = SOCKS5_WRITE_AUTH_REQUEST;
+
+    // Send an initial SOCKS 5 greeting
+    LOGDEBUG(("socks5: sending auth methods"));
+    WriteUint8(0x05); // version -- 5
+    WriteUint8(0x01); // # auth methods -- 1
+    WriteUint8(0x00); // we don't support authentication
+
+    return PR_SUCCESS;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadV5AuthResponse()
+{
+    NS_ABORT_IF_FALSE(mState == SOCKS5_READ_AUTH_RESPONSE,
+                      "Handling SOCKS 5 auth method reply in wrong state!");
+    NS_ABORT_IF_FALSE(mDataLength == 2,
+                      "SOCKS 5 auth method reply must be 2 bytes!");
+
+    LOGDEBUG(("socks5: checking auth method reply"));
+
+    // Check version number
+    if (ReadUint8() != 0x05) {
+        LOGERROR(("socks5: unexpected version in the reply"));
+        HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+        return PR_FAILURE;
     }
 
-    if ((response[0] != 0x00) && (response[0] != 0x04)) {
-        // Novell BorderManager sends a response of type 4, should be zero
-        // According to the spec. Cope with this brokenness.        
-        // it's not a SOCKS 4 reply or version 0 of the reply code
-        LOGERROR(("Not a SOCKS 4 reply. Expected: 0; received: %x.", response[0]));
-        return NS_ERROR_FAILURE;
+    // Make sure our authentication choice was accepted
+    if (ReadUint8() != 0x00) {
+        LOGERROR(("socks5: server did not accept our authentication method"));
+        HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+        return PR_FAILURE;
+    }
+
+    return WriteV5ConnectRequest();
+}
+
+PRStatus
+nsSOCKSSocketInfo::WriteV5ConnectRequest()
+{
+    // Send SOCKS 5 connect request
+    PRNetAddr *addr = &mDestinationAddr;
+    PRInt32 proxy_resolve;
+    proxy_resolve = mFlags & nsISocketProvider::PROXY_RESOLVES_HOST;
+
+    LOGDEBUG(("socks5: sending connection request (socks5 resolve? %s)",
+             proxy_resolve? "yes" : "no"));
+
+    mDataLength = 0;
+    mState = SOCKS5_WRITE_CONNECT_REQUEST;
+
+    WriteUint8(0x05); // version -- 5
+    WriteUint8(0x01); // command -- connect
+    WriteUint8(0x00); // reserved
+   
+    // Add the address to the SOCKS 5 request. SOCKS 5 supports several
+    // address types, so we pick the one that works best for us.
+    if (proxy_resolve) {
+        // Add the host name. Only a single byte is used to store the length,
+        // so we must prevent long names from being used.
+        if (mDestinationHost.Length() > MAX_HOSTNAME_LEN) {
+            LOGERROR(("socks5: destination host name is too long!"));
+            HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+            return PR_FAILURE;
+        }
+        WriteUint8(0x03); // addr type -- domainname
+        WriteUint8(mDestinationHost.Length()); // name length
+        WriteString(mDestinationHost);
+    } else if (PR_NetAddrFamily(addr) == PR_AF_INET) {
+        WriteUint8(0x01); // addr type -- IPv4
+        WriteNetAddr(addr);
+    } else if (PR_NetAddrFamily(addr) == PR_AF_INET6) {
+        WriteUint8(0x04); // addr type -- IPv6
+        WriteNetAddr(addr);
+    } else {
+        LOGERROR(("socks5: destination address of unknown type!"));
+        HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+        return PR_FAILURE;
+    }
+
+    WriteNetPort(addr); // port
+
+    return PR_SUCCESS;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadV5AddrTypeAndLength(PRUint8 *type, PRUint32 *len)
+{
+    NS_ABORT_IF_FALSE(mState == SOCKS5_READ_CONNECT_RESPONSE_TOP ||
+                      mState == SOCKS5_READ_CONNECT_RESPONSE_BOTTOM,
+                      "Invalid state!");
+    NS_ABORT_IF_FALSE(mDataLength >= 5,
+                      "SOCKS 5 connection reply must be at least 5 bytes!");
+ 
+    // Seek to the address location 
+    mReadOffset = 3;
+   
+    *type = ReadUint8();
+
+    switch (*type) {
+        case 0x01: // ipv4
+            *len = 4 - 1;
+            break;
+        case 0x04: // ipv6
+            *len = 16 - 1;
+            break;
+        case 0x03: // fqdn
+            *len = ReadUint8();
+            break;
+        default:   // wrong address type
+            LOGERROR(("socks5: wrong address type in connection reply!"));
+            return PR_FAILURE;
+    }
+
+    return PR_SUCCESS;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadV5ConnectResponseTop()
+{
+    PRUint8 res;
+    PRUint32 len;
+
+    NS_ABORT_IF_FALSE(mState == SOCKS5_READ_CONNECT_RESPONSE_TOP,
+                      "Invalid state!");
+    NS_ABORT_IF_FALSE(mDataLength == 5,
+                      "SOCKS 5 connection reply must be exactly 5 bytes!");
+
+    LOGDEBUG(("socks5: checking connection reply"));
+
+    // Check version number
+    if (ReadUint8() != 0x05) {
+        LOGERROR(("socks5: unexpected version in the reply"));
+        HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+        return PR_FAILURE;
+    }
+
+    // Check response
+    res = ReadUint8();
+    if (res != 0x00) {
+        PRErrorCode c = PR_CONNECT_REFUSED_ERROR;
+
+        switch (res) {
+            case 0x01:
+                LOGERROR(("socks5: connect failed: "
+                          "01, General SOCKS server failure."));
+                break;
+            case 0x02:
+                LOGERROR(("socks5: connect failed: "
+                          "02, Connection not allowed by ruleset."));
+                break;
+            case 0x03:
+                LOGERROR(("socks5: connect failed: 03, Network unreachable."));
+                c = PR_NETWORK_UNREACHABLE_ERROR;
+                break;
+            case 0x04:
+                LOGERROR(("socks5: connect failed: 04, Host unreachable."));
+                break;
+            case 0x05:
+                LOGERROR(("socks5: connect failed: 05, Connection refused."));
+                break;
+            case 0x06:  
+                LOGERROR(("socks5: connect failed: 06, TTL expired."));
+                c = PR_CONNECT_TIMEOUT_ERROR;
+                break;
+            case 0x07:
+                LOGERROR(("socks5: connect failed: "
+                          "07, Command not supported."));
+                break;
+            case 0x08:
+                LOGERROR(("socks5: connect failed: "
+                          "08, Address type not supported."));
+                c = PR_BAD_ADDRESS_ERROR;
+                break;
+            default:
+                LOGERROR(("socks5: connect failed."));
+                break;
+        }
+
+        HandshakeFinished(c);
+        return PR_FAILURE;
+    }
+
+    if (ReadV5AddrTypeAndLength(&res, &len) != PR_SUCCESS) {
+        HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+        return PR_FAILURE;
+    }
+
+    mState = SOCKS5_READ_CONNECT_RESPONSE_BOTTOM;
+    WantRead(len + 2);
+
+    return PR_SUCCESS;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadV5ConnectResponseBottom()
+{
+    PRUint8 type;
+    PRUint32 len;
+
+    NS_ABORT_IF_FALSE(mState == SOCKS5_READ_CONNECT_RESPONSE_BOTTOM,
+                      "Invalid state!");
+
+    if (ReadV5AddrTypeAndLength(&type, &len) != PR_SUCCESS) {
+        HandshakeFinished(PR_BAD_ADDRESS_ERROR);
+        return PR_FAILURE;
     }
 
-    if (response[1] != 0x5A) { // = 90: request granted
-        // connect request not granted
-        LOGERROR(("Connection request refused. Expected: 90; received: %d.", response[1]));
-        return NS_ERROR_FAILURE;
+    NS_ABORT_IF_FALSE(mDataLength == 7+len,
+                      "SOCKS 5 unexpected length of connection reply!");
+
+    LOGDEBUG(("socks5: loading source addr and port"));
+    // Read what the proxy says is our source address
+    switch (type) {
+        case 0x01: // ipv4
+            ReadNetAddr(&mExternalProxyAddr, PR_AF_INET);
+            break;
+        case 0x04: // ipv6
+            ReadNetAddr(&mExternalProxyAddr, PR_AF_INET6);
+            break;
+        case 0x03: // fqdn (skip)
+            mReadOffset += len;
+            mExternalProxyAddr.raw.family = PR_AF_INET;
+            break;
+    }
+
+    ReadNetPort(&mExternalProxyAddr);
+
+    LOGDEBUG(("socks5: connected!"));
+    HandshakeFinished();
+
+    return PR_SUCCESS;
+}
+
+void
+nsSOCKSSocketInfo::SetConnectTimeout(PRIntervalTime to)
+{
+    mTimeout = to;
+}
+
+PRStatus
+nsSOCKSSocketInfo::DoHandshake(PRFileDesc *fd, PRInt16 oflags)
+{
+    LOGDEBUG(("socks: DoHandshake(), state = %d", mState));
+
+    switch (mState) {
+        case SOCKS_INITIAL:
+            return ConnectToProxy(fd);
+        case SOCKS_CONNECTING_TO_PROXY:
+            return ContinueConnectingToProxy(fd, oflags);
+        case SOCKS4_WRITE_CONNECT_REQUEST:
+            if (WriteToSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            WantRead(8);
+            mState = SOCKS4_READ_CONNECT_RESPONSE;
+            return PR_SUCCESS;
+        case SOCKS4_READ_CONNECT_RESPONSE:
+            if (ReadFromSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            return ReadV4ConnectResponse();
+
+        case SOCKS5_WRITE_AUTH_REQUEST:
+            if (WriteToSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            WantRead(2);
+            mState = SOCKS5_READ_AUTH_RESPONSE;
+            return PR_SUCCESS;
+        case SOCKS5_READ_AUTH_RESPONSE:
+            if (ReadFromSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            return ReadV5AuthResponse();
+        case SOCKS5_WRITE_CONNECT_REQUEST:
+            if (WriteToSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+
+            // The SOCKS 5 response to the connection request is variable
+            // length. First, we'll read enough to tell how long the response
+            // is, and will read the rest later.
+            WantRead(5);
+            mState = SOCKS5_READ_CONNECT_RESPONSE_TOP;
+            return PR_SUCCESS;
+        case SOCKS5_READ_CONNECT_RESPONSE_TOP:
+            if (ReadFromSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            return ReadV5ConnectResponseTop();
+        case SOCKS5_READ_CONNECT_RESPONSE_BOTTOM:
+            if (ReadFromSocket(fd) != PR_SUCCESS)
+                return PR_FAILURE;
+            return ReadV5ConnectResponseBottom();
+
+        case SOCKS_CONNECTED:
+            LOGERROR(("socks: already connected"));
+            HandshakeFinished(PR_IS_CONNECTED_ERROR);
+            return PR_FAILURE;
+        case SOCKS_FAILED:
+            LOGERROR(("socks: already failed"));
+            return PR_FAILURE;
     }
+
+    LOGERROR(("socks: executing handshake in invalid state, %d", mState));
+    HandshakeFinished(PR_INVALID_STATE_ERROR);
+
+    return PR_FAILURE;
+}
+
+PRInt16
+nsSOCKSSocketInfo::GetPollFlags() const
+{
+    switch (mState) {
+        case SOCKS_CONNECTING_TO_PROXY:
+            return PR_POLL_EXCEPT | PR_POLL_WRITE;
+        case SOCKS4_WRITE_CONNECT_REQUEST:
+        case SOCKS5_WRITE_AUTH_REQUEST:
+        case SOCKS5_WRITE_CONNECT_REQUEST:
+            return PR_POLL_WRITE;
+        case SOCKS4_READ_CONNECT_RESPONSE:
+        case SOCKS5_READ_AUTH_RESPONSE:
+        case SOCKS5_READ_CONNECT_RESPONSE_TOP:
+        case SOCKS5_READ_CONNECT_RESPONSE_BOTTOM:
+            return PR_POLL_READ;
+        default:
+            break;
+    }
+
+    return 0;
+}
+
+inline void
+nsSOCKSSocketInfo::WriteUint8(PRUint8 v)
+{
+    NS_ABORT_IF_FALSE(mDataLength + sizeof(v) <= BUFFER_SIZE,
+                      "Can't write that much data!");
+    mData[mDataLength] = v;
+    mDataLength += sizeof(v);
+}
+
+inline void
+nsSOCKSSocketInfo::WriteUint16(PRUint16 v)
+{
+    NS_ABORT_IF_FALSE(mDataLength + sizeof(v) <= BUFFER_SIZE,
+                      "Can't write that much data!");
+    memcpy(mData + mDataLength, &v, sizeof(v));
+    mDataLength += sizeof(v);
+}
+
+inline void
+nsSOCKSSocketInfo::WriteUint32(PRUint32 v)
+{
+    NS_ABORT_IF_FALSE(mDataLength + sizeof(v) <= BUFFER_SIZE,
+                      "Can't write that much data!");
+    memcpy(mData + mDataLength, &v, sizeof(v));
+    mDataLength += sizeof(v);
+}
+
+void
+nsSOCKSSocketInfo::WriteNetAddr(const PRNetAddr *addr)
+{
+    const char *ip = NULL;
+    PRUint32 len = 0;
+
+    if (PR_NetAddrFamily(addr) == PR_AF_INET) {
+        ip = (const char*)&addr->inet.ip;
+        len = sizeof(addr->inet.ip);
+    } else if (PR_NetAddrFamily(addr) == PR_AF_INET6) {
+        ip = (const char*)addr->ipv6.ip.pr_s6_addr;
+        len = sizeof(addr->ipv6.ip.pr_s6_addr);
+    }
+
+    NS_ABORT_IF_FALSE(ip != NULL, "Unknown address");
+    NS_ABORT_IF_FALSE(mDataLength + len <= BUFFER_SIZE,
+                      "Can't write that much data!");
  
-    return NS_OK;
+    memcpy(mData + mDataLength, ip, len);
+    mDataLength += len;
+}
 
+void
+nsSOCKSSocketInfo::WriteNetPort(const PRNetAddr *addr)
+{
+    WriteUint16(PR_NetAddrInetPort(addr));
+}
+
+void
+nsSOCKSSocketInfo::WriteString(const nsACString &str)
+{
+    NS_ABORT_IF_FALSE(mDataLength + str.Length() <= BUFFER_SIZE,
+                      "Can't write that much data!");
+    memcpy(mData + mDataLength, str.Data(), str.Length());
+    mDataLength += str.Length();
 }
 
+inline PRUint8
+nsSOCKSSocketInfo::ReadUint8()
+{
+    PRUint8 rv;
+    NS_ABORT_IF_FALSE(mReadOffset + sizeof(rv) <= mDataLength,
+                      "Not enough space to pop a uint8!");
+    rv = mData[mReadOffset];
+    mReadOffset += sizeof(rv);
+    return rv;
+}
+
+inline PRUint16
+nsSOCKSSocketInfo::ReadUint16()
+{
+    PRUint16 rv;
+    NS_ABORT_IF_FALSE(mReadOffset + sizeof(rv) <= mDataLength,
+                      "Not enough space to pop a uint16!");
+    memcpy(&rv, mData + mReadOffset, sizeof(rv));
+    mReadOffset += sizeof(rv);
+    return rv;
+}
+
+inline PRUint32
+nsSOCKSSocketInfo::ReadUint32()
+{
+    PRUint32 rv;
+    NS_ABORT_IF_FALSE(mReadOffset + sizeof(rv) <= mDataLength,
+                      "Not enough space to pop a uint32!");
+    memcpy(&rv, mData + mReadOffset, sizeof(rv));
+    mReadOffset += sizeof(rv);
+    return rv;
+}
+
+void
+nsSOCKSSocketInfo::ReadNetAddr(PRNetAddr *addr, PRUint16 fam)
+{
+    PRUint32 amt;
+    const PRUint8 *ip = mData + mReadOffset;
+
+    addr->raw.family = fam;
+    if (fam == PR_AF_INET) {
+        amt = sizeof(addr->inet.ip);
+        NS_ABORT_IF_FALSE(mReadOffset + amt <= mDataLength,
+                          "Not enough space to pop an ipv4 addr!");
+        memcpy(&addr->inet.ip, ip, amt);
+    } else if (fam == PR_AF_INET6) {
+        amt = sizeof(addr->ipv6.ip.pr_s6_addr);
+        NS_ABORT_IF_FALSE(mReadOffset + amt <= mDataLength,
+                          "Not enough space to pop an ipv6 addr!");
+        memcpy(addr->ipv6.ip.pr_s6_addr, ip, amt);
+    }
+
+    mReadOffset += amt;
+}
+
+void
+nsSOCKSSocketInfo::ReadNetPort(PRNetAddr *addr)
+{
+    addr->inet.port = ReadUint16();
+}
+
+void
+nsSOCKSSocketInfo::WantRead(PRUint32 sz)
+{
+    NS_ABORT_IF_FALSE(mDataIoPtr == NULL,
+                      "WantRead() called while I/O already in progress!");
+    NS_ABORT_IF_FALSE(mDataLength + sz <= BUFFER_SIZE,
+                      "Can't read that much data!");
+    mAmountToRead = sz;
+}
+
+PRStatus
+nsSOCKSSocketInfo::ReadFromSocket(PRFileDesc *fd)
+{
+    PRInt32 rc;
+    const PRUint8 *end;
+
+    if (!mAmountToRead) {
+        LOGDEBUG(("socks: ReadFromSocket(), nothing to do"));
+        return PR_SUCCESS;
+    }
+
+    if (!mDataIoPtr) {
+        mDataIoPtr = mData + mDataLength;
+        mDataLength += mAmountToRead;
+    }
+
+    end = mData + mDataLength;
+
+    while (mDataIoPtr < end) {
+        rc = PR_Read(fd, mDataIoPtr, end - mDataIoPtr);
+        if (rc <= 0) {
+            if (rc == 0) {
+                LOGERROR(("socks: proxy server closed connection"));
+                HandshakeFinished(PR_CONNECT_REFUSED_ERROR);
+                return PR_FAILURE;
+            } else if (PR_GetError() == PR_WOULD_BLOCK_ERROR) {
+                LOGDEBUG(("socks: ReadFromSocket(), want read"));
+            }
+            break;
+        }
+
+        mDataIoPtr += rc;
+    }
+
+    LOGDEBUG(("socks: ReadFromSocket(), have %u bytes total",
+             unsigned(mDataIoPtr - mData)));
+    if (mDataIoPtr == end) {
+        mDataIoPtr = nsnull;
+        mAmountToRead = 0;
+        mReadOffset = 0;
+        return PR_SUCCESS;
+    }
+
+    return PR_FAILURE;
+}
+
+PRStatus
+nsSOCKSSocketInfo::WriteToSocket(PRFileDesc *fd)
+{
+    PRInt32 rc;
+    const PRUint8 *end;
+
+    if (!mDataLength) {
+        LOGDEBUG(("socks: WriteToSocket(), nothing to do"));
+        return PR_SUCCESS;
+    }
+
+    if (!mDataIoPtr)
+        mDataIoPtr = mData;
+
+    end = mData + mDataLength;
+
+    while (mDataIoPtr < end) {
+        rc = PR_Write(fd, mDataIoPtr, end - mDataIoPtr);
+        if (rc < 0) {
+            if (PR_GetError() == PR_WOULD_BLOCK_ERROR) {
+                LOGDEBUG(("socks: WriteToSocket(), want write"));
+            }
+            break;
+        }
+        
+        mDataIoPtr += rc;
+    }
+
+    if (mDataIoPtr == end) {
+        mDataIoPtr = nsnull;
+        mDataLength = 0;
+        mReadOffset = 0;
+        return PR_SUCCESS;
+    }
+    
+    return PR_FAILURE;
+}
 
 static PRStatus
-nsSOCKSIOLayerConnect(PRFileDesc *fd, const PRNetAddr *addr, PRIntervalTime /*timeout*/)
+nsSOCKSIOLayerConnect(PRFileDesc *fd, const PRNetAddr *addr, PRIntervalTime to)
 {
+    PRStatus status;
+    PRNetAddr dst;
 
+    nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
+    if (info == NULL) return PR_FAILURE;
+
+    if (PR_NetAddrFamily(addr) == PR_AF_INET6 &&
+        PR_IsNetAddrType(addr, PR_IpAddrV4Mapped)) {
+        const PRUint8 *srcp;
+
+        LOGDEBUG(("socks: converting ipv4-mapped ipv6 address to ipv4"));
+
+        // copied from _PR_ConvertToIpv4NetAddr()
+        PR_InitializeNetAddr(PR_IpAddrAny, 0, &dst);
+        srcp = addr->ipv6.ip.pr_s6_addr;
+        memcpy(&dst.inet.ip, srcp + 12, 4);
+        dst.inet.family = PR_AF_INET;
+        dst.inet.port = addr->ipv6.port;
+    } else {
+        memcpy(&dst, addr, sizeof(dst));
+    }
+
+    info->SetDestinationAddr(&dst);
+    info->SetConnectTimeout(to);
+
+    do {
+        status = info->DoHandshake(fd, -1);
+    } while (status == PR_SUCCESS && !info->IsConnected());
+
+    return status;
+}
+
+static PRStatus
+nsSOCKSIOLayerConnectContinue(PRFileDesc *fd, PRInt16 oflags)
+{
     PRStatus status;
 
     nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
     if (info == NULL) return PR_FAILURE;
 
-    // First, we need to look up our proxy...
-    const nsCString &proxyHost = info->ProxyHost();
-
-    if (proxyHost.IsEmpty())
-        return PR_FAILURE;
-
-    PRInt32 socksVersion = info->Version();
-
-    LOGDEBUG(("nsSOCKSIOLayerConnect SOCKS %u; proxyHost: %s.", socksVersion, proxyHost.get()));
+    do { 
+        status = info->DoHandshake(fd, oflags);
+    } while (status == PR_SUCCESS && !info->IsConnected());
 
-    // Sync resolve the proxy hostname.
-    PRNetAddr proxyAddr;
-    nsCOMPtr<nsIDNSRecord> rec;
-    nsresult rv;
-    {
-        nsCOMPtr<nsIDNSService> dns = do_GetService(NS_DNSSERVICE_CONTRACTID);
-        if (!dns)
-            return PR_FAILURE;
-
-        rv = dns->Resolve(proxyHost, 0, getter_AddRefs(rec));
-        if (NS_FAILED(rv))
-            return PR_FAILURE;
-    }
-
-    info->SetInternalProxyAddr(&proxyAddr);
+    return status;
+}
 
-    // For now, we'll do this as a blocking connect,
-    // but with nspr 4.1, the necessary functions to
-    // do a non-blocking connect will be available
-
-    // Preserve the non-blocking state of the socket
-    PRBool nonblocking;
-    PRSocketOptionData sockopt;
-    sockopt.option = PR_SockOpt_Nonblocking;
-    status = PR_GetSocketOption(fd, &sockopt);
+static PRInt16
+nsSOCKSIOLayerPoll(PRFileDesc *fd, PRInt16 in_flags, PRInt16 *out_flags)
+{
+    nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
+    if (info == NULL) return PR_FAILURE;
 
-    if (PR_SUCCESS != status) {
-        LOGERROR(("PR_GetSocketOption() failed. status = %x.", status));
-        return status;
-    }
-
-    // Store blocking option
-    nonblocking = sockopt.value.non_blocking;
-
-    sockopt.option = PR_SockOpt_Nonblocking;
-    sockopt.value.non_blocking = PR_FALSE;
-    status = PR_SetSocketOption(fd, &sockopt);
-
-    if (PR_SUCCESS != status) {
-        LOGERROR(("PR_SetSocketOption() failed. status = %x.", status));
-        return status;
+    if (!info->IsConnected()) {
+        *out_flags = 0;
+        return info->GetPollFlags();
     }
 
-    // Now setup sockopts, so we can restore the value later.
-    sockopt.option = PR_SockOpt_Nonblocking;
-    sockopt.value.non_blocking = nonblocking;
-
-    // This connectWait should be long enough to connect to local proxy
-    // servers, but not much longer. Since this protocol negotiation
-    // uses blocking network calls, the app can appear to hang for a maximum
-    // of this time if the user presses the STOP button during the SOCKS
-    // connection negotiation. Note that this value only applies to the
-    // connecting to the SOCKS server: once the SOCKS connection has been
-    // established, the value is not used anywhere else.
-    PRIntervalTime connectWait = PR_SecondsToInterval(10);
-
-    // Connect to the proxy server.
-    PRInt32 addresses = 0;
-    do {
-        rv = rec->GetNextAddr(info->ProxyPort(), &proxyAddr);
-        if (NS_FAILED(rv)) {
-            status = PR_FAILURE;
-            break;
-        }
-        ++addresses;
-        status = fd->lower->methods->connect(fd->lower, &proxyAddr, connectWait);
-    } while (PR_SUCCESS != status);
-
-    if (PR_SUCCESS != status) {
-        LOGERROR(("Failed to TCP connect to the proxy server (%s): timeout = %d, status = %x, tried %d addresses.", proxyHost.get(), connectWait, status, addresses));
-        PR_SetSocketOption(fd, &sockopt);
-        return status;
-    }
-
-
-    // We are now connected to the SOCKS proxy server.
-    // Now we will negotiate a connection to the desired server.
-
-    // External IP address returned from ConnectSOCKS5(). Not supported in SOCKS4.
-    PRNetAddr extAddr;
-    PR_InitializeNetAddr(PR_IpAddrNull, 0, &extAddr);
-
-    NS_ASSERTION((socksVersion == 4) || (socksVersion == 5), "SOCKS Version must be selected");
-
-    // Try to connect via SOCKS 5.
-    if (socksVersion == 5) {
-        rv = ConnectSOCKS5(fd, addr, &extAddr, connectWait);
-
-        if (NS_FAILED(rv)) {
-            PR_SetSocketOption(fd, &sockopt);
-            return PR_FAILURE;
-        }
-
-    }
-
-    // Try to connect via SOCKS 4.
-    else {
-        rv = ConnectSOCKS4(fd, addr, connectWait);
-
-        if (NS_FAILED(rv)) {
-            PR_SetSocketOption(fd, &sockopt);
-            return PR_FAILURE;
-        }
-
-    }
-
-
-    info->SetDestinationAddr((PRNetAddr*)addr);
-    info->SetExternalProxyAddr(&extAddr);
-
-    // restore non-blocking option
-    PR_SetSocketOption(fd, &sockopt);
-
-    // we're set-up and connected.
-    // this socket can be used as normal now.
-
-    return PR_SUCCESS;
+    return fd->lower->methods->poll(fd->lower, in_flags, out_flags);
 }
 
 static PRStatus
 nsSOCKSIOLayerClose(PRFileDesc *fd)
 {
     nsSOCKSSocketInfo * info = (nsSOCKSSocketInfo*) fd->secret;
     PRDescIdentity id = PR_GetLayersIdentity(fd);
 
@@ -880,16 +1115,18 @@ nsSOCKSIOLayerAddToSocket(PRInt32 family
 
 
     if (firstTime)
     {
         nsSOCKSIOLayerIdentity		= PR_GetUniqueIdentity("SOCKS layer");
         nsSOCKSIOLayerMethods		= *PR_GetDefaultIOMethods();
 
         nsSOCKSIOLayerMethods.connect	= nsSOCKSIOLayerConnect;
+        nsSOCKSIOLayerMethods.connectcontinue	= nsSOCKSIOLayerConnectContinue;
+        nsSOCKSIOLayerMethods.poll	= nsSOCKSIOLayerPoll;
         nsSOCKSIOLayerMethods.bind	= nsSOCKSIOLayerBind;
         nsSOCKSIOLayerMethods.acceptread = nsSOCKSIOLayerAcceptRead;
         nsSOCKSIOLayerMethods.getsockname = nsSOCKSIOLayerGetName;
         nsSOCKSIOLayerMethods.getpeername = nsSOCKSIOLayerGetPeerName;
         nsSOCKSIOLayerMethods.accept	= nsSOCKSIOLayerAccept;
         nsSOCKSIOLayerMethods.listen	= nsSOCKSIOLayerListen;
         nsSOCKSIOLayerMethods.close	= nsSOCKSIOLayerClose;