[tor-commits] [flashproxy/master] Better factoring of SOCKS requests.

[dcf at torproject.org]

commit 824d87ebdd24ffb90f5f500383d9c7926561b3cc
Author: David Fifield <david at bamsoftware.com>
Date:   Mon Nov 26 23:38:39 2012 -0800

    Better factoring of SOCKS requests.
---
 websocket-transport/socks.go            |   19 ++++++++++++++
 websocket-transport/websocket-client.go |   42 +++++++++++++-----------------
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/websocket-transport/socks.go b/websocket-transport/socks.go
index e8ef51f..89d83ad 100644
--- a/websocket-transport/socks.go
+++ b/websocket-transport/socks.go
@@ -18,6 +18,25 @@ const (
 	socksRequestFailed   = 0x5b
 )
 
+// Read a SOCKS4a connect request, and call the given connect callback with the
+// requested destination string. If the callback returns an error, sends a SOCKS
+// request failed message. Otherwise, sends a SOCKS request granted message for
+// the destination address returned by the callback.
+func AwaitSocks4aConnect(conn *net.TCPConn, connect func(string) (*net.TCPAddr, error)) error {
+	dest, err := ReadSocks4aConnect(conn)
+	if err != nil {
+		SendSocks4aResponseFailed(conn)
+		return err
+	}
+	destAddr, err := connect(dest)
+	if err != nil {
+		SendSocks4aResponseFailed(conn)
+		return err
+	}
+	SendSocks4aResponseGranted(conn, destAddr)
+	return nil
+}
+
 // Read a SOCKS4a connect request. Returns a "host:port" string.
 func ReadSocks4aConnect(s io.Reader) (string, error) {
 	r := bufio.NewReader(s)
diff --git a/websocket-transport/websocket-client.go b/websocket-transport/websocket-client.go
index 1259a50..129e90a 100644
--- a/websocket-transport/websocket-client.go
+++ b/websocket-transport/websocket-client.go
@@ -101,36 +101,30 @@ func handleConnection(conn *net.TCPConn) error {
 		handlerChan <- -1
 	}()
 
-	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
-	dest, err := ReadSocks4aConnect(conn)
-	if err != nil {
-		SendSocks4aResponseFailed(conn)
-		return err
-	}
-	// Disable deadline.
-	conn.SetDeadline(time.Time{})
-	logDebug("SOCKS request for %s", dest)
-
-	// We need the parsed IP and port for the SOCKS reply.
-	destAddr, err := net.ResolveTCPAddr("tcp", dest)
-	if err != nil {
-		SendSocks4aResponseFailed(conn)
-		return err
-	}
+	var ws *websocket.Conn
 
-	wsUrl := url.URL{Scheme: "ws", Host: dest}
-	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
+	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
+	err := AwaitSocks4aConnect(conn, func(dest string) (*net.TCPAddr, error) {
+		// Disable deadline.
+		conn.SetDeadline(time.Time{})
+		logDebug("SOCKS request for %s", dest)
+		destAddr, err := net.ResolveTCPAddr("tcp", dest)
+		if err != nil {
+			return nil, err
+		}
+		wsUrl := url.URL{Scheme: "ws", Host: dest}
+		ws, err = websocket.Dial(wsUrl.String(), "", wsUrl.String())
+		if err != nil {
+			return nil, err
+		}
+		logDebug("WebSocket connection to %s", ws.Config().Location.String())
+		return destAddr, nil
+	})
 	if err != nil {
-		SendSocks4aResponseFailed(conn)
 		return err
 	}
 	defer ws.Close()
-	logDebug("WebSocket connection to %s", ws.Config().Location.String())
-
-	SendSocks4aResponseGranted(conn, destAddr)
-
 	proxy(conn, ws)
-
 	return nil
 }