83 lines
1.2 KiB
Go
83 lines
1.2 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"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
|
|
ctx context.Context
|
|
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 {
|
|
conn.Close()
|
|
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.ReadChan <- nil
|
|
conn.close = true
|
|
return conn.conn.Close()
|
|
}
|
|
|
|
func (conn *Conn) IsClosed() bool {
|
|
return conn.close
|
|
}
|