初始化提交

util/ws/websocket.go

@@ -0,0 +1,76 @@
+package ws
+
+import "github.com/gorilla/websocket"
+
+type Packet struct {
+	isBinary bool
+	data     []byte
+}
+
+func Text(buf []byte) *Packet {
+	return &Packet{
+		isBinary: false,
+		data:     buf,
+	}
+}
+
+func Binary(buf []byte) *Packet {
+	return &Packet{
+		isBinary: true,
+		data:     buf,
+	}
+}
+
+type Conn struct {
+	WriteChan chan *Packet
+	ReadChan  chan []byte
+	conn      *websocket.Conn
+	close     bool
+}
+
+func Wrap(conn *websocket.Conn) *Conn {
+	return &Conn{
+		WriteChan: make(chan *Packet),
+		ReadChan:  make(chan []byte),
+		conn:      conn,
+	}
+}
+
+func (conn *Conn) ReadLoop() {
+	for !conn.close {
+		_, message, err := conn.conn.ReadMessage()
+		if err != nil {
+			break
+		}
+		conn.ReadChan <- message
+	}
+}
+
+func (conn *Conn) WriteLoop() {
+	for {
+		select {
+		case message := <-conn.WriteChan:
+			typ := websocket.TextMessage
+			if message.isBinary {
+				typ = websocket.BinaryMessage
+			}
+			err := conn.conn.WriteMessage(typ, message.data)
+			if err != nil {
+				break
+			}
+		default:
+			if conn.close {
+				return
+			}
+		}
+	}
+}
+
+func (conn *Conn) Close() error {
+	conn.close = true
+	return conn.conn.Close()
+}
+
+func (conn *Conn) IsClosed() bool {
+	return conn.close
+}