130 lines
2.4 KiB
Go
130 lines
2.4 KiB
Go
package websocket
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
ws "github.com/gorilla/websocket"
|
|
)
|
|
|
|
// GracefulCloseTimeout is the time to wait trying to gracefully close a
|
|
// connection before simply cutting it.
|
|
var GracefulCloseTimeout = 100 * time.Millisecond
|
|
|
|
var _ net.Conn = (*Conn)(nil)
|
|
|
|
// Conn implements net.Conn interface for gorilla/websocket.
|
|
type Conn struct {
|
|
*ws.Conn
|
|
DefaultMessageType int
|
|
reader io.Reader
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
func (c *Conn) Read(b []byte) (int, error) {
|
|
if c.reader == nil {
|
|
if err := c.prepNextReader(); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
for {
|
|
n, err := c.reader.Read(b)
|
|
switch err {
|
|
case io.EOF:
|
|
c.reader = nil
|
|
|
|
if n > 0 {
|
|
return n, nil
|
|
}
|
|
|
|
if err := c.prepNextReader(); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
// explicitly looping
|
|
default:
|
|
return n, err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Conn) prepNextReader() error {
|
|
t, r, err := c.Conn.NextReader()
|
|
if err != nil {
|
|
if wserr, ok := err.(*ws.CloseError); ok {
|
|
if wserr.Code == 1000 || wserr.Code == 1005 {
|
|
return io.EOF
|
|
}
|
|
}
|
|
return err
|
|
}
|
|
|
|
if t == ws.CloseMessage {
|
|
return io.EOF
|
|
}
|
|
|
|
c.reader = r
|
|
return nil
|
|
}
|
|
|
|
func (c *Conn) Write(b []byte) (n int, err error) {
|
|
if err := c.Conn.WriteMessage(c.DefaultMessageType, b); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return len(b), nil
|
|
}
|
|
|
|
// Close closes the connection. Only the first call to Close will receive the
|
|
// close error, subsequent and concurrent calls will return nil.
|
|
// This method is thread-safe.
|
|
func (c *Conn) Close() error {
|
|
var err error
|
|
c.closeOnce.Do(func() {
|
|
err1 := c.Conn.WriteControl(ws.CloseMessage, nil, time.Now().Add(GracefulCloseTimeout))
|
|
err2 := c.Conn.Close()
|
|
switch {
|
|
case err1 != nil:
|
|
err = err1
|
|
case err2 != nil:
|
|
err = err2
|
|
}
|
|
})
|
|
return err
|
|
}
|
|
|
|
func (c *Conn) LocalAddr() net.Addr {
|
|
return NewAddr(c.Conn.LocalAddr().String())
|
|
}
|
|
|
|
func (c *Conn) RemoteAddr() net.Addr {
|
|
return NewAddr(c.Conn.RemoteAddr().String())
|
|
}
|
|
|
|
func (c *Conn) SetDeadline(t time.Time) error {
|
|
if err := c.SetReadDeadline(t); err != nil {
|
|
return err
|
|
}
|
|
|
|
return c.SetWriteDeadline(t)
|
|
}
|
|
|
|
func (c *Conn) SetReadDeadline(t time.Time) error {
|
|
return c.Conn.SetReadDeadline(t)
|
|
}
|
|
|
|
func (c *Conn) SetWriteDeadline(t time.Time) error {
|
|
return c.Conn.SetWriteDeadline(t)
|
|
}
|
|
|
|
// NewConn creates a Conn given a regular gorilla/websocket Conn.
|
|
func NewConn(raw *ws.Conn) *Conn {
|
|
return &Conn{
|
|
Conn: raw,
|
|
DefaultMessageType: ws.BinaryMessage,
|
|
}
|
|
}
|