From 9b26b5d46a5c9092d3b2bbfb80f60622b4d04ce7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=84=9A=E8=80=85?= <guweigang@bullsoft.org>
Date: Sat, 25 Apr 2026 03:04:01 +0800
Subject: [PATCH] net.websocket: add server on_attached hook after handshake
 attach (#26955)

---
 vlib/net/websocket/events.v           | 38 +++++++++++++++++++++++++++
 vlib/net/websocket/websocket_server.v |  8 ++++++
 vlib/net/websocket/websocket_test.v   | 23 ++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/vlib/net/websocket/events.v b/vlib/net/websocket/events.v
index de1dc5f5f..e6c51676e 100644
--- a/vlib/net/websocket/events.v
+++ b/vlib/net/websocket/events.v
@@ -32,7 +32,16 @@ struct CloseEventHandler {
 	ref      voidptr // referenced object
 }
 
+struct AttachedEventHandler {
+	handler  ServerAttachedFn  = unsafe { nil }
+	handler2 ServerAttachedFn2 = unsafe { nil }
+	is_ref   bool
+	ref      voidptr
+}
+
 pub type AcceptClientFn = fn (mut c ServerClient) !bool
+pub type ServerAttachedFn = fn (mut c ServerClient) !
+pub type ServerAttachedFn2 = fn (mut c ServerClient, v voidptr) !
 
 pub type SocketMessageFn = fn (mut c Client, msg &Message) !
 
@@ -90,6 +99,25 @@ pub fn (mut s Server) on_close_ref(fun SocketCloseFn2, ref voidptr) {
 	}
 }
 
+// on_attached registers a callback after a client has been attached to the
+// server, the handshake response has been written, and callbacks are set up,
+// but before the blocking listen loop begins.
+pub fn (mut s Server) on_attached(fun ServerAttachedFn) {
+	s.attached_callbacks << AttachedEventHandler{
+		handler: fun
+	}
+}
+
+// on_attached_ref registers a callback after a client has been attached to the
+// server and provides a reference object.
+pub fn (mut s Server) on_attached_ref(fun ServerAttachedFn2, ref voidptr) {
+	s.attached_callbacks << AttachedEventHandler{
+		handler2: fun
+		ref:      ref
+		is_ref:   true
+	}
+}
+
 // on_message registers a callback on new messages
 pub fn (mut ws Client) on_message(fun SocketMessageFn) {
 	ws.message_callbacks << MessageEventHandler{
@@ -166,6 +194,16 @@ fn (mut s Server) send_connect_event(mut c ServerClient) !bool {
 	return res
 }
 
+fn (mut s Server) send_attached_event(mut c ServerClient) ! {
+	for ev_handler in s.attached_callbacks {
+		if !ev_handler.is_ref {
+			ev_handler.handler(mut c)!
+		} else {
+			ev_handler.handler2(mut c, ev_handler.ref)!
+		}
+	}
+}
+
 // send_message_event invokes the on_message callback
 fn (mut ws Client) send_message_event(msg &Message) {
 	ws.debug_log('sending on_message event')
diff --git a/vlib/net/websocket/websocket_server.v b/vlib/net/websocket/websocket_server.v
index 380cb6ff0..7b9b9f1fa 100644
--- a/vlib/net/websocket/websocket_server.v
+++ b/vlib/net/websocket/websocket_server.v
@@ -20,6 +20,7 @@ mut:
 	logger                  &log.Logger      = default_logger
 	ls                      &net.TcpListener = unsafe { nil } // listener used to get incoming connection to socket
 	accept_client_callbacks []AcceptClientFn      // accept client callback functions
+	attached_callbacks      []AttachedEventHandler
 	message_callbacks       []MessageEventHandler // new message callback functions
 	close_callbacks         []CloseEventHandler   // close message callback functions
 pub:
@@ -183,6 +184,13 @@ fn (mut s Server) attach_client(mut server_client ServerClient, handshake_respon
 		s.server_state.clients[server_client.client.id] = unsafe { server_client }
 	}
 	s.setup_callbacks(mut server_client)
+	s.send_attached_event(mut server_client) or {
+		lock s.server_state {
+			s.server_state.clients.delete(server_client.client.id)
+		}
+		server_client.client.shutdown_socket() or {}
+		return err
+	}
 }
 
 // setup_callbacks initialize all callback functions
diff --git a/vlib/net/websocket/websocket_test.v b/vlib/net/websocket/websocket_test.v
index 821164893..23c5afa21 100644
--- a/vlib/net/websocket/websocket_test.v
+++ b/vlib/net/websocket/websocket_test.v
@@ -190,3 +190,26 @@ fn test_on_close_when_client_closing_connection() ! {
 	time.sleep(1000 * time.millisecond)
 	assert test_results.nr_closes == 1
 }
+
+fn attached_fail_cb(mut sc websocket.ServerClient) ! {
+	return error('attached hook failure')
+}
+
+fn client_close_ignore_cb(mut cli websocket.Client, code int, reason string) ! {}
+
+fn test_on_attached_error_does_not_leak_client() ! {
+	mut ws := websocket.new_server(.ip, 30005, '')
+	ws.on_attached(attached_fail_cb)
+	start_server_in_thread_and_wait_till_it_is_ready_to_accept_connections(mut ws)
+
+	mut client := websocket.new_client('ws://localhost:30005')!
+	client.on_close(client_close_ignore_cb)
+	client.connect()!
+	spawn client.listen()
+	time.sleep(500 * time.millisecond)
+
+	client_count := rlock ws.server_state {
+		ws.server_state.clients.len
+	}
+	assert client_count == 0
+}
-- 
2.39.5