Cloak/internal/util/websocket.go

83 lines
1.5 KiB
Go
Raw Normal View History

2019-09-01 00:33:34 +00:00
package util
import (
2019-09-01 19:23:45 +00:00
"errors"
2019-09-01 00:33:34 +00:00
"github.com/gorilla/websocket"
2019-09-01 19:23:45 +00:00
"io"
"net"
"sync"
2019-09-01 00:33:34 +00:00
"time"
)
// WebSocketConn implements io.ReadWriteCloser
// it makes websocket.Conn binary-oriented
2019-09-01 00:33:34 +00:00
type WebSocketConn struct {
*websocket.Conn
2019-09-01 19:23:45 +00:00
writeM sync.Mutex
2019-09-01 00:33:34 +00:00
}
func (ws *WebSocketConn) Write(data []byte) (int, error) {
2019-09-01 19:23:45 +00:00
ws.writeM.Lock()
2019-09-01 00:33:34 +00:00
err := ws.WriteMessage(websocket.BinaryMessage, data)
2019-09-01 19:23:45 +00:00
ws.writeM.Unlock()
2019-09-01 00:33:34 +00:00
if err != nil {
return 0, err
} else {
return len(data), nil
}
}
2019-09-01 19:23:45 +00:00
func (ws *WebSocketConn) Read(buf []byte) (n int, err error) {
t, r, err := ws.NextReader()
2019-09-01 00:33:34 +00:00
if err != nil {
return 0, err
}
2019-09-01 19:23:45 +00:00
if t != websocket.BinaryMessage {
return 0, nil
}
// Read until io.EOL for one full message
for {
var read int
read, err = r.Read(buf[n:])
if err != nil {
if err == io.EOF {
err = nil
break
} else {
break
}
} else {
// There may be data available to read but n == len(buf)-1, read==0 because buffer is full
if read == 0 {
err = errors.New("nothing more is read. message may be larger than buffer")
break
}
}
n += read
}
return
}
func (ws *WebSocketConn) Close() error {
ws.writeM.Lock()
defer ws.writeM.Unlock()
return ws.Conn.Close()
2019-09-01 00:33:34 +00:00
}
func (ws *WebSocketConn) SetDeadline(t time.Time) error {
err := ws.SetReadDeadline(t)
if err != nil {
return err
}
err = ws.SetWriteDeadline(t)
if err != nil {
return err
}
return nil
}
2019-09-01 19:23:45 +00:00
// ws unit reader
func ReadWebSocket(conn net.Conn, buffer []byte) (n int, err error) {
return conn.Read(buffer)
}