bwilliams
f0c171c8a8
One of the two is required. --http-upgrade upgrades an HTTP connection to a WebSocket on /ws, as before, for job-search serve; --url and --origin go with it. --direct connects straight to ws-test serve over plain TCP, one JSON message per line, at --addr (default from server.host and server.port). --ping, --message, --listen and stdin work the same in both modes. The README's ws-test section explains the two modes.
116 lines
2.8 KiB
Go
116 lines
2.8 KiB
Go
package ws
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
)
|
|
|
|
// maxDirectLineBytes is the longest line a directConn reads.
|
|
const maxDirectLineBytes = 1 << 20
|
|
|
|
// errBinaryMessage is returned for a binary message, which a directConn
|
|
// can't send.
|
|
var errBinaryMessage = errors.New("only text messages can be sent directly")
|
|
|
|
// directConn is a plain TCP connection to ws-test serve, without an HTTP
|
|
// upgrade: each message is one line of JSON in both directions (see
|
|
// websocket.Hub.ServeTCP). It has the methods dial uses of a
|
|
// *websocket.Conn.
|
|
type directConn struct {
|
|
conn net.Conn
|
|
scanner *bufio.Scanner
|
|
}
|
|
|
|
// dialDirect connects to addr, giving up when ctx is done.
|
|
func dialDirect(ctx context.Context, addr string) (*directConn, error) {
|
|
var dialer net.Dialer
|
|
conn, err := dialer.DialContext(ctx, "tcp", addr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
scanner := bufio.NewScanner(conn)
|
|
scanner.Buffer(nil, maxDirectLineBytes)
|
|
|
|
return &directConn{
|
|
conn: conn,
|
|
scanner: scanner,
|
|
}, nil
|
|
}
|
|
|
|
// Read returns the next line as a text message, skipping the blank lines the
|
|
// server sends to check the connection. Cancelling ctx interrupts it. The
|
|
// server closing the connection is reported as io.EOF.
|
|
func (d *directConn) Read(ctx context.Context) (websocket.MessageType, []byte, error) {
|
|
stop := context.AfterFunc(ctx, func() {
|
|
d.conn.SetReadDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the read
|
|
})
|
|
defer stop()
|
|
|
|
for d.scanner.Scan() {
|
|
line := strings.TrimSuffix(d.scanner.Text(), "\r")
|
|
if strings.TrimSpace(line) == "" {
|
|
continue
|
|
}
|
|
return websocket.MessageText, []byte(line), nil
|
|
}
|
|
|
|
err := d.scanner.Err()
|
|
switch {
|
|
case ctx.Err() != nil:
|
|
return 0, nil, ctx.Err()
|
|
case err == nil:
|
|
return 0, nil, io.EOF
|
|
}
|
|
|
|
return 0, nil, err
|
|
}
|
|
|
|
// Write sends payload as one line. Cancelling ctx interrupts it.
|
|
func (d *directConn) Write(
|
|
ctx context.Context,
|
|
messageType websocket.MessageType,
|
|
payload []byte,
|
|
) error {
|
|
if messageType != websocket.MessageText {
|
|
return errBinaryMessage
|
|
}
|
|
|
|
stop := context.AfterFunc(ctx, func() {
|
|
d.conn.SetWriteDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the write
|
|
})
|
|
defer stop()
|
|
|
|
line := make([]byte, 0, len(payload)+1)
|
|
line = append(append(line, payload...), '\n')
|
|
_, err := d.conn.Write(line)
|
|
if err != nil && ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
// Close closes the connection. Plain TCP has no close handshake, so code and
|
|
// reason aren't sent.
|
|
func (d *directConn) Close(code websocket.StatusCode, reason string) error {
|
|
return d.CloseNow()
|
|
}
|
|
|
|
// CloseNow closes the connection. Closing it again isn't an error.
|
|
func (d *directConn) CloseNow() error {
|
|
err := d.conn.Close()
|
|
if errors.Is(err, net.ErrClosed) {
|
|
return nil
|
|
}
|
|
|
|
return err
|
|
}
|