[tor-commits] [flashproxy/master] Add a websocketConn with Read and Write methods.

[dcf at torproject.org]

commit e05c31d80217403d0a01b46b5954ff816bd34eb0
Author: David Fifield <david at bamsoftware.com>
Date:   Mon Nov 26 01:16:46 2012 -0800

    Add a websocketConn with Read and Write methods.
---
 websocket-transport/websocket-server.go |   65 +++++++++++++++++++++++++++++++
 1 files changed, 65 insertions(+), 0 deletions(-)

diff --git a/websocket-transport/websocket-server.go b/websocket-transport/websocket-server.go
index 3874ea9..c005d7e 100644
--- a/websocket-transport/websocket-server.go
+++ b/websocket-transport/websocket-server.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"encoding/base64"
+	"errors"
 	"fmt"
 	"net"
 	"net/http"
@@ -18,6 +20,69 @@ func logDebug(format string, v ...interface{}) {
 	fmt.Fprintf(os.Stderr, format+"\n", v...)
 }
 
+type websocketConn struct {
+	Ws *websocket
+	Base64 bool
+	messageBuf []byte
+}
+
+func (conn *websocketConn) Read(b []byte) (n int, err error) {
+	for len(conn.messageBuf) == 0 {
+		var m websocketMessage
+		m, err = conn.Ws.ReadMessage()
+		if err != nil {
+			return
+		}
+		if conn.Base64 {
+			if m.Opcode != 1 {
+				err = errors.New(fmt.Sprintf("got non-text opcode %d with the base64 subprotocol", m.Opcode))
+				return
+			}
+			conn.messageBuf = make([]byte, base64.StdEncoding.DecodedLen(len(m.Payload)))
+			var num int
+			num, err = base64.StdEncoding.Decode(conn.messageBuf, m.Payload)
+			if err != nil {
+				return
+			}
+			conn.messageBuf = conn.messageBuf[:num]
+		} else {
+			if m.Opcode != 2 {
+				err = errors.New(fmt.Sprintf("got non-binary opcode %d with no subprotocol", m.Opcode))
+				return
+			}
+			conn.messageBuf = m.Payload
+		}
+	}
+
+	n = copy(b, conn.messageBuf)
+	conn.messageBuf = conn.messageBuf[n:]
+
+	return
+}
+
+func (conn *websocketConn) Write(b []byte) (n int, err error) {
+	if conn.Base64 {
+		buf := make([]byte, base64.StdEncoding.EncodedLen(len(b)))
+		base64.StdEncoding.Encode(buf, b)
+		err = conn.Ws.WriteMessage(1, buf)
+		if err != nil {
+			return
+		}
+		n = len(b)
+	} else {
+		err = conn.Ws.WriteMessage(2, b)
+		n = len(b)
+	}
+	return
+}
+
+func NewWebsocketConn(ws *websocket) websocketConn {
+	var conn websocketConn
+	conn.Ws = ws
+	conn.Base64 = (ws.Subprotocol == "base64")
+	return conn
+}
+
 func websocketHandler(ws *websocket) {
 	fmt.Printf("blah\n")
 }