From 75197d34e3f51804a03a3a58236290ab9f87d6c6 Mon Sep 17 00:00:00 2001
From: Philippe PITTOLI <p.pittoli@unistra.fr>
Date: Mon, 5 Aug 2019 01:52:54 +0200
Subject: [PATCH] websockets: handle errors

---
 src/websocketc.cr | 103 +++++++++++++++++++++-------------------------
 src/websocketd.cr |  16 ++++++-
 2 files changed, 62 insertions(+), 57 deletions(-)

diff --git a/src/websocketc.cr b/src/websocketc.cr
index 583a840..5b3aea5 100644
--- a/src/websocketc.cr
+++ b/src/websocketc.cr
@@ -19,62 +19,55 @@ OptionParser.parse! do |parser|
 	end
 end
 
-ws = WebSocket.new(URI.parse(uri))
-
-# HTTP::WebSocket.new(URI.parse("wss://websocket.example.com/chat"))       # Creates a new WebSocket with TLS to `websocket.example.com`
-# HTTP::WebSocket.new(URI.parse("http://websocket.example.com:8080/chat")) # Creates a new WebSocket to `websocket.example.com` on port `8080`
-# HTTP::WebSocket.new(URI.parse("ws://websocket.example.com/chat"),        # Creates a new WebSocket to `websocket.example.com` with an Authorization header
-#   HTTP::Headers{"Authorization" => "Bearer authtoken"})
-
-m = ws.read
-if m.nil?
-	puts "oh no"
-	exit
-end
-# puts "message: #{String.new(m)}"
-
-# puts "sending pong"
-ws.send "pong".to_slice
-m = ws.read
-if m.nil?
-	puts "oh no"
-	exit
-end
-# puts "message: #{String.new(m)}"
-
-# puts "sending coucou"
-tosend = to_message 2, "coucou"
-# print_hexa(String.new(tosend), "#{CBLUE}Sending message hexa#{CRESET}")
-ws.send tosend
-
-m = ws.read
-if m.nil?
-	puts "oh no"
-	exit
-end
-puts "message: #{String.new(m)}"
-
-# ws.on_message do |message|
-# 	if message == "are we websocket yet?"
-# 		puts "sending pong"
-# 		ws.send "pong\n"
-# 	elsif message.chomp == "OK"
-# 		puts "sending coucou"
-# 		m = to_message 2, "coucou"
-# 		print_hexa(String.new(m), "#{CBLUE}Sending message hexa#{CRESET}")
-# 		ws.send m
-# 	else
-# 		print_hexa message, "#{CORANGE}Receving a message#{CRESET}"
-# 	end
-# end
-
-ws.on_close do |socket|
-	puts "socket is closing"
-	# socket.close
+def read_then_print(ws : WebSocket)
+	m = read ws
+	puts "message: #{String.new(m)}"
 end
 
-# ws.run
+def read_then_print_hexa(ws : WebSocket)
+	m = read ws
+	print_hexa(String.new(m), "#{CBLUE}Received message hexa#{CRESET}")
+end
 
-ws.close
+def read(ws : WebSocket)
+	m = ws.read
+	if m.nil?
+		raise "empty message"
+	end
 
-ws.read
+	m
+end
+
+def send_with_announce(ws : WebSocket, m : String)
+	puts "sending #{m}"
+	send ws, m
+end
+
+def send(ws : WebSocket, m : String)
+	ws.send m
+end
+
+def send(ws : WebSocket, m : Slice)
+	ws.send m
+end
+
+
+begin
+	ws = WebSocket.new(URI.parse(uri))
+
+	ws.on_close do |socket|
+		puts "socket is closing"
+		exit 0
+	end
+
+	read ws
+	send ws, "pong"
+	read ws
+	send ws, to_message(2, "coucou")
+	read ws
+
+	ws.close
+	ws.read
+rescue e
+	puts "Exception: #{e}"
+end
diff --git a/src/websocketd.cr b/src/websocketd.cr
index ef85584..1c8e78f 100644
--- a/src/websocketd.cr
+++ b/src/websocketd.cr
@@ -208,7 +208,13 @@ def websocket_switching_procedure (activefd : Int, context : InstanceStorage)
 			# The client is a WebSocket on top of a TCP connection
 			client = context.socklist[activefd]
 			wsclient = context.wssocklist[activefd]
-			message = wsclient.read
+			begin
+				message = wsclient.read
+			rescue e
+				puts "#{CRED}Exception (receiving a message)#{CRESET} #{e}"
+				closing_client activefd, context
+				return
+			end
 
 			if wsclient.closed?
 				# puts "#{CBLUE}client is closed#{CRESET}"
@@ -290,7 +296,13 @@ service.loop do |event|
 			# since it's an external communication
 			# we have to read the message here, it's not handled by default in libipc
 			wsclient = context.wssocklist[activefd]
-			message = wsclient.read
+			begin
+				message = wsclient.read
+			rescue e
+				puts "#{CRED}Exception (receiving a message)#{CRESET} #{e}"
+				closing_client activefd, context
+				next
+			end
 
 			if message.nil?
 				puts "#{CBLUE}disconnection of client#{CRESET} #{event.connection.fd}"