1 files changed, 32 insertions, 12 deletions

diff --git a/xP/xP.go b/xP/xP.go
index bc7ebc2..b858ace 100644
--- a/xP/xP.go
+++ b/xP/xP.go
@@ -16,7 +16,7 @@ import (
 	"os"
 	"time"
 
-	"golang.org/x/net/websocket"
+	"nhooyr.io/websocket"
 )
 
 var (
@@ -27,16 +27,21 @@ var (
 
 func clientToRelay(
 	ctx context.Context, ws *websocket.Conn, conn net.Conn) bool {
-	var j string
-	if err := websocket.Message.Receive(ws, &j); err != nil {
+	t, b, err := ws.Read(ctx)
+	if err != nil {
 		log.Println("Command receive failed: " + err.Error())
 		return false
 	}
+	if t != websocket.MessageText {
+		log.Println("Command receive failed: " +
+			"binary messages are not supported")
+		return false
+	}
 
-	log.Printf("?> %s\n", j)
+	log.Printf("?> %s\n", b)
 
 	var m RelayCommandMessage
-	if err := json.Unmarshal([]byte(j), &m); err != nil {
+	if err := json.Unmarshal(b, &m); err != nil {
 		log.Println("Command unmarshalling failed: " + err.Error())
 		return false
 	}
@@ -85,7 +90,7 @@ func relayToClient(
 		log.Println("Event marshalling failed: " + err.Error())
 		return false
 	}
-	if err := websocket.Message.Send(ws, string(j)); err != nil {
+	if err := ws.Write(ctx, websocket.MessageText, j); err != nil {
 		log.Println("Event send failed: " + err.Error())
 		return false
 	}
@@ -94,7 +99,7 @@ func relayToClient(
 	return true
 }
 
-func errorToClient(ws *websocket.Conn, err error) bool {
+func errorToClient(ctx context.Context, ws *websocket.Conn, err error) bool {
 	j, err := json.Marshal(&RelayEventMessage{
 		EventSeq: 0,
 		Data: RelayEventData{
@@ -109,22 +114,37 @@ func errorToClient(ws *websocket.Conn, err error) bool {
 		log.Println("Event marshalling failed: " + err.Error())
 		return false
 	}
-	if err := websocket.Message.Send(ws, string(j)); err != nil {
+	if err := ws.Write(ctx, websocket.MessageText, j); err != nil {
 		log.Println("Event send failed: " + err.Error())
 		return false
 	}
 	return true
 }
 
-func handleWebSocket(ws *websocket.Conn) {
+func handleWS(w http.ResponseWriter, r *http.Request) {
+	ws, err := websocket.Accept(w, r, &websocket.AcceptOptions{
+		InsecureSkipVerify: true,
+		CompressionMode:    websocket.CompressionContextTakeover,
+		// This is for the payload, and happens to trigger on all messages.
+		CompressionThreshold: 16,
+	})
+	if err != nil {
+		log.Println("Client rejected: " + err.Error())
+		return
+	}
+
+	defer ws.Close(websocket.StatusGoingAway, "Goodbye")
+
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+
 	conn, err := net.Dial("tcp", addressConnect)
 	if err != nil {
-		errorToClient(ws, err)
+		errorToClient(ctx, ws, err)
 		return
 	}
 
 	// We don't need to intervene, so it's just two separate pipes so far.
-	ctx, cancel := context.WithCancel(ws.Request().Context())
 	go func() {
 		for clientToRelay(ctx, ws, conn) {
 		}
@@ -184,7 +204,7 @@ func main() {
 		addressWS = os.Args[3]
 	}
 
-	http.Handle("/ws", websocket.Handler(handleWebSocket))
+	http.Handle("/ws", http.HandlerFunc(handleWS))
 	http.Handle("/", http.HandlerFunc(handleDefault))
 
 	s := &http.Server{