commit 2021b406a3d5249502b0752415cc0f4db1e333d2 Author: David Fifield david@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.