From e4790b5fa4d2e29892ab162a3148311c738320b3 Mon Sep 17 00:00:00 2001
From: Fumitoshi Ukai <ukai@google.com>
Date: Thu, 22 Sep 2011 21:49:24 -0400
Subject: [PATCH] websocket: add mutex to make websocket full-duplex

One benefit of websocket is that it is full-duplex so that it could
send and receive at the same time.
This CL makes websocket goroutine safe, so user could use websocket
both on goroutine for read and on goroutine for write.

R=golang-dev, rsc
CC=golang-dev
https://golang.org/cl/5058043
---
 src/pkg/websocket/hixie.go     |  2 ++
 src/pkg/websocket/hybi.go      |  4 ++++
 src/pkg/websocket/websocket.go | 11 +++++++++++
 3 files changed, 17 insertions(+)

diff --git a/src/pkg/websocket/hixie.go b/src/pkg/websocket/hixie.go
index b755a5c6a6..43de8a7800 100644
--- a/src/pkg/websocket/hixie.go
+++ b/src/pkg/websocket/hixie.go
@@ -235,6 +235,8 @@ func (handler *hixiFrameHandler) HandleFrame(frame frameReader) (r frameReader,
 }
 
 func (handler *hixiFrameHandler) WriteClose(_ int) (err os.Error) {
+	handler.conn.wio.Lock()
+	defer handler.conn.wio.Unlock()
 	closingFrame := []byte{'\xff', '\x00'}
 	handler.conn.buf.Write(closingFrame)
 	return handler.conn.buf.Flush()
diff --git a/src/pkg/websocket/hybi.go b/src/pkg/websocket/hybi.go
index cad47182d6..c4d990d6d8 100644
--- a/src/pkg/websocket/hybi.go
+++ b/src/pkg/websocket/hybi.go
@@ -288,6 +288,8 @@ func (handler *hybiFrameHandler) HandleFrame(frame frameReader) (r frameReader,
 }
 
 func (handler *hybiFrameHandler) WriteClose(status int) (err os.Error) {
+	handler.conn.wio.Lock()
+	defer handler.conn.wio.Unlock()
 	w, err := handler.conn.frameWriterFactory.NewFrameWriter(CloseFrame)
 	if err != nil {
 		return err
@@ -300,6 +302,8 @@ func (handler *hybiFrameHandler) WriteClose(status int) (err os.Error) {
 }
 
 func (handler *hybiFrameHandler) WritePong(msg []byte) (n int, err os.Error) {
+	handler.conn.wio.Lock()
+	defer handler.conn.wio.Unlock()
 	w, err := handler.conn.frameWriterFactory.NewFrameWriter(PongFrame)
 	if err != nil {
 		return 0, err
diff --git a/src/pkg/websocket/websocket.go b/src/pkg/websocket/websocket.go
index d57d1149c8..1855705c99 100644
--- a/src/pkg/websocket/websocket.go
+++ b/src/pkg/websocket/websocket.go
@@ -15,6 +15,7 @@ import (
 	"json"
 	"net"
 	"os"
+	"sync"
 	"url"
 )
 
@@ -147,9 +148,11 @@ type Conn struct {
 	buf *bufio.ReadWriter
 	rwc io.ReadWriteCloser
 
+	rio sync.Mutex
 	frameReaderFactory
 	frameReader
 
+	wio sync.Mutex
 	frameWriterFactory
 
 	frameHandler
@@ -163,6 +166,8 @@ type Conn struct {
 // will read the rest of the frame data.
 // it reads Text frame or Binary frame.
 func (ws *Conn) Read(msg []byte) (n int, err os.Error) {
+	ws.rio.Lock()
+	defer ws.rio.Unlock()
 again:
 	if ws.frameReader == nil {
 		frame, err := ws.frameReaderFactory.NewFrameReader()
@@ -191,6 +196,8 @@ again:
 // Write implements the io.Writer interface:
 // it writes data as a frame to the WebSocket connection.
 func (ws *Conn) Write(msg []byte) (n int, err os.Error) {
+	ws.wio.Lock()
+	defer ws.wio.Unlock()
 	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
 	if err != nil {
 		return 0, err
@@ -279,6 +286,8 @@ func (cd Codec) Send(ws *Conn, v interface{}) (err os.Error) {
 	if err != nil {
 		return err
 	}
+	ws.wio.Lock()
+	defer ws.wio.Unlock()
 	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
 	_, err = w.Write(data)
 	w.Close()
@@ -287,6 +296,8 @@ func (cd Codec) Send(ws *Conn, v interface{}) (err os.Error) {
 
 // Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
 func (cd Codec) Receive(ws *Conn, v interface{}) (err os.Error) {
+	ws.rio.Lock()
+	defer ws.rio.Unlock()
 	if ws.frameReader != nil {
 		_, err = io.Copy(ioutil.Discard, ws.frameReader)
 		if err != nil {