Skip to content
Snippets Groups Projects
Commit 01e28aa4 authored by David Fifield's avatar David Fifield
Browse files

Rewrite websocketconn with synchronous pipes.

Makes the following changes:
 * permits concurrent Read/Write/Close
 * converts certain CloseErrors into io.EOF

https://bugs.torproject.org/33144
parent 5708a1d5
No related branches found
Tags webext-0.0.6
No related merge requests found
......@@ -10,61 +10,103 @@ import (
// An abstraction that makes an underlying WebSocket connection look like an
// io.ReadWriteCloser.
type Conn struct {
Ws *websocket.Conn
r io.Reader
ws *websocket.Conn
Reader io.Reader
Writer io.Writer
}
// Implements io.Reader.
func (conn *Conn) Read(b []byte) (n int, err error) {
var opCode int
if conn.r == nil {
// New message
var r io.Reader
for {
if opCode, r, err = conn.Ws.NextReader(); err != nil {
return
}
if opCode != websocket.BinaryMessage && opCode != websocket.TextMessage {
continue
}
conn.r = r
break
}
}
n, err = conn.r.Read(b)
if err == io.EOF {
// Message finished
conn.r = nil
err = nil
}
return
return conn.Reader.Read(b)
}
// Implements io.Writer.
func (conn *Conn) Write(b []byte) (n int, err error) {
var w io.WriteCloser
if w, err = conn.Ws.NextWriter(websocket.BinaryMessage); err != nil {
return
}
if n, err = w.Write(b); err != nil {
return
}
err = w.Close()
return
return conn.Writer.Write(b)
}
// Implements io.Closer.
func (conn *Conn) Close() error {
// Ignore any error in trying to write a Close frame.
_ = conn.Ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))
return conn.Ws.Close()
_ = conn.ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))
return conn.ws.Close()
}
func readLoop(w io.Writer, ws *websocket.Conn) error {
for {
messageType, r, err := ws.NextReader()
if err != nil {
return err
}
if messageType != websocket.BinaryMessage && messageType != websocket.TextMessage {
continue
}
_, err = io.Copy(w, r)
if err != nil {
return err
}
}
return nil
}
func writeLoop(ws *websocket.Conn, r io.Reader) error {
for {
var buf [2048]byte
n, err := r.Read(buf[:])
if err != nil {
return err
}
data := buf[:n]
w, err := ws.NextWriter(websocket.BinaryMessage)
if err != nil {
return err
}
n, err = w.Write(data)
if err != nil {
return err
}
err = w.Close()
if err != nil {
return err
}
}
}
// websocket.Conn methods start returning websocket.CloseError after the
// connection has been closed. We want to instead interpret that as io.EOF, just
// as you would find with a normal net.Conn. This only converts
// websocket.CloseErrors with known codes; other codes like CloseProtocolError
// and CloseAbnormalClosure will still be reported as anomalous.
func closeErrorToEOF(err error) error {
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseNoStatusReceived) {
err = io.EOF
}
return err
}
// Create a new Conn.
func New(ws *websocket.Conn) *Conn {
var conn Conn
conn.Ws = ws
return &conn
// Set up synchronous pipes to serialize reads and writes to the
// underlying websocket.Conn.
//
// https://godoc.org/github.com/gorilla/websocket#hdr-Concurrency
// "Connections support one concurrent reader and one concurrent writer.
// Applications are responsible for ensuring that no more than one
// goroutine calls the write methods (NextWriter, etc.) concurrently and
// that no more than one goroutine calls the read methods (NextReader,
// etc.) concurrently. The Close and WriteControl methods can be called
// concurrently with all other methods."
pr1, pw1 := io.Pipe()
go func() {
pw1.CloseWithError(closeErrorToEOF(readLoop(pw1, ws)))
}()
pr2, pw2 := io.Pipe()
go func() {
pr2.CloseWithError(closeErrorToEOF(writeLoop(ws, pr2)))
}()
return &Conn{
ws: ws,
Reader: pr1,
Writer: pw2,
}
}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment