a81d316635
All checks were successful
gitea/video-server-backend/pipeline/head This commit looks good
60 lines
1.2 KiB
Go
60 lines
1.2 KiB
Go
package bufferedsocket
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type BufferedSocket struct {
|
|
conn *websocket.Conn
|
|
bufferSize int
|
|
buffer []byte
|
|
bufferLen int
|
|
}
|
|
|
|
func NewBufferedSocket(conn *websocket.Conn, bufferSize int) *BufferedSocket {
|
|
return &BufferedSocket{
|
|
conn: conn,
|
|
bufferSize: bufferSize,
|
|
buffer: make([]byte, 0, bufferSize),
|
|
}
|
|
}
|
|
|
|
func (bs *BufferedSocket) WriteMessage(message []byte) error {
|
|
if len(message) > bs.bufferSize {
|
|
return errors.New("message larger than buffer size")
|
|
}
|
|
|
|
if bs.bufferLen+len(message) > bs.bufferSize {
|
|
err := bs.Flush()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
bs.buffer = append(bs.buffer, message...)
|
|
bs.bufferLen += len(message)
|
|
return nil
|
|
}
|
|
|
|
func (bs *BufferedSocket) Flush() error {
|
|
if bs.bufferLen > 0 {
|
|
err := bs.conn.WriteMessage(websocket.TextMessage, bs.buffer[:bs.bufferLen])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
bs.buffer = make([]byte, 0, bs.bufferSize)
|
|
bs.bufferLen = 0
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (bs *BufferedSocket) Close() error {
|
|
err := bs.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
bs.conn.Close()
|
|
return nil
|
|
}
|