websocket: pass the instance of client to handler

Since handler are created before creating a client and a handler may
send a response to server after receiving a message from server, we pass
the current client to handler as parameter.

3 files changed, 48 insertions, 37 deletions

diff --git a/lib/websocket/client.go b/lib/websocket/client.go
index 0cc1e85a..d6dadcb6 100644
--- a/lib/websocket/client.go
+++ b/lib/websocket/client.go
@@ -57,7 +57,7 @@ var (
 //
 //	cl := &Client{
 //		Endpoint: "ws://127.0.0.1:9001",
-//		HandleText: func(frame *Frame) error {
+//		HandleText: func(cl *Client, frame *Frame) error {
 //			// Process response from request or broadcast from
 //			// server.
 //			return nil
@@ -175,7 +175,7 @@ func (cl *Client) Connect() (err error) {
 //
 // dummyHandle define dummy handle for HandleText and HandleBin.
 //
-func (cl *Client) dummyHandle(frame *Frame) error {
+func dummyHandle(cl *Client, frame *Frame) error {
 	return nil
 }
 
@@ -185,16 +185,16 @@ func (cl *Client) dummyHandle(frame *Frame) error {
 //
 func (cl *Client) init() (err error) {
 	if cl.HandleBin == nil {
-		cl.HandleBin = cl.dummyHandle
+		cl.HandleBin = dummyHandle
 	}
 	if cl.handleClose == nil {
-		cl.handleClose = cl.onClose
+		cl.handleClose = clientOnClose
 	}
 	if cl.handlePing == nil {
-		cl.handlePing = cl.onPing
+		cl.handlePing = clientOnPing
 	}
 	if cl.HandleText == nil {
-		cl.HandleText = cl.dummyHandle
+		cl.HandleText = dummyHandle
 	}
 
 	err = cl.parseURI()
@@ -317,9 +317,9 @@ func (cl *Client) handleBadRequest() {
 }
 
 //
-// onClose request from server.
+// clientOnClose request from server.
 //
-func (cl *Client) onClose(frame *Frame) error {
+func clientOnClose(cl *Client, frame *Frame) error {
 	switch {
 	case frame.closeCode == 0:
 		frame.closeCode = StatusBadRequest
@@ -436,9 +436,9 @@ func (cl *Client) handleFragment(frame *Frame) (isInvalid bool) {
 			cl.handleInvalidData()
 			return true
 		}
-		err = cl.HandleText(frame)
+		err = cl.HandleText(cl, frame)
 	} else {
-		err = cl.HandleBin(frame)
+		err = cl.HandleBin(cl, frame)
 	}
 	if err != nil {
 		cl.handleBadRequest()
@@ -471,19 +471,19 @@ func (cl *Client) handleFrame(frame *Frame) (isClosing bool) {
 		cl.handleBadRequest()
 		return true
 	case OpcodeClose:
-		cl.handleClose(frame)
+		cl.handleClose(cl, frame)
 		return true
 	case OpcodePing:
-		_ = cl.handlePing(frame)
+		_ = cl.handlePing(cl, frame)
 	case OpcodePong:
 		if cl.handlePong != nil {
-			_ = cl.handlePong(frame)
+			_ = cl.handlePong(cl, frame)
 		}
 	case OpcodeControlRsvB, OpcodeControlRsvC, OpcodeControlRsvD, OpcodeControlRsvE, OpcodeControlRsvF:
 		if cl.HandleRsvControl != nil {
-			_ = cl.HandleRsvControl(frame)
+			_ = cl.HandleRsvControl(cl, frame)
 		} else {
-			cl.handleClose(frame)
+			cl.handleClose(cl, frame)
 			isClosing = true
 		}
 	}
@@ -492,15 +492,25 @@ func (cl *Client) handleFrame(frame *Frame) (isClosing bool) {
 }
 
 func (cl *Client) handleHandshake(ctx context.Context, resp []byte) (err error) {
+	if debug.Value >= 3 {
+		max := 512
+		if len(resp) < 512 {
+			max = len(resp)
+		}
+		fmt.Printf("websocket: Client.handleHandshake:\n%s\n--\n", resp[:max])
+	}
 	httpBuf := bufio.NewReader(bytes.NewBuffer(resp))
 
 	httpRes, err := http.ReadResponse(httpBuf, nil)
-	httpRes.Body.Close()
 	if err != nil {
+		fmt.Printf("websocket: Client.handleHandshake: http.ReadResponse")
 		return err
 	}
 
+	httpRes.Body.Close()
+
 	if httpRes.StatusCode != http.StatusSwitchingProtocols {
+		fmt.Printf("websocket: Client.handleHandshake: status code: %d\n", httpRes.StatusCode)
 		err = fmt.Errorf(httpRes.Status)
 		return err
 	}
@@ -649,9 +659,10 @@ func (cl *Client) Quit() {
 }
 
 //
-// onPing default handler when client receive control PING frame from server.
+// clientOnPing default handler when client receive control PING frame from
+// server.
 //
-func (cl *Client) onPing(frame *Frame) error {
+func clientOnPing(cl *Client, frame *Frame) error {
 	if frame == nil {
 		return nil
 	}
diff --git a/lib/websocket/client_test.go b/lib/websocket/client_test.go
index c6f9f7d9..af2cbdb5 100644
--- a/lib/websocket/client_test.go
+++ b/lib/websocket/client_test.go
@@ -133,7 +133,7 @@ func TestClientPing(t *testing.T) {
 			}
 		}
 
-		testClient.handleClose = func(got *Frame) error {
+		testClient.handleClose = func(cl *Client, got *Frame) error {
 			exp := c.expClose
 			test.Assert(t, "close", exp, got, true)
 
@@ -141,13 +141,13 @@ func TestClientPing(t *testing.T) {
 				got.payload = got.payload[2:]
 			}
 
-			testClient.SendClose(got.closeCode, got.payload)
-			testClient.Quit()
+			cl.SendClose(got.closeCode, got.payload)
+			cl.Quit()
 			wg.Done()
 			return nil
 		}
 
-		testClient.handlePong = func(got *Frame) (err error) {
+		testClient.handlePong = func(cl *Client, got *Frame) (err error) {
 			exp := c.exp
 			test.Assert(t, "handlePong", exp, got, true)
 			wg.Done()
@@ -268,16 +268,16 @@ func TestClientText(t *testing.T) {
 			}
 		}
 
-		testClient.handleClose = func(got *Frame) error {
+		testClient.handleClose = func(cl *Client, got *Frame) error {
 			exp := c.expClose
 			test.Assert(t, "close", exp, got, true)
-			testClient.SendClose(got.closeCode, got.payload)
-			testClient.Quit()
+			cl.SendClose(got.closeCode, got.payload)
+			cl.Quit()
 			wg.Done()
 			return nil
 		}
 
-		testClient.HandleText = func(got *Frame) error {
+		testClient.HandleText = func(cl *Client, got *Frame) error {
 			exp := c.exp
 			test.Assert(t, "text", exp, got, true)
 			wg.Done()
@@ -401,16 +401,16 @@ func TestClientFragmentation(t *testing.T) {
 			}
 		}
 
-		testClient.handleClose = func(got *Frame) error {
+		testClient.handleClose = func(cl *Client, got *Frame) error {
 			exp := c.expClose
 			test.Assert(t, "close", exp, got, true)
-			testClient.SendClose(got.closeCode, got.payload)
-			testClient.Quit()
+			cl.SendClose(got.closeCode, got.payload)
+			cl.Quit()
 			wg.Done()
 			return nil
 		}
 
-		testClient.HandleText = func(got *Frame) error {
+		testClient.HandleText = func(cl *Client, got *Frame) error {
 			exp := c.exp
 			test.Assert(t, "text", exp, got, true)
 			wg.Done()
@@ -472,7 +472,7 @@ func TestClientFragmentation2(t *testing.T) {
 		payload: []byte("Shulhan"),
 	}}
 
-	testClient.handlePong = func(got *Frame) error {
+	testClient.handlePong = func(cl *Client, got *Frame) error {
 		exp := &Frame{
 			fin:        frameIsFinished,
 			opcode:     OpcodePong,
@@ -485,7 +485,7 @@ func TestClientFragmentation2(t *testing.T) {
 		return nil
 	}
 
-	testClient.HandleText = func(got *Frame) error {
+	testClient.HandleText = func(cl *Client, got *Frame) error {
 		exp := &Frame{
 			fin:        frameIsFinished,
 			opcode:     OpcodeText,
@@ -557,7 +557,7 @@ func TestClientSendBin(t *testing.T) {
 			}
 		}
 
-		testClient.HandleBin = func(got *Frame) error {
+		testClient.HandleBin = func(cl *Client, got *Frame) error {
 			exp := c.exp
 			test.Assert(t, "HandleBin", exp, got, true)
 			wg.Done()
@@ -621,7 +621,7 @@ func TestClientSendPing(t *testing.T) {
 		c := c
 		t.Log(c.desc)
 
-		testClient.handlePong = func(got *Frame) error {
+		testClient.handlePong = func(cl *Client, got *Frame) error {
 			exp := c.exp
 			test.Assert(t, "handlePong", exp, got, true)
 			wg.Done()
@@ -657,7 +657,7 @@ func TestClientSendClose(t *testing.T) {
 		t.Fatal("TestClientSendClose: Connect: " + err.Error())
 	}
 
-	testClient.handleClose = func(got *Frame) error {
+	testClient.handleClose = func(cl *Client, got *Frame) error {
 		exp := &Frame{
 			fin:        frameIsFinished,
 			opcode:     OpcodeClose,
@@ -667,7 +667,7 @@ func TestClientSendClose(t *testing.T) {
 			isComplete: true,
 		}
 		test.Assert(t, "handleClose", exp, got, true)
-		testClient.Quit()
+		cl.Quit()
 		wg.Done()
 		return nil
 	}
diff --git a/lib/websocket/handler.go b/lib/websocket/handler.go
index 02030561..ec2356a1 100644
--- a/lib/websocket/handler.go
+++ b/lib/websocket/handler.go
@@ -18,7 +18,7 @@ type clientRawHandler func(ctx context.Context, resp []byte) (err error)
 // ClientHandler define a callback type for client to handle packet from
 // server (either broadcast or from response of request) in the form of frame.
 //
-type ClientHandler func(frame *Frame) (err error)
+type ClientHandler func(cl *Client, frame *Frame) (err error)
 
 // HandlerAuthFn define server callback type to handle authentication request.
 type HandlerAuthFn func(req *Handshake) (ctx context.Context, err error)