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

common/websocketconn/websocketconn.go

+78 −36
Changes for common/websocketconn/websocketconn.go: 78 added lines, 36 removed lines.
Original line number Original line Diff line number Diff line
@@ -10,61 +10,103 @@ import ( 
// An abstraction that makes an underlying WebSocket connection look like an

// An abstraction that makes an underlying WebSocket connection look like an

// io.ReadWriteCloser.

// io.ReadWriteCloser.

type Conn struct {

type Conn struct {

	Ws *websocket.Conn

	ws     *websocket.Conn

	r  io.Reader

	Reader io.Reader

	Writer io.Writer

}

}





// Implements io.Reader.

// Implements io.Reader.

func (conn *Conn) Read(b []byte) (n int, err error) {

func (conn *Conn) Read(b []byte) (n int, err error) {

	var opCode int

	return conn.Reader.Read(b)

	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

// Implements io.Writer.

			break

func (conn *Conn) Write(b []byte) (n int, err error) {

	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()

}

}





	n, err = conn.r.Read(b)

func readLoop(w io.Writer, ws *websocket.Conn) error {

	if err == io.EOF {

	for {

		// Message finished

		messageType, r, err := ws.NextReader()

		conn.r = nil

		if err != nil {

		err = nil

			return err

		}

		if messageType != websocket.BinaryMessage && messageType != websocket.TextMessage {

			continue

		}

		_, err = io.Copy(w, r)

		if err != nil {

			return err

		}

		}

	return

	}

	return nil

}

}





// Implements io.Writer.

func writeLoop(ws *websocket.Conn, r io.Reader) error {

func (conn *Conn) Write(b []byte) (n int, err error) {

	for {

	var w io.WriteCloser

		var buf [2048]byte

	if w, err = conn.Ws.NextWriter(websocket.BinaryMessage); err != nil {

		n, err := r.Read(buf[:])

		return

		if err != nil {

			return err

		}

		data := buf[:n]

		w, err := ws.NextWriter(websocket.BinaryMessage)

		if err != nil {

			return err

		}

		}

	if n, err = w.Write(b); err != nil {

		n, err = w.Write(data)

		return

		if err != nil {

			return err

		}

		}

		err = w.Close()

		err = w.Close()

	return

		if err != nil {

			return err

		}

	}

}

}





// Implements io.Closer.

// websocket.Conn methods start returning websocket.CloseError after the

func (conn *Conn) Close() error {

// connection has been closed. We want to instead interpret that as io.EOF, just

	// Ignore any error in trying to write a Close frame.

// as you would find with a normal net.Conn. This only converts

	_ = conn.Ws.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))

// websocket.CloseErrors with known codes; other codes like CloseProtocolError

	return conn.Ws.Close()

// 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.

// Create a new Conn.

func New(ws *websocket.Conn) *Conn {

func New(ws *websocket.Conn) *Conn {

	var conn Conn

	// Set up synchronous pipes to serialize reads and writes to the

	conn.Ws = ws

	// underlying websocket.Conn.

	return &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,

	}

}
 
}