68 lines
1.3 KiB
Go
68 lines
1.3 KiB
Go
package client
|
|
|
|
import (
|
|
"io"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
var (
|
|
// DefaultReadBufferSize is the default maximum read buffer size.
|
|
DefaultReadBufferSize = 25 * 1024 * 1024
|
|
|
|
// DefaultWriteBufferSize is the default maximum write buffer size.
|
|
DefaultWriteBufferSize = 10 * 1024 * 1024
|
|
)
|
|
|
|
// Transport is the common interface to send/receive messages to a target.
|
|
type Transport interface {
|
|
Read() ([]byte, error)
|
|
Write([]byte) error
|
|
io.Closer
|
|
}
|
|
|
|
// Conn wraps a gorilla/websocket.Conn connection.
|
|
type Conn struct {
|
|
*websocket.Conn
|
|
}
|
|
|
|
// Read reads the next websocket message.
|
|
func (c *Conn) Read() ([]byte, error) {
|
|
_, buf, err := c.ReadMessage()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
// Write writes a websocket message.
|
|
func (c *Conn) Write(buf []byte) error {
|
|
return c.WriteMessage(websocket.TextMessage, buf)
|
|
}
|
|
|
|
// Dial dials the specified target's websocket URL.
|
|
//
|
|
// Note: uses gorilla/websocket.
|
|
func Dial(urlstr string, opts ...DialOption) (Transport, error) {
|
|
d := &websocket.Dialer{
|
|
ReadBufferSize: DefaultReadBufferSize,
|
|
WriteBufferSize: DefaultWriteBufferSize,
|
|
}
|
|
|
|
// apply opts
|
|
for _, o := range opts {
|
|
o(d)
|
|
}
|
|
|
|
// connect
|
|
conn, _, err := d.Dial(urlstr, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Conn{conn}, nil
|
|
}
|
|
|
|
// DialOption is a dial option.
|
|
type DialOption func(*websocket.Dialer)
|