[tor-commits] [flashproxy/master] Limit sent WebSocket messages to 1500 bytes.

[dcf at torproject.org]

commit 2f1243f897318e5b731d9e520d36b8ad37929fc0
Author: David Fifield <david at bamsoftware.com>
Date:   Sun Nov 11 18:50:19 2012 -0800

    Limit sent WebSocket messages to 1500 bytes.
---
 websocket-transport/websocket-client.go |   75 +++++++++++++++++++++++++++---
 1 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/websocket-transport/websocket-client.go b/websocket-transport/websocket-client.go
index 5c1c882..f8ae8ef 100644
--- a/websocket-transport/websocket-client.go
+++ b/websocket-transport/websocket-client.go
@@ -11,26 +11,83 @@ import (
 )
 
 const socksTimeout = 2
+const bufSiz = 1500
 
 func logDebug(format string, v ...interface{}) {
 	fmt.Fprintf(os.Stderr, format+"\n", v...)
 }
 
-func proxy(local *net.TCPConn, ws *websocket.Conn) error {
+func proxy(local *net.TCPConn, ws *websocket.Conn) {
+	var localToWs chan bool
+	var wsToLocal chan bool
+
 	// Local-to-WebSocket read loop.
+	localToWs = make(chan bool, 1)
 	go func() {
-		n, err := io.Copy(ws, local)
-		logDebug("end local-to-WebSocket %d %s", n, err)
+		buf := make([]byte, bufSiz)
+		var err error
+		for {
+			n, er := local.Read(buf[:])
+			if n > 0 {
+				ew := websocket.Message.Send(ws, buf[:n])
+				if ew != nil {
+					err = ew
+					break
+				}
+			}
+			if er != nil {
+				err = er
+				break
+			}
+		}
+		if err != nil && err != io.EOF {
+			logDebug("%s", err)
+		}
+		local.CloseRead()
+		ws.Close()
+
+		localToWs <- true
 	}()
 
 	// WebSocket-to-local read loop.
+	wsToLocal = make(chan bool, 1)
 	go func() {
-		n, err := io.Copy(local, ws)
-		logDebug("end WebSocket-to-local %d %s", n, err)
+		var buf []byte
+		var err error
+		for {
+			er := websocket.Message.Receive(ws, &buf)
+			if er != nil {
+				err = er
+				break
+			}
+			n, ew := local.Write(buf)
+			if ew != nil {
+				err = ew
+				break
+			}
+			if n != len(buf) {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if err != nil && err != io.EOF {
+			logDebug("%s", err)
+		}
+		local.CloseWrite()
+		ws.Close()
+
+		wsToLocal <- true
 	}()
 
-	select {}
-	return nil
+	// Select twice, once for each read loop.
+	select {
+	case <-localToWs:
+	case <-wsToLocal:
+	}
+	select {
+	case <-localToWs:
+	case <-wsToLocal:
+	}
 }
 
 func handleConnection(conn *net.TCPConn) error {
@@ -64,7 +121,9 @@ func handleConnection(conn *net.TCPConn) error {
 
 	sendSocks4aResponseGranted(conn, destAddr)
 
-	return proxy(conn, ws)
+	proxy(conn, ws)
+
+	return nil
 }
 
 func socksAcceptLoop(ln *net.TCPListener) error {