change ws package

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

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=

main.go

@@ -1,14 +1,21 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"log"
 	"net"
 	"regexp"
 
 	"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()
 	}
 }