websocket: enable using an already existing connection (from vweb or …

…another http server) (#20103)
vlang · Dec 19, 2023 · 5be5cd9 · 5be5cd9
1 parent db6ae6e
commit 5be5cd9
Show file tree

Hide file tree

Showing 5 changed files with 156 additions and 7 deletions.
diff --git a/examples/vweb/vweb_websocket/assets/websocket_client.js b/examples/vweb/vweb_websocket/assets/websocket_client.js
@@ -0,0 +1,22 @@
+const messageList = document.getElementById('message-list');
+const protocol = location.protocol === 'https:' ? 'wss' : 'ws';
+const socket = new WebSocket(`${protocol}://${location.host}/ws`);
+let i = 0;
+
+function send(message) {
+  messageList.innerHTML += `<li>&gt; ${message}</li>`;
+  socket.send(message);
+}
+
+socket.addEventListener("open", (event) => {
+  console.log('Connected to WS server');
+  send('Hey everyone !');
+});
+
+socket.addEventListener("message", (event) => {
+  const { data } = event;
+  messageList.innerHTML += `<li>&lt; ${data}</li>`;
+  setTimeout(() => {
+    send(`Roger ${i++}`);
+  }, 3000);
+});
diff --git a/examples/vweb/vweb_websocket/index.html b/examples/vweb/vweb_websocket/index.html
@@ -0,0 +1,11 @@
+<!doctype html>
+<html lang=en>
+<head>
+<meta charset=utf-8>
+<title>vweb websocket example page</title>
+</head>
+<body>
+  <ol id="message-list"></ol>
+  <script type="text/javascript" src="websocket_client.js"></script>
+</body>
+</html>
diff --git a/examples/vweb/vweb_websocket/vweb_websocket.v b/examples/vweb/vweb_websocket/vweb_websocket.v
@@ -0,0 +1,77 @@
+module main
+
+import log
+import net.http
+import net.websocket
+import term
+import vweb
+
+const http_port = 8080
+
+struct App {
+	vweb.Context
+mut:
+	wss &websocket.Server @[vweb_global]
+}
+
+fn slog(message string) {
+	eprintln(term.colorize(term.bright_yellow, message))
+}
+
+fn clog(message string) {
+	eprintln(term.colorize(term.cyan, message))
+}
+
+fn wlog(message string) {
+	eprintln(term.colorize(term.bright_blue, message))
+}
+
+fn main() {
+	mut app := new_app() or { panic(err) }
+	vweb.run(app, http_port)
+}
+
+fn new_app() !&App {
+	mut app := &App{
+		wss: new_websocker_server()!
+	}
+	app.handle_static('assets', true)
+	return app
+}
+
+fn new_websocker_server() !&websocket.Server {
+	mut wss := &websocket.Server{
+		logger: &log.Log{
+			level: .debug
+		}
+	}
+	wss.on_connect(fn (mut server_client websocket.ServerClient) !bool {
+		slog('ws.on_connect, server_client.client_key: ${server_client.client_key}')
+		return true
+	})!
+	wss.on_message(fn (mut ws websocket.Client, msg &websocket.Message) ! {
+		slog('s.on_message msg.opcode: ${msg.opcode} | msg.payload: ${msg.payload}')
+		ws.write(msg.payload, msg.opcode) or {
+			eprintln('ws.write err: ${err}')
+			return err
+		}
+	})
+	wss.on_close(fn (mut ws websocket.Client, code int, reason string) ! {
+		slog('s.on_close code: ${code}, reason: ${reason}')
+	})
+	slog('Websocket Server initialized')
+	return wss
+}
+
+pub fn (mut app App) index() vweb.Result {
+	return $vweb.html()
+}
+
+pub fn (mut app App) ws() !vweb.Result {
+	key := app.req.header.get(http.CommonHeader.sec_websocket_key)!
+	app.wss.handle_handshake(mut app.conn, key) or {
+		wlog('handle_handshake error: ${err.msg()}')
+		return err
+	}
+	return app.text('')
+}
diff --git a/vlib/net/http/header.v b/vlib/net/http/header.v
@@ -103,6 +103,7 @@ pub enum CommonHeader {
 	sec_fetch_site
 	sec_fetch_user
 	sec_websocket_accept
+	sec_websocket_key
 	server
 	server_timing
 	set_cookie
@@ -209,6 +210,7 @@ pub fn (h CommonHeader) str() string {
 		.sec_fetch_site { 'Sec-Fetch-Site' }
 		.sec_fetch_user { 'Sec-Fetch-User' }
 		.sec_websocket_accept { 'Sec-WebSocket-Accept' }
+		.sec_websocket_key { 'Sec-WebSocket-Key' }
 		.server { 'Server' }
 		.server_timing { 'Server-Timing' }
 		.set_cookie { 'Set-Cookie' }
@@ -314,6 +316,7 @@ const common_header_map = {
 	'sec-fetch-site':                      .sec_fetch_site
 	'sec-fetch-user':                      .sec_fetch_user
 	'sec-websocket-accept':                .sec_websocket_accept
+	'sec_websocket_key':                   .sec_websocket_key
 	'server':                              .server
 	'server-timing':                       .server_timing
 	'set-cookie':                          .set_cookie

diff --git a/vlib/net/websocket/websocket_server.v b/vlib/net/websocket/websocket_server.v
@@ -134,22 +134,58 @@ fn (mut s Server) serve_client(mut c Client) ! {
 		c.logger.debug('server-> End serve client (${c.id})')
 	}
 	mut handshake_response, mut server_client := s.handle_server_handshake(mut c)!
+	s.attach_client(mut server_client, handshake_response)!
+	c.listen() or {
+		s.logger.error(err.msg())
+		return err
+	}
+}
+
+// handle_handshake use an existing connection to respond to the handshake for a given key
+pub fn (mut s Server) handle_handshake(mut conn net.TcpConn, key string) !&ServerClient {
+	mut c := &Client{
+		is_server: true
+		conn: conn
+		is_ssl: false
+		logger: &log.Log{
+			level: .debug
+		}
+		client_state: ClientState{
+			state: .open
+		}
+		last_pong_ut: time.now().unix
+		id: rand.uuid_v4()
+	}
+	mut server_client := &ServerClient{
+		resource_name: 'GET'
+		client_key: key
+		client: unsafe { c }
+		server: unsafe { &s }
+	}
+	digest := create_key_challenge_response(key)!
+	handshake_response := 'HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ${digest}\r\n\r\n'
+	s.attach_client(mut server_client, handshake_response)!
+	spawn s.handle_ping()
+	c.listen() or {
+		s.logger.error(err.msg())
+		return err
+	}
+	return server_client
+}
+
+fn (mut s Server) attach_client(mut server_client ServerClient, handshake_response string) ! {
 	accept := s.send_connect_event(mut server_client)!
 	if !accept {
 		s.logger.debug('server-> client not accepted')
-		c.shutdown_socket()!
+		server_client.client.shutdown_socket()!
 		return
 	}
 	// the client is accepted
-	c.socket_write(handshake_response.bytes())!
+	server_client.client.socket_write(handshake_response.bytes())!
 	lock s.server_state {
-		s.server_state.clients[server_client.client.id] = server_client
+		s.server_state.clients[server_client.client.id] = unsafe { server_client }
 	}
 	s.setup_callbacks(mut server_client)
-	c.listen() or {
-		s.logger.error(err.msg())
-		return err
-	}
 }
 
 // setup_callbacks initialize all callback functions