feat: port to coder/websocket + fix conn

This commit is contained in:
minish 2025-12-29 21:40:10 -05:00
parent 3b9bd07185
commit 60635005e7
Signed by: min
SSH Key Fingerprint: SHA256:h4k7JNrfe1dzv1WE3oGVeAY9DPSZXIu3/j89+6DtHWE
3 changed files with 17 additions and 15 deletions

2
go.mod
View File

@ -2,4 +2,4 @@ module git.min.rip/min/websocket-client-go
go 1.23.6 go 1.23.6
require golang.org/x/net v0.35.0 require github.com/coder/websocket v1.8.14

4
go.sum
View File

@ -1,2 +1,2 @@
golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8= github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g=
golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk= github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg=

View File

@ -2,13 +2,15 @@ package websocket
import ( import (
"context" "context"
"encoding/json" "errors"
"fmt" "fmt"
"net/http"
"net/url" "net/url"
"sync" "sync"
"time" "time"
"golang.org/x/net/websocket" "github.com/coder/websocket"
"github.com/coder/websocket/wsjson"
) )
type WebSocketClient struct { type WebSocketClient struct {
@ -42,8 +44,9 @@ func (ws *WebSocketClient) Connect(
attempts int, attempts int,
interval time.Duration, interval time.Duration,
) (err error) { ) (err error) {
ws.Disconnect() if ws.conn != nil {
ws.ctx, ws.cancel = context.WithCancel(context.Background()) return errors.New("connection already open")
}
ws.URL, err = url.Parse(webSocketUrl) ws.URL, err = url.Parse(webSocketUrl)
if err != nil { if err != nil {
@ -53,7 +56,9 @@ func (ws *WebSocketClient) Connect(
// attempt retry // attempt retry
for range attempts { for range attempts {
origin := ws.URL.Scheme + "://" + ws.URL.Host origin := ws.URL.Scheme + "://" + ws.URL.Host
ws.conn, err = websocket.Dial(ws.URL.String(), "", origin) header := make(http.Header)
header.Add("Origin", origin)
ws.conn, _, err = websocket.Dial(ws.ctx, ws.URL.String(), &websocket.DialOptions{HTTPHeader: header})
if err == nil { if err == nil {
break break
} }
@ -77,7 +82,8 @@ func (ws *WebSocketClient) Connect(
return return
default: default:
var payload []byte var payload []byte
if err := websocket.Message.Receive(ws.conn, &payload); err != nil { _, payload, err := ws.conn.Read(ws.ctx)
if err != nil {
if ws.onClose != nil { if ws.onClose != nil {
ws.onClose(ws) ws.onClose(ws)
} }
@ -93,7 +99,7 @@ func (ws *WebSocketClient) Connect(
func (ws *WebSocketClient) Disconnect() { func (ws *WebSocketClient) Disconnect() {
if ws.conn != nil { if ws.conn != nil {
ws.conn.Close() ws.conn.CloseNow()
} }
ws.cancel() ws.cancel()
@ -105,9 +111,5 @@ func (ws *WebSocketClient) Disconnect() {
} }
func (ws *WebSocketClient) SendJSON(v any) error { func (ws *WebSocketClient) SendJSON(v any) error {
bytes, err := json.Marshal(v) return wsjson.Write(ws.ctx, ws.conn, v)
if err != nil {
return fmt.Errorf("failed to marshal json: %v", err)
}
return websocket.Message.Send(ws.conn, string(bytes))
} }