Add timeout for initial handshake.

c86b1c87 · Juliusz Chroboczek · b201c3d9 · c86b1c87
Commit c86b1c87 authored Apr 25, 2020 by Juliusz Chroboczek
Show whitespace changes
Inline Side-by-side

Showing with 12 additions and 0 deletions

client.go client.go +12 -0

No files found.
--- a/client.go
+++ b/client.go
@@ -108,11 +108,23 @@ type closeMessage struct {

 func startClient(conn *websocket.Conn) (err error) {
 	var m clientMessage
+
+	err = conn.SetReadDeadline(time.Now().Add(15 * time.Second))
+	if err != nil {
+		conn.Close()
+		return
+	}
 	err = conn.ReadJSON(&m)
 	if err != nil {
 		conn.Close()
 		return
 	}
+	err = conn.SetReadDeadline(time.Time{})
+	if err != nil {
+		conn.Close()
+		return
+	}
+
 	if m.Type != "handshake" {
 		conn.Close()
 		return