client fix

http.go

@@ -27,6 +27,7 @@ import (
 	"mime"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -267,6 +268,23 @@ func (t *httpServerConn) RemoteAddr() string {
 // SetWriteDeadline does nothing and always returns nil.
 func (t *httpServerConn) SetWriteDeadline(time.Time) error { return nil }
 
+type WebsocketServer struct {
+	s *Server
+}
+
+func (s *WebsocketServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if isWebsocket(r) {
+		s.s.WebsocketHandler([]string{"*"}).ServeHTTP(w, r)
+		return
+	}
+	s.s.ServeHTTP(w, r)
+}
+
+func isWebsocket(r *http.Request) bool {
+	return strings.EqualFold(r.Header.Get("Upgrade"), "websocket") &&
+		strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade")
+}
+
 // ServeHTTP serves JSON-RPC requests over HTTP.
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// Permit dumb empty requests for remote health-checks (AWS)

websocket.go

@@ -160,10 +160,10 @@ func heartbeat(ctx context.Context, c *websocket.Conn, d time.Duration) {
 func newWebsocketCodec(ctx context.Context, c *websocket.Conn, host string, req http.Header) ServerCodec {
 	c.SetReadLimit(wsMessageSizeLimit)
 	jsonWriter := func(v any) error {
-		return wsjson.Write(ctx, c, v)
+		return wsjson.Write(context.Background(), c, v)
 	}
 	jsonReader := func(v any) error {
-		return wsjson.Read(ctx, c, v)
+		return wsjson.Read(context.Background(), c, v)
 	}
 	conn := websocket.NetConn(ctx, c, websocket.MessageText)
 	wc := &websocketCodec{