From 2d3508c8720dfa7cccd650cd50a75072c8ff87c5 Mon Sep 17 00:00:00 2001
From: starryskye <skye@blueoxygen.net>
Date: Mon, 13 Feb 2023 05:28:41 -0700
Subject: [PATCH] net.websocket: call close event (#17300)

---
 vlib/net/websocket/websocket_client.v |  1 +
 vlib/net/websocket/websocket_test.v   | 54 +++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/vlib/net/websocket/websocket_client.v b/vlib/net/websocket/websocket_client.v
index 28ded9c204..d2167b317b 100644
--- a/vlib/net/websocket/websocket_client.v
+++ b/vlib/net/websocket/websocket_client.v
@@ -337,6 +337,7 @@ pub fn (mut ws Client) close(code int, message string) ! {
 	defer {
 		ws.shutdown_socket() or {}
 		ws.reset_state() or {}
+		ws.send_close_event(code, message)
 	}
 	ws.set_state(.closing)
 	// mut code32 := 0
diff --git a/vlib/net/websocket/websocket_test.v b/vlib/net/websocket/websocket_test.v
index b0392c9e83..b9161bc0db 100644
--- a/vlib/net/websocket/websocket_test.v
+++ b/vlib/net/websocket/websocket_test.v
@@ -8,6 +8,7 @@ struct WebsocketTestResults {
 pub mut:
 	nr_messages      int
 	nr_pong_received int
+	nr_closes        int
 }
 
 // Do not run these tests everytime, since they are flaky.
@@ -126,3 +127,56 @@ fn ws_test(family net.AddrFamily, uri string) ! {
 	assert test_results.nr_pong_received >= 2
 	assert test_results.nr_messages == 2
 }
+
+fn test_on_close_when_server_closing_connection() ! {
+	port := 30000 + rand.intn(1024) or { 0 }
+	mut ws := websocket.new_server(.ip, port, '')
+	ws.on_message(fn (mut cli websocket.Client, msg &websocket.Message) ! {
+		if msg.opcode == .text_frame {
+			cli.close(1000, 'closing connection')!
+		}
+	})
+
+	mut test_results := WebsocketTestResults{}
+	ws.on_close_ref(fn (mut cli websocket.Client, code int, reason string, mut res WebsocketTestResults) ! {
+		res.nr_closes++
+	}, test_results)
+	spawn ws.listen()
+
+	mut client := websocket.new_client('ws://localhost:${port}')!
+
+	client.connect()!
+	spawn client.listen()
+
+	time.sleep(1000 * time.millisecond)
+
+	client.write_string('a message')!
+
+	time.sleep(1000 * time.millisecond)
+
+	assert test_results.nr_closes == 1
+}
+
+fn test_on_close_when_client_closing_connection() ! {
+	port := 30000 + rand.intn(1024) or { 0 }
+	mut ws := websocket.new_server(.ip, port, '')
+	spawn ws.listen()
+
+	mut client := websocket.new_client('ws://localhost:${port}')!
+
+	mut test_results := WebsocketTestResults{}
+	client.on_close_ref(fn (mut cli websocket.Client, code int, reason string, mut res WebsocketTestResults) ! {
+		res.nr_closes++
+	}, test_results)
+
+	client.connect()!
+	spawn client.listen()
+
+	time.sleep(1000 * time.millisecond)
+
+	client.close(1000, 'closing connection')!
+
+	time.sleep(1000 * time.millisecond)
+
+	assert test_results.nr_closes == 1
+}