From ebbc7bd47173423b465dd0ff282772018c6c60ac Mon Sep 17 00:00:00 2001
From: Delyan Angelov <delian66@gmail.com>
Date: Wed, 22 Jul 2020 18:42:57 +0300
Subject: [PATCH] examples: fix crash when running examples/ws/client.v outside
 of valgrind

---
 examples/ws/client.v           |  1 +
 vlib/net/websocket/handshake.v |  4 ++--
 vlib/net/websocket/io.v        | 15 ++++++++++-----
 vlib/net/websocket/ws.v        |  5 +----
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/examples/ws/client.v b/examples/ws/client.v
index 83dd2b4b70..044c3875a8 100644
--- a/examples/ws/client.v
+++ b/examples/ws/client.v
@@ -10,6 +10,7 @@ fn main() {
 }
 
 fn ws_test(uri string) {
+	println('connecting to $uri ...')
 	mut ws := websocket.new(uri)
 	ws.subscriber.subscribe('on_open', on_open)
 	ws.subscriber.subscribe('on_message', on_message)
diff --git a/vlib/net/websocket/handshake.v b/vlib/net/websocket/handshake.v
index 97589a97c9..587cfcb71b 100644
--- a/vlib/net/websocket/handshake.v
+++ b/vlib/net/websocket/handshake.v
@@ -6,7 +6,7 @@ fn (mut ws Client) read_handshake(seckey string) {
 	max_buffer := 1024
 	buffer_size := 1
 	mut buffer := malloc(max_buffer)
-	for bytes_read <= max_buffer {
+	for bytes_read < max_buffer - 1 {
 		mut res := 0
 		unsafe {
 			res = ws.read_from_server(buffer + bytes_read, buffer_size)
@@ -21,7 +21,7 @@ fn (mut ws Client) read_handshake(seckey string) {
 		}
 		bytes_read += buffer_size
 	}
-	buffer[max_buffer + 1] = `\0`
+	buffer[max_buffer - 1] = `\0`
 	ws.handshake_handler(string(byteptr(buffer)), seckey)
 }
 
diff --git a/vlib/net/websocket/io.v b/vlib/net/websocket/io.v
index dfe31059a5..e8bb7e2d32 100644
--- a/vlib/net/websocket/io.v
+++ b/vlib/net/websocket/io.v
@@ -3,16 +3,21 @@ module websocket
 fn (mut ws Client) write_to_server(buf voidptr, len int) int {
 	mut bytes_written := 0
 	ws.write_lock.m_lock()
-	bytes_written = if ws.is_ssl { C.SSL_write(ws.ssl, buf, len) } else { C.write(ws.socket.sockfd,
-			buf, len) }
+	if ws.is_ssl { 
+		bytes_written = C.SSL_write(ws.ssl, buf, len)
+	} else {
+		bytes_written = C.write(ws.socket.sockfd, buf, len)
+	}        
 	ws.write_lock.unlock()
 	return bytes_written
 }
 
 fn (ws &Client) read_from_server(buffer byteptr, buffer_size int) int {
-	return if ws.is_ssl {
-		C.SSL_read(ws.ssl, buffer, buffer_size)
+	mut res := 0
+	if ws.is_ssl {
+		res = C.SSL_read(ws.ssl, buffer, buffer_size)
 	} else {
-		C.read(ws.socket.sockfd, buffer, buffer_size)
+		res = C.read(ws.socket.sockfd, buffer, buffer_size)
 	}
+	return res
 }
diff --git a/vlib/net/websocket/ws.v b/vlib/net/websocket/ws.v
index 1693bbd30b..738456175d 100644
--- a/vlib/net/websocket/ws.v
+++ b/vlib/net/websocket/ws.v
@@ -239,13 +239,11 @@ pub fn (mut ws Client) close(code int, message string) {
 	}
 }
 
+// write will send the payload to the websocket server. NB: *it will not free the payload*, the caller is responsible for that.
 pub fn (mut ws Client) write(payload byteptr, payload_len int, code OPCode) int {
 	mut bytes_written := -1
 	if ws.state != .open {
 		ws.send_error_event('WebSocket closed. Cannot write.')
-		unsafe {
-			free(payload)
-		}
 		return -1
 	}
 	header_len := 6 + if payload_len > 125 { 2 } else { 0 } + if payload_len > 0xffff { 6 } else { 0 }
@@ -305,7 +303,6 @@ pub fn (mut ws Client) write(payload byteptr, payload_len int, code OPCode) int
 	ws.log.debug('write: $bytes_written bytes written.')
 	free_data:
 	unsafe {
-		free(payload)
 		frame_buf.free()
 		header.free()
 		masking_key.free()