[tor-commits] [flashproxy/master] readFrame.

[dcf at torproject.org]

commit 2021b406a3d5249502b0752415cc0f4db1e333d2
Author: David Fifield <david at bamsoftware.com>
Date:   Sun Nov 25 23:48:53 2012 -0800

    readFrame.
---
 websocket-transport/websocket-server.go |    1 +
 websocket-transport/websocket.go        |   93 +++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 0 deletions(-)

diff --git a/websocket-transport/websocket-server.go b/websocket-transport/websocket-server.go
index a2f3652..3874ea9 100644
--- a/websocket-transport/websocket-server.go
+++ b/websocket-transport/websocket-server.go
@@ -30,6 +30,7 @@ func startListener(addr *net.TCPAddr) (*net.TCPListener, error) {
 	go func() {
 		var config websocketConfig
 		config.Subprotocols = []string{"base64"}
+		config.MaxMessageSize = 1500
 		http.Handle("/", config.Handler(websocketHandler))
 		err = http.Serve(ln, nil)
 		if err != nil {
diff --git a/websocket-transport/websocket.go b/websocket-transport/websocket.go
index 9ac9754..fa0de7e 100644
--- a/websocket-transport/websocket.go
+++ b/websocket-transport/websocket.go
@@ -4,7 +4,10 @@ import (
 	"bufio"
 	"crypto/sha1"
 	"encoding/base64"
+	"encoding/binary"
+	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"strings"
@@ -12,6 +15,14 @@ import (
 
 type websocketConfig struct {
 	Subprotocols []string
+	MaxMessageSize uint64
+}
+
+func (config *websocketConfig) maxMessageSize() uint64 {
+	if config.MaxMessageSize == 0 {
+		return 64000
+	}
+	return config.MaxMessageSize
 }
 
 type websocket struct {
@@ -19,9 +30,90 @@ type websocket struct {
 	Bufrw       *bufio.ReadWriter
 	// Whether we are a client or a server implications for masking.
 	IsClient    bool
+	MaxMessageSize uint64
 	Subprotocol string
 }
 
+type websocketFrame struct {
+	Fin bool
+	Opcode byte
+	Payload []byte
+}
+
+func (frame *websocketFrame) IsControl() bool {
+	return (frame.Opcode & 0x08) != 0
+}
+
+func applyMask(payload []byte, maskKey [4]byte) {
+	for i, _ := range payload {
+		payload[i] = payload[i] ^ maskKey[i % 4]
+	}
+}
+
+func (ws *websocket) ReadFrame() (frame websocketFrame, err error) {
+	var b byte
+	err = binary.Read(ws.Bufrw, binary.BigEndian, &b)
+	if err != nil {
+		return
+	}
+	frame.Fin = (b & 0x80) != 0
+	frame.Opcode = b & 0x0f
+	err = binary.Read(ws.Bufrw, binary.BigEndian, &b)
+	if err != nil {
+		return
+	}
+	masked := (b & 0x80) != 0
+
+	payloadLen := uint64(b & 0x7f)
+	if payloadLen == 126 {
+		var short uint16
+		err = binary.Read(ws.Bufrw, binary.BigEndian, &short)
+		if err != nil {
+			return
+		}
+		payloadLen = uint64(short)
+	} else if payloadLen == 127 {
+		var long uint64
+		err = binary.Read(ws.Bufrw, binary.BigEndian, &long)
+		if err != nil {
+			return
+		}
+		payloadLen = long
+	}
+	if payloadLen > ws.MaxMessageSize {
+		err = errors.New(fmt.Sprintf("frame payload length of %d exceeds maximum of %d", payloadLen, ws.MaxMessageSize))
+		return
+	}
+
+	maskKey := [4]byte{}
+	if masked {
+		if ws.IsClient {
+			err = errors.New("client got masked frame")
+			return
+		}
+		err = binary.Read(ws.Bufrw, binary.BigEndian, &maskKey)
+		if err != nil {
+			return
+		}
+	} else {
+		if !ws.IsClient {
+			err = errors.New("server got unmasked frame")
+			return
+		}
+	}
+
+	frame.Payload = make([]byte, payloadLen)
+	_, err = io.ReadFull(ws.Bufrw, frame.Payload)
+	if err != nil {
+		return
+	}
+	if masked {
+		applyMask(frame.Payload, maskKey)
+	}
+
+	return frame, nil
+}
+
 func commaSplit(s string) []string {
 	var result []string
 	if strings.TrimSpace(s) == "" {
@@ -129,6 +221,7 @@ func (handler *WebSocketHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.
 	ws.Conn = conn
 	ws.Bufrw = bufrw
 	ws.IsClient = false
+	ws.MaxMessageSize = handler.config.MaxMessageSize
 
 	// See RFC 6455 section 4.2.2, item 5 for these steps.