gotty/server/ws_wrapper.go

42 lines
715 B
Go
Raw Normal View History

2017-08-24 05:40:28 +00:00
package server
import (
2023-10-21 15:49:19 +00:00
"io"
2022-03-31 17:42:51 +00:00
2017-08-24 05:40:28 +00:00
"github.com/gorilla/websocket"
2022-03-31 17:42:51 +00:00
"github.com/pkg/errors"
2017-08-24 05:40:28 +00:00
)
type wsWrapper struct {
*websocket.Conn
}
func (wsw *wsWrapper) Write(p []byte) (n int, err error) {
writer, err := wsw.Conn.NextWriter(websocket.TextMessage)
if err != nil {
return 0, err
}
defer writer.Close()
return writer.Write(p)
}
func (wsw *wsWrapper) Read(p []byte) (n int, err error) {
for {
msgType, reader, err := wsw.Conn.NextReader()
if err != nil {
return 0, err
}
if msgType != websocket.TextMessage {
continue
}
2023-10-21 15:49:19 +00:00
b, err := io.ReadAll(reader)
2022-03-31 17:42:51 +00:00
if len(b) > len(p) {
return 0, errors.Wrapf(err, "Client message exceeded buffer size")
}
n = copy(p, b)
return n, err
2017-08-24 05:40:28 +00:00
}
}