AI bug fixes

Author: lanrat

certgraph.go

@@ -165,7 +165,7 @@ func main() {
 
 	// create the output directory if it does not exist
 	if len(config.savePath) > 0 {
-		err := os.MkdirAll(config.savePath, 0777)
+		err := os.MkdirAll(config.savePath, 0755)
 		if err != nil {
 			fmt.Fprintln(os.Stderr, err)
 			return
@@ -215,7 +215,7 @@ func getDriverSingle(name string) (driver.Driver, error) {
 	case "censys":
 		d, err = censys.Driver(config.savePath, config.includeCTSubdomains, config.includeCTExpired)
 	default:
-		return nil, fmt.Errorf("unknown driver name: %s", config.driver)
+		return nil, fmt.Errorf("unknown driver name: %s", name)
 	}
 	return d, err
 }
@@ -272,8 +272,7 @@ func breathFirstSearch(roots []string) {
 	}()
 	// thread to start all other threads from DomainChan
 	go func() {
-		for {
-			domainNode := <-domainNodeInputChan
+		for domainNode := range domainNodeInputChan {
 
 			// depth check
 			if domainNode.Depth > config.maxDepth {
@@ -341,6 +340,7 @@ func breathFirstSearch(roots []string) {
 	}()
 
 	wg.Wait() // wait for querying to finish
+	close(domainNodeInputChan) // close input channel to signal goroutine to exit
 	close(domainNodeOutputChan)
 	<-done // wait for save to finish
 }

dns/ns.go

@@ -4,11 +4,12 @@ package dns
 import (
 	"context"
 	"net"
+	"sync"
 	"time"
 )
 
 var (
-	dnsCache    = make(map[string]bool)
+	dnsCache    = &sync.Map{}
 	dnsResolver = &net.Resolver{}
 )
 
@@ -76,13 +77,12 @@ func HasRecordsCache(domain string, timeout time.Duration) (bool, error) {
 	if err != nil {
 		return false, err
 	}
-	hasDNS, found := dnsCache[domain]
-	if found {
-		return hasDNS, nil
+	if cached, found := dnsCache.Load(domain); found {
+		return cached.(bool), nil
 	}
 	hasRecords, err := HasRecords(domain, timeout)
-	if err != nil {
-		dnsCache[domain] = hasRecords
+	if err == nil {
+		dnsCache.Store(domain, hasRecords)
 	}
 	return hasRecords, err
 }

driver/censys/censys.go

@@ -203,7 +203,11 @@ func (d *censys) QueryDomain(domain string) (driver.Result, error) {
 	}
 
 	for _, r := range resp.Results {
-		fp := fingerprint.FromHexHash(r.Fingerprint)
+		fp, err := fingerprint.FromHexHash(r.Fingerprint)
+		if err != nil {
+			log.Printf("censys: invalid fingerprint %s: %v", r.Fingerprint, err)
+			continue
+		}
 		results.fingerprints.Add(domain, fp)
 	}
 

driver/crtsh/crtsh.go

@@ -79,6 +79,11 @@ func Driver(maxQueryResults int, timeout time.Duration, savePath string, include
 		return nil, err
 	}
 
+	// Configure connection pool to prevent resource leaks
+	d.db.SetMaxOpenConns(10)
+	d.db.SetMaxIdleConns(2)
+	d.db.SetConnMaxLifetime(time.Hour)
+
 	err = d.setSQLTimeout(d.timeout.Seconds())
 
 	return d, err
@@ -88,6 +93,14 @@ func (d *crtsh) GetName() string {
 	return driverName
 }
 
+// Close closes the database connection
+func (d *crtsh) Close() error {
+	if d.db != nil {
+		return d.db.Close()
+	}
+	return nil
+}
+
 func (d *crtsh) setSQLTimeout(sec float64) error {
 	_, err := d.db.Exec(fmt.Sprintf("SET statement_timeout TO %f;", (1000 * sec)))
 	return err
@@ -161,6 +174,7 @@ func (d *crtsh) QueryDomain(domain string) (driver.Result, error) {
 	if err != nil {
 		return results, err
 	}
+	defer func() { _ = rows.Close() }()
 
 	for rows.Next() {
 		var hash []byte
@@ -202,6 +216,7 @@ func (d *crtsh) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error
 	if err != nil {
 		return certNode, err
 	}
+	defer func() { _ = rows.Close() }()
 
 	for rows.Next() {
 		var domain string
@@ -214,7 +229,7 @@ func (d *crtsh) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error
 
 	if d.save {
 		var rawCert []byte
-		queryStr = `SELECT certificate FORM certificate_and_identities WHERE digest(certificate, 'sha256') = $1;`
+		queryStr = `SELECT certificate FROM certificate_and_identities WHERE digest(certificate, 'sha256') = $1;`
 		row := d.db.QueryRow(queryStr, fp[:])
 		err = row.Scan(&rawCert)
 		if err != nil {

driver/http/http.go

@@ -142,6 +142,9 @@ func (c *httpCertDriver) dialTLS(network, addr string) (net.Conn, error) {
 	connState := conn.ConnectionState()
 
 	// only look at leaf certificate which is valid for domain, rest of cert chain is ignored
+	if len(connState.PeerCertificates) == 0 {
+		return conn, fmt.Errorf("no peer certificates found")
+	}
 	certResult := driver.NewCertResult(connState.PeerCertificates[0])
 	c.certs[certResult.Fingerprint] = certResult
 	host, _, err := net.SplitHostPort(addr)

driver/smtp/smtp.go

@@ -129,6 +129,9 @@ func (d *smtpDriver) QueryDomain(host string) (driver.Result, error) {
 	}
 
 	// only look at leaf certificate which is valid for domain, rest of cert chain is ignored
+	if len(certs) == 0 {
+		return results, fmt.Errorf("no certificates found")
+	}
 	certResult := driver.NewCertResult(certs[0])
 	results.certs[certResult.Fingerprint] = certResult
 	results.fingerprints.Add(host, certResult.Fingerprint)

fingerprint/fingerprint.go

@@ -35,21 +35,21 @@ func FromRawCertBytes(data []byte) Fingerprint {
 }
 
 // FromB64Hash returns a Fingerprint from a base64 encoded hash string
-func FromB64Hash(hash string) Fingerprint {
+func FromB64Hash(hash string) (Fingerprint, error) {
 	data, err := base64.StdEncoding.DecodeString(hash)
 	if err != nil {
-		panic(err)
+		return Fingerprint{}, err
 	}
-	return FromHashBytes(data)
+	return FromHashBytes(data), nil
 }
 
 // FromHexHash returns a Fingerprint from a hex encoded hash string
-func FromHexHash(hash string) Fingerprint {
+func FromHexHash(hash string) (Fingerprint, error) {
 	decoded, err := hex.DecodeString(hash)
 	if err != nil {
-		panic(err)
+		return Fingerprint{}, err
 	}
-	return FromHashBytes(decoded)
+	return FromHashBytes(decoded), nil
 }
 
 // B64Encode returns the b64 string of a Fingerprint

fingerprint/fp_test.go

@@ -61,7 +61,10 @@ func TestFromRawCertBytes(t *testing.T) {
 
 func TestFromB64Hash(t *testing.T) {
 
-	fp := fingerprint.FromB64Hash(fpHashB64)
+	fp, err := fingerprint.FromB64Hash(fpHashB64)
+	if err != nil {
+		t.Fatalf("FromB64Hash failed: %v", err)
+	}
 
 	uppercaseHash := strings.ToUpper(fpHashHex)
 	hashHex := fp.HexString()
@@ -79,7 +82,10 @@ func TestFromB64Hash(t *testing.T) {
 
 func TestFromHexHash(t *testing.T) {
 
-	fp := fingerprint.FromHexHash(fpHashHex)
+	fp, err := fingerprint.FromHexHash(fpHashHex)
+	if err != nil {
+		t.Fatalf("FromHexHash failed: %v", err)
+	}
 
 	hashB64 := fp.B64Encode()
 

graph/graph.go

@@ -4,6 +4,7 @@ package graph
 import (
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/lanrat/certgraph/fingerprint"
 )
@@ -12,7 +13,7 @@ import (
 type CertGraph struct {
 	domains    sync.Map
 	certs      sync.Map
-	numDomains int
+	numDomains int64
 	depth      uint
 }
 
@@ -31,7 +32,7 @@ func (graph *CertGraph) AddCert(certNode *CertNode) {
 
 // AddDomain add a DomainNode to the graph
 func (graph *CertGraph) AddDomain(domainNode *DomainNode) {
-	graph.numDomains++
+	atomic.AddInt64(&graph.numDomains, 1)
 	// save the new maximum depth if greather then current
 	if domainNode.Depth > graph.depth {
 		graph.depth = domainNode.Depth
@@ -44,7 +45,7 @@ func (graph *CertGraph) AddDomain(domainNode *DomainNode) {
 
 // NumDomains returns the number of domains in the graph
 func (graph *CertGraph) NumDomains() int {
-	return graph.numDomains
+	return int(atomic.LoadInt64(&graph.numDomains))
 }
 
 // DomainDepth returns the maximum depth of the graph from the initial root domains