From a50e0f0522e29b796200ff52db7c6ef9e04a41c8 Mon Sep 17 00:00:00 2001
From: Ken <ken0x0a+github@gmail.com>
Date: Tue, 5 Jul 2022 12:40:23 +0900
Subject: [PATCH] net.websocket: allow timeout to be configured (#14941)

---
 vlib/net/websocket/io.v               | 27 ++++++---------------------
 vlib/net/websocket/websocket_client.v | 21 ++++++++++++++-------
 2 files changed, 20 insertions(+), 28 deletions(-)

diff --git a/vlib/net/websocket/io.v b/vlib/net/websocket/io.v
index 2d0db657b2..4606304a6a 100644
--- a/vlib/net/websocket/io.v
+++ b/vlib/net/websocket/io.v
@@ -1,7 +1,6 @@
 module websocket
 
 import net
-import time
 
 // socket_read reads from socket into the provided buffer
 fn (mut ws Client) socket_read(mut buffer []u8) ?int {
@@ -13,15 +12,8 @@ fn (mut ws Client) socket_read(mut buffer []u8) ?int {
 			r := ws.ssl_conn.read(mut buffer)?
 			return r
 		} else {
-			for {
-				r := ws.conn.read(mut buffer) or {
-					if err.code() == net.err_timed_out_code {
-						continue
-					}
-					return err
-				}
-				return r
-			}
+			r := ws.conn.read(mut buffer)?
+			return r
 		}
 	}
 	return none
@@ -37,15 +29,8 @@ fn (mut ws Client) socket_read_ptr(buf_ptr &u8, len int) ?int {
 			r := ws.ssl_conn.socket_read_into_ptr(buf_ptr, len)?
 			return r
 		} else {
-			for {
-				r := ws.conn.read_ptr(buf_ptr, len) or {
-					if err.code() == net.err_timed_out_code {
-						continue
-					}
-					return err
-				}
-				return r
-			}
+			r := ws.conn.read_ptr(buf_ptr, len)?
+			return r
 		}
 	}
 	return none
@@ -91,8 +76,8 @@ fn (mut ws Client) dial_socket() ?&net.TcpConn {
 	mut t := net.dial_tcp(tcp_address)?
 	optval := int(1)
 	t.sock.set_option_int(.keep_alive, optval)?
-	t.set_read_timeout(30 * time.second)
-	t.set_write_timeout(30 * time.second)
+	t.set_read_timeout(ws.read_timeout)
+	t.set_write_timeout(ws.write_timeout)
 	if ws.is_ssl {
 		ws.ssl_conn.connect(mut t, ws.uri.hostname)?
 	}
diff --git a/vlib/net/websocket/websocket_client.v b/vlib/net/websocket/websocket_client.v
index 95aa4a2ce5..35ebe95fbd 100644
--- a/vlib/net/websocket/websocket_client.v
+++ b/vlib/net/websocket/websocket_client.v
@@ -27,9 +27,11 @@ mut:
 	open_callbacks    []OpenEventHandler    // all callbacks on_open
 	close_callbacks   []CloseEventHandler   // all callbacks on_close
 pub:
-	is_ssl bool   // true if secure socket is used
-	uri    Uri    // uri of current connection
-	id     string // unique id of client
+	is_ssl        bool   // true if secure socket is used
+	uri           Uri    // uri of current connection
+	id            string // unique id of client
+	read_timeout  i64
+	write_timeout i64
 pub mut:
 	header            http.Header  // headers that will be passed when connecting
 	conn              &net.TcpConn // underlying TCP socket connection
@@ -73,8 +75,14 @@ pub enum OPCode {
 	pong = 0x0A
 }
 
+[params]
+pub struct ClientOpt {
+	read_timeout  i64 = 30 * time.second
+	write_timeout i64 = 30 * time.second
+}
+
 // new_client instance a new websocket client
-pub fn new_client(address string) ?&Client {
+pub fn new_client(address string, opt ClientOpt) ?&Client {
 	uri := parse_uri(address)?
 	return &Client{
 		conn: 0
@@ -88,6 +96,8 @@ pub fn new_client(address string) ?&Client {
 		state: .closed
 		id: rand.uuid_v4()
 		header: http.new_header()
+		read_timeout: opt.read_timeout
+		write_timeout: opt.write_timeout
 	}
 }
 
@@ -97,9 +107,6 @@ pub fn (mut ws Client) connect() ? {
 	ws.set_state(.connecting)
 	ws.logger.info('connecting to host $ws.uri')
 	ws.conn = ws.dial_socket()?
-	// Todo: make setting configurable
-	ws.conn.set_read_timeout(time.second * 30)
-	ws.conn.set_write_timeout(time.second * 30)
 	ws.handshake()?
 	ws.set_state(.open)
 	ws.logger.info('successfully connected to host $ws.uri')