drivers implement context

Author: lanrat

certgraph.go

@@ -11,6 +11,7 @@
 package main
 
 import (
+	"context"
 	"embed"
 	"encoding/json"
 	"flag"
@@ -389,7 +390,9 @@ func visit(domainNode *graph.DomainNode) {
 
 	// perform cert search
 	// TODO do pagination in multiple threads to not block on long searches
-	results, err := certDriver.QueryDomain(domainNode.Domain)
+	ctx, cancel := context.WithTimeout(context.Background(), config.timeout)
+	defer cancel()
+	results, err := certDriver.QueryDomain(ctx, domainNode.Domain)
 	if err != nil {
 		// this is VERY common to error, usually this is a DNS or tcp connection related issue
 		// we will skip the domain if we can't query it
@@ -430,7 +433,7 @@ func visit(domainNode *graph.DomainNode) {
 			processedCerts[fp] = true
 
 			// get cert details
-			certResult, err := results.QueryCert(fp)
+			certResult, err := results.QueryCert(ctx, fp)
 			if err != nil {
 				v("QueryCert", err)
 				continue

driver/censys/censys.go

@@ -7,6 +7,7 @@ package censys
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"flag"
@@ -66,8 +67,8 @@ func (c *censysCertDriver) GetRelated() ([]string, error) {
 	return nil, nil // Return nil instead of empty slice for better memory efficiency
 }
 
-func (c *censysCertDriver) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
-	return c.driver.QueryCert(fp)
+func (c *censysCertDriver) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+	return c.driver.QueryCert(ctx, fp)
 }
 
 // TODO support pagination
@@ -188,7 +189,7 @@ func (d *censys) jsonRequest(method, url string, request, response interface{})
 	return nil
 }
 
-func (d *censys) QueryDomain(domain string) (driver.Result, error) {
+func (d *censys) QueryDomain(ctx context.Context, domain string) (driver.Result, error) {
 	results := &censysCertDriver{
 		host:         domain,
 		fingerprints: make(driver.FingerprintMap),
@@ -218,7 +219,7 @@ func (d *censys) QueryDomain(domain string) (driver.Result, error) {
 	return results, nil
 }
 
-func (d *censys) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+func (d *censys) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
 	certNode := new(driver.CertResult)
 	certNode.Fingerprint = fp
 	certNode.Domains = make([]string, 0, 5)

driver/crtsh/crtsh.go

@@ -8,6 +8,7 @@
 package crtsh
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"log"
@@ -57,8 +58,8 @@ func (c *crtshCertDriver) GetRelated() ([]string, error) {
 	return nil, nil // Return nil instead of empty slice for better memory efficiency
 }
 
-func (c *crtshCertDriver) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
-	return c.driver.QueryCert(fp)
+func (c *crtshCertDriver) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+	return c.driver.QueryCert(ctx, fp)
 }
 
 // Driver creates a new CT driver for crt.sh
@@ -108,7 +109,7 @@ func (d *crtsh) setSQLTimeout(sec float64) error {
 	return err
 }
 
-func (d *crtsh) QueryDomain(domain string) (driver.Result, error) {
+func (d *crtsh) QueryDomain(ctx context.Context, domain string) (driver.Result, error) {
 	results := &crtshCertDriver{
 		host:         domain,
 		fingerprints: make(driver.FingerprintMap),
@@ -164,7 +165,7 @@ func (d *crtsh) QueryDomain(domain string) (driver.Result, error) {
 		if debug {
 			log.Printf("QueryDomain try %d: %s", try, queryStr)
 		}
-		rows, err = d.db.Query(queryStr, d.includeExpired, d.includeSubdomains, d.queryLimit, domain)
+		rows, err = d.db.QueryContext(ctx, queryStr, d.includeExpired, d.includeSubdomains, d.queryLimit, domain)
 		if err == nil {
 			break
 		}
@@ -202,7 +203,7 @@ func (d *crtsh) QueryDomain(domain string) (driver.Result, error) {
 	return results, nil
 }
 
-func (d *crtsh) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+func (d *crtsh) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
 	certNode := new(driver.CertResult)
 	certNode.Fingerprint = fp
 	certNode.Domains = make([]string, 0, 5)
@@ -217,7 +218,7 @@ func (d *crtsh) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error
 	for try < 5 {
 		// this is a hack while crt.sh gets there stuff together
 		try++
-		rows, err = d.db.Query(queryStr, fp[:])
+		rows, err = d.db.QueryContext(ctx, queryStr, fp[:])
 		if err == nil {
 			break
 		}

driver/driver.go

@@ -2,6 +2,7 @@
 package driver
 
 import (
+	"context"
 	"crypto/x509"
 	"sort"
 	"strings"
@@ -10,8 +11,6 @@ import (
 	"github.com/lanrat/certgraph/status"
 )
 
-// TODO add context instead of timeout on all requests
-
 // Drivers contains all the drivers that have been registered
 var Drivers []string
 
@@ -27,7 +26,7 @@ type Driver interface {
 	// QueryDomain is the main entrypoint for Driver Searching
 	// The domain provided will return a CertDriver instance which can be used to query the
 	// certificates for the provided domain using the driver
-	QueryDomain(domain string) (Result, error)
+	QueryDomain(ctx context.Context, domain string) (Result, error)
 
 	// GetName returns the name of the driver
 	GetName() string
@@ -46,7 +45,7 @@ type Result interface {
 	GetFingerprints() (FingerprintMap, error)
 
 	// QueryCert returns the details of the provided certificate or an error if not found
-	QueryCert(fp fingerprint.Fingerprint) (*CertResult, error)
+	QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*CertResult, error)
 }
 
 // FingerprintMap stores a mapping of domains to Fingerprints returned from the driver

driver/example.go

@@ -1,10 +1,14 @@
 package driver
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 // Example provides a simple entrypoint to test a driver on an individual domain
 func Example(domain string, driver Driver) error {
-	certDriver, err := driver.QueryDomain(domain)
+	ctx := context.Background()
+	certDriver, err := driver.QueryDomain(ctx, domain)
 	if err != nil {
 		return err
 	}
@@ -27,7 +31,7 @@ func Example(domain string, driver Driver) error {
 	for domain, fingerprints := range fingerprintMap {
 		for i := range fingerprints {
 			fmt.Printf("%s: %s\n", domain, fingerprints[i].HexString())
-			cert, err := certDriver.QueryCert(fingerprints[i])
+			cert, err := certDriver.QueryCert(ctx, fingerprints[i])
 			if err != nil {
 				return err
 			}

driver/http/http.go

@@ -2,6 +2,7 @@
 package http
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
@@ -58,7 +59,7 @@ func (c *httpCertDriver) GetRelated() ([]string, error) {
 
 // QueryCert retrieves certificate details for a specific fingerprint.
 // Returns an error if the certificate was not found in this HTTP query.
-func (c *httpCertDriver) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+func (c *httpCertDriver) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
 	cert, found := c.certs[fp]
 	if found {
 		return cert, nil
@@ -121,10 +122,15 @@ func (d *httpDriver) newHTTPCertDriver() *httpCertDriver {
 
 // QueryDomain discovers certificates for a domain through HTTPS connections.
 // Follows redirects and collects certificates from all encountered servers.
-func (d *httpDriver) QueryDomain(host string) (driver.Result, error) {
+func (d *httpDriver) QueryDomain(ctx context.Context, host string) (driver.Result, error) {
 	results := d.newHTTPCertDriver()
 
-	resp, err := results.client.Get(fmt.Sprintf("https://%s", host))
+	req, err := http.NewRequestWithContext(ctx, "GET", fmt.Sprintf("https://%s", host), nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+	
+	resp, err := results.client.Do(req)
 	fullStatus := status.CheckNetErr(err)
 	if fullStatus != status.GOOD {
 		return results, err // in some rare cases this error can be ignored

driver/multi/multi_driver.go

@@ -2,6 +2,7 @@
 package multi
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
@@ -39,14 +40,14 @@ func (d *multiDriver) GetName() string {
 
 // QueryDomain executes domain queries against all drivers concurrently.
 // Returns a merged result containing certificates and status information from all drivers.
-func (d *multiDriver) QueryDomain(domain string) (driver.Result, error) {
+func (d *multiDriver) QueryDomain(ctx context.Context, domain string) (driver.Result, error) {
 	r := newResult(domain)
-	var group errgroup.Group
+	group, ctx := errgroup.WithContext(ctx)
 	for _, d := range d.drivers {
 		goFunc := func(localDriver driver.Driver) func() error {
 			return func() error {
 				return func(localDriver driver.Driver) error {
-					result, err := localDriver.QueryDomain(domain)
+					result, err := localDriver.QueryDomain(ctx, domain)
 					if err != nil {
 						return err
 					}
@@ -104,9 +105,9 @@ func (c *multiResult) add(r driver.Result) error {
 
 // QueryCert attempts to retrieve certificate details from any of the drivers.
 // Returns the first successful result found among the combined drivers.
-func (c *multiResult) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+func (c *multiResult) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
 	for _, result := range c.results {
-		cr, err := result.QueryCert(fp)
+		cr, err := result.QueryCert(ctx, fp)
 		if err != nil {
 			return nil, err
 		}

driver/smtp/smtp.go

@@ -60,7 +60,7 @@ func (c *smtpCertDriver) GetRelated() ([]string, error) {
 
 // QueryCert retrieves certificate details for a specific fingerprint.
 // Returns an error if the certificate was not found in this SMTP query.
-func (c *smtpCertDriver) QueryCert(fp fingerprint.Fingerprint) (*driver.CertResult, error) {
+func (c *smtpCertDriver) QueryCert(ctx context.Context, fp fingerprint.Fingerprint) (*driver.CertResult, error) {
 	cert, found := c.certs[fp]
 	if found {
 		return cert, nil
@@ -119,7 +119,7 @@ func (d *smtpDriver) smtpGetCerts(host string) ([]*x509.Certificate, error) {
 
 // QueryDomain discovers certificates for a domain through SMTP STARTTLS.
 // Also performs MX record lookups to find related mail server domains.
-func (d *smtpDriver) QueryDomain(host string) (driver.Result, error) {
+func (d *smtpDriver) QueryDomain(ctx context.Context, host string) (driver.Result, error) {
 	results := &smtpCertDriver{
 		host:         host,
 		status:       make(status.Map),