Merge pull request #4 from Issif/memory_leak

Issif · Jun 10, 2020 · 79ad191 · 79ad191
2 parents c6c43b8 + 6c617f0
commit 79ad191
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/lib/lib.go b/lib/lib.go
@@ -139,12 +139,18 @@ func IsMatchingCert(config *Configuration, result *Result, reg *regexp.Regexp) b
 
 // LoopCertStream gathers messages from CertStream
 func LoopCertStream(config *Configuration) {
+	dial := ws.Dialer{
+		ReadBufferSize:  8192,
+		WriteBufferSize: 512,
+		Timeout:         1 * time.Second,
+	}
 	for {
-		conn, _, _, err := ws.DefaultDialer.Dial(context.Background(), certInput)
-		defer conn.Close()
+		// conn, _, _, err := ws.DefaultDialer.Dial(context.Background(), certInput)
+		conn, _, _, err := dial.Dial(context.Background(), certInput)
 		if err != nil {
 			log.Warn("Error connecting to CertStream! Sleeping a few seconds and reconnecting...")
 			time.Sleep(1 * time.Second)
+			conn.Close()
 			continue
 		}
 		for {
@@ -155,6 +161,7 @@ func LoopCertStream(config *Configuration) {
 			}
 			config.Messages <- msg
 		}
+		conn.Close()
 	}
 }