From 5736db8678036adfc9ca0468507d274b2fd1563d Mon Sep 17 00:00:00 2001
From: Issif <thomas.labarussias@qonto.eu>
Date: Thu, 16 Apr 2020 18:22:41 +0200
Subject: [PATCH] change ws package

---
 go.mod  |  3 +++
 go.sum  |  6 ++++++
 main.go | 18 ++++++++++++++----
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/go.mod b/go.mod
index 9675a1e..0a3b691 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,9 @@ go 1.14
 require (
 	github.com/CaliDog/certstream-go v0.0.0-20180219203951-6016c5462366
 	github.com/falcosecurity/falcosidekick v0.0.0-20200128210241-efc9dea5ac3a
+	github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee // indirect
+	github.com/gobwas/pool v0.2.0 // indirect
+	github.com/gobwas/ws v1.0.3
 	github.com/google/pprof v0.0.0-20200413000643-b1a96885c1c6 // indirect
 	github.com/gorilla/websocket v1.4.2
 	github.com/jmoiron/jsonq v0.0.0-20150511023944-e874b168d07e
diff --git a/go.sum b/go.sum
index fb04270..29813a1 100644
--- a/go.sum
+++ b/go.sum
@@ -40,6 +40,12 @@ github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0=
+github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
+github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8=
+github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
+github.com/gobwas/ws v1.0.3 h1:ZOigqf7iBxkA4jdQ3am7ATzdlOFp9YzA6NmuvEEZc9g=
+github.com/gobwas/ws v1.0.3/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
diff --git a/main.go b/main.go
index eb0d3a0..d8a441b 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"log"
 	"net"
@@ -8,7 +9,13 @@ import (
 
 	"time"
 
-	"github.com/gorilla/websocket"
+	"net/http"
+	_ "net/http/pprof"
+
+	_ "expvar"
+
+	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
 )
 
 type result struct {
@@ -50,24 +57,26 @@ func init() {
 }
 
 func main() {
+	go http.ListenAndServe("localhost:6060", nil)
+
 	msgChan := make(chan []byte, 10)
 	for i := 0; i < config.Workers; i++ {
 		go certCheckWorker(msgChan)
 	}
 
 	for {
-		ws, _, err := websocket.DefaultDialer.Dial("wss://certstream.calidog.io", nil)
-		defer ws.Close()
+		conn, _, _, err := ws.DefaultDialer.Dial(context.Background(), "wss://certstream.calidog.io")
 
 		if err != nil {
 			if config.DisplayErrors == "true" {
 				log.Println("[ERROR] : Error connecting to certstream! Sleeping a few seconds and reconnecting...")
 			}
+			conn.Close()
 			time.Sleep(1 * time.Second)
 			continue
 		}
 		for {
-			_, msg, err := ws.ReadMessage()
+			msg, _, err := wsutil.ReadServerData(conn)
 			if err != nil {
 				if config.DisplayErrors == "true" {
 					log.Println("[ERROR] : Error reading message from CertStream")
@@ -76,6 +85,7 @@ func main() {
 			}
 			msgChan <- msg
 		}
+		conn.Close()
 	}
 }