From fce54343ea6ce797a98c48d5646dd3170ae3bd47 Mon Sep 17 00:00:00 2001
From: Edmund Huber <edmund.pgh@gmail.com>
Date: Sat, 6 Jun 2015 16:11:06 -0700
Subject: [PATCH] touchups to irc/websocket.go per review comments

---
 irc/server.go    |  6 +----
 irc/websocket.go | 57 ++++++++++++++++--------------------------------
 2 files changed, 20 insertions(+), 43 deletions(-)

diff --git a/irc/server.go b/irc/server.go
index 9b05ba2f..33fa2fb7 100644
--- a/irc/server.go
+++ b/irc/server.go
@@ -232,11 +232,7 @@ func (s *Server) wslisten(addr string) {
 			return
 		}
 
-		wsc := WSContainer{
-			conn: ws,
-		}
-
-		s.newConns <- wsc
+		s.newConns <- WSContainer{ws}
 	})
 	go func() {
 		Log.info.Printf("%s listening on %s", s, addr)
diff --git a/irc/websocket.go b/irc/websocket.go
index 74e21c5d..6bf38f64 100644
--- a/irc/websocket.go
+++ b/irc/websocket.go
@@ -2,7 +2,6 @@ package irc
 
 import (
 	"github.com/gorilla/websocket"
-	"net"
 	"net/http"
 	"time"
 )
@@ -10,56 +9,38 @@ import (
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
-	/* If a WS session contains sensitive information, and you choose to use
-	   cookies for authentication (during the HTTP(S) upgrade request), then
-	   you should check that Origin is a domain under your control. If it
-	   isn't, then it is possible for users of your site, visiting a naughty
-	   Origin, to have a WS opened using their credentials. See
-	   http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html#main.
-	   We don't care about Origin because the (IRC) authentication is contained
-	   in the WS stream -- the WS session is not privileged when it is opened.
-	*/
+	// If a WS session contains sensitive information, and you choose to use
+	// cookies for authentication (during the HTTP(S) upgrade request), then
+	// you should check that Origin is a domain under your control. If it
+	// isn't, then it is possible for users of your site, visiting a naughty
+	// Origin, to have a WS opened using their credentials. See
+	// http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html#main.
+	// We don't care about Origin because the (IRC) authentication is contained
+	// in the WS stream -- the WS session is not privileged when it is opened.
 	CheckOrigin:     func(r *http.Request) bool { return true },
 }
 
 type WSContainer struct {
-	conn *websocket.Conn
-}
-
-func (this WSContainer) Close() error {
-	return this.conn.Close()
-}
-
-func (this WSContainer) LocalAddr() net.Addr {
-	return this.conn.LocalAddr()
-}
-
-func (this WSContainer) RemoteAddr() net.Addr {
-	return this.conn.RemoteAddr()
+	*websocket.Conn
 }
 
 func (this WSContainer) Read(msg []byte) (int, error) {
-	_, tmp, err := this.conn.ReadMessage()
-	str := (string)(tmp)
-	n := copy(msg, ([]byte)(str+CRLF+CRLF))
-	return n, err
+	ty, bytes, err := this.ReadMessage()
+	if ty == websocket.TextMessage {
+		n := copy(msg, []byte(string(bytes)+CRLF+CRLF))
+		return n, err
+	}
+	// Binary, and other kinds of messages, are thrown away.
+	return 0, nil
 }
 
 func (this WSContainer) Write(msg []byte) (int, error) {
-	err := this.conn.WriteMessage(1, msg)
+	err := this.WriteMessage(websocket.TextMessage, msg)
 	return len(msg), err
 }
 
 func (this WSContainer) SetDeadline(t time.Time) error {
-	err := this.conn.SetWriteDeadline(t)
-	err = this.conn.SetReadDeadline(t)
+	err := this.SetWriteDeadline(t)
+	err = this.SetReadDeadline(t)
 	return err
 }
-
-func (this WSContainer) SetReadDeadline(t time.Time) error {
-	return this.conn.SetReadDeadline(t)
-}
-
-func (this WSContainer) SetWriteDeadline(t time.Time) error {
-	return this.conn.SetWriteDeadline(t)
-}