WIP - New interface for modules to given greater control over connection (L4/TLS/etc) establishment

Dockerfile

@@ -13,7 +13,7 @@ RUN go mod download && go mod verify
 
 # Build the actual app
 COPY . .
-RUN make all
+RUN make clean; make all
 
 ## Runtime image ##
 FROM alpine:3.21 AS run

bin/default_modules.go

@@ -2,57 +2,56 @@ package bin
 
 import (
 	"github.com/zmap/zgrab2"
-	"github.com/zmap/zgrab2/modules"
-	"github.com/zmap/zgrab2/modules/bacnet"
-	"github.com/zmap/zgrab2/modules/banner"
-	"github.com/zmap/zgrab2/modules/dnp3"
-	"github.com/zmap/zgrab2/modules/fox"
-	"github.com/zmap/zgrab2/modules/ftp"
-	"github.com/zmap/zgrab2/modules/http"
-	"github.com/zmap/zgrab2/modules/imap"
-	"github.com/zmap/zgrab2/modules/ipp"
-	"github.com/zmap/zgrab2/modules/modbus"
-	"github.com/zmap/zgrab2/modules/mongodb"
-	"github.com/zmap/zgrab2/modules/mssql"
-	"github.com/zmap/zgrab2/modules/mysql"
+	//"github.com/zmap/zgrab2/modules/bacnet"
+	//"github.com/zmap/zgrab2/modules/banner"
+	//"github.com/zmap/zgrab2/modules/dnp3"
+	//"github.com/zmap/zgrab2/modules/fox"
+	//"github.com/zmap/zgrab2/modules/ftp"
+	//"github.com/zmap/zgrab2/modules/http"
+	//"github.com/zmap/zgrab2/modules/imap"
+	//"github.com/zmap/zgrab2/modules/ipp"
+	//"github.com/zmap/zgrab2/modules/modbus"
+	//"github.com/zmap/zgrab2/modules/mongodb"
+	//"github.com/zmap/zgrab2/modules/mssql"
+	//"github.com/zmap/zgrab2/modules/mysql"
 	"github.com/zmap/zgrab2/modules/ntp"
-	"github.com/zmap/zgrab2/modules/oracle"
-	"github.com/zmap/zgrab2/modules/pop3"
-	"github.com/zmap/zgrab2/modules/postgres"
-	"github.com/zmap/zgrab2/modules/redis"
-	"github.com/zmap/zgrab2/modules/siemens"
-	"github.com/zmap/zgrab2/modules/smb"
-	"github.com/zmap/zgrab2/modules/smtp"
-	"github.com/zmap/zgrab2/modules/telnet"
+	//"github.com/zmap/zgrab2/modules/oracle"
+	//"github.com/zmap/zgrab2/modules/pop3"
+	//"github.com/zmap/zgrab2/modules/postgres"
+	//"github.com/zmap/zgrab2/modules/redis"
+	//"github.com/zmap/zgrab2/modules/siemens"
+	//"github.com/zmap/zgrab2/modules/smb"
+	//"github.com/zmap/zgrab2/modules/smtp"
+	//"github.com/zmap/zgrab2/modules/telnet"
 )
 
 var defaultModules zgrab2.ModuleSet
 
 func init() {
 	defaultModules = map[string]zgrab2.ScanModule{
-		"bacnet":   &bacnet.Module{},
-		"banner":   &banner.Module{},
-		"dnp3":     &dnp3.Module{},
-		"fox":      &fox.Module{},
-		"ftp":      &ftp.Module{},
-		"http":     &http.Module{},
-		"imap":     &imap.Module{},
-		"ipp":      &ipp.Module{},
-		"modbus":   &modbus.Module{},
-		"mongodb":  &mongodb.Module{},
-		"mssql":    &mssql.Module{},
-		"mysql":    &mysql.Module{},
-		"ntp":      &ntp.Module{},
-		"oracle":   &oracle.Module{},
-		"pop3":     &pop3.Module{},
-		"postgres": &postgres.Module{},
-		"redis":    &redis.Module{},
-		"siemens":  &siemens.Module{},
-		"smb":      &smb.Module{},
-		"smtp":     &smtp.Module{},
-		"ssh":      &modules.SSHModule{},
-		"telnet":   &telnet.Module{},
-		"tls":      &modules.TLSModule{},
+		//"bacnet":   &bacnet.Module{},
+		//"banner":   &banner.Module{},
+		//"dnp3":     &dnp3.Module{},
+		//"fox":      &fox.Module{},
+		//"ftp":      &ftp.Module{},
+		//"http":     &http.Module{},
+		//"imap":     &imap.Module{},
+		//"ipp":      &ipp.Module{},
+		//"modbus":   &modbus.Module{},
+		//"mongodb":  &mongodb.Module{},
+		//"mssql":    &mssql.Module{},
+		//"mysql":    &mysql.Module{},
+		"ntp": &ntp.Module{},
+		//"oracle":   &oracle.Module{},
+		//"pop3":     &pop3.Module{},
+		//"postgres": &postgres.Module{},
+		//"redis":    &redis.Module{},
+		//"siemens":  &siemens.Module{},
+		//"smb":      &smb.Module{},
+		//"smtp":     &smtp.Module{},
+		//"ssh":      &modules.SSHModule{},
+		//"telnet":   &telnet.Module{},
+		//"tls":      &modules.TLSModule{},
 	}
 }
 

conn.go

@@ -164,13 +164,6 @@ func (c *TimeoutConnection) SetDeadline(deadline time.Time) error {
 	return nil
 }
 
-// GetTimeoutDialFunc returns a DialFunc that dials with the given timeout
-func GetTimeoutDialFunc(timeout time.Duration) func(string, string) (net.Conn, error) {
-	return func(proto, target string) (net.Conn, error) {
-		return DialTimeoutConnection(proto, target, timeout, 0)
-	}
-}
-
 // Close the underlying connection.
 func (c *TimeoutConnection) Close() error {
 	return c.Conn.Close()
@@ -233,7 +226,7 @@ func NewTimeoutConnection(ctx context.Context, conn net.Conn, timeout, readTimeo
 }
 
 // DialTimeoutConnectionEx dials the target and returns a net.Conn that uses the configured timeouts for Read/Write operations.
-func DialTimeoutConnectionEx(proto string, target string, dialTimeout, sessionTimeout, readTimeout, writeTimeout time.Duration, bytesReadLimit int) (net.Conn, error) {
+func DialTimeoutConnectionEx(ctx context.Context, proto string, target string, dialTimeout, sessionTimeout, readTimeout, writeTimeout time.Duration, bytesReadLimit int) (net.Conn, error) {
 	var conn net.Conn
 	var err error
 	if dialTimeout > 0 {
@@ -247,12 +240,12 @@ func DialTimeoutConnectionEx(proto string, target string, dialTimeout, sessionTi
 		}
 		return nil, err
 	}
-	return NewTimeoutConnection(context.Background(), conn, sessionTimeout, readTimeout, writeTimeout, bytesReadLimit), nil
+	return NewTimeoutConnection(ctx, conn, sessionTimeout, readTimeout, writeTimeout, bytesReadLimit), nil
 }
 
 // DialTimeoutConnection dials the target and returns a net.Conn that uses the configured single timeout for all operations.
-func DialTimeoutConnection(proto string, target string, timeout time.Duration, bytesReadLimit int) (net.Conn, error) {
-	return DialTimeoutConnectionEx(proto, target, timeout, timeout, timeout, timeout, bytesReadLimit)
+func DialTimeoutConnection(ctx context.Context, proto string, target string, timeout time.Duration, bytesReadLimit int) (net.Conn, error) {
+	return DialTimeoutConnectionEx(ctx, proto, target, timeout, timeout, timeout, timeout, bytesReadLimit)
 }
 
 // Dialer provides Dial and DialContext methods to get connections with the given timeout.
@@ -316,7 +309,7 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.
 
 // Dial returns a connection with the configured timeout.
 func (d *Dialer) Dial(proto string, target string) (net.Conn, error) {
-	return DialTimeoutConnectionEx(proto, target, d.ConnectTimeout, d.Timeout, d.ReadTimeout, d.WriteTimeout, 0)
+	return DialTimeoutConnectionEx(context.Background(), proto, target, d.ConnectTimeout, d.Timeout, d.ReadTimeout, d.WriteTimeout, 0)
 }
 
 // GetTimeoutConnectionDialer gets a Dialer that dials connections with the given timeout.

conn_bytelimit_test.go

@@ -257,7 +257,7 @@ type directDial struct {
 }
 
 func (d *directDial) connect(ctx context.Context, t *testing.T, port int, idx int) (*TimeoutConnection, error) {
-	conn, err := DialTimeoutConnectionEx("tcp", fmt.Sprintf("127.0.0.1:%d", port), time.Second, time.Second, time.Second, time.Second, d.limit)
+	conn, err := DialTimeoutConnectionEx(context.Background(), "tcp", fmt.Sprintf("127.0.0.1:%d", port), time.Second, time.Second, time.Second, time.Second, d.limit)
 	var ret *TimeoutConnection
 	if conn != nil {
 		ret = conn.(*TimeoutConnection)

conn_timeout_test.go

@@ -180,7 +180,7 @@ func (cfg *connTimeoutTestConfig) dialerDial() (*TimeoutConnection, error) {
 
 // Dial a connection to the configured endpoint using a DialTimeoutConnectionEx
 func (cfg *connTimeoutTestConfig) directDial() (*TimeoutConnection, error) {
-	ret, err := DialTimeoutConnectionEx("tcp", cfg.getEndpoint(), cfg.connectTimeout, cfg.timeout, cfg.readTimeout, cfg.writeTimeout, 0)
+	ret, err := DialTimeoutConnectionEx(context.Background(), "tcp", cfg.getEndpoint(), cfg.connectTimeout, cfg.timeout, cfg.readTimeout, cfg.writeTimeout, 0)
 	if err != nil {
 		return nil, err
 	}

go.mod

@@ -7,6 +7,7 @@ toolchain go1.23.5
 require (
 	github.com/hdm/jarm-go v0.0.7
 	github.com/prometheus/client_golang v1.21.1
+	github.com/quic-go/quic-go v0.50.0
 	github.com/rabbitmq/amqp091-go v1.10.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/zmap/zcrypto v0.0.0-20250129210703-03c45d0bae98
@@ -23,15 +24,23 @@ require (
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
-	github.com/klauspost/compress v1.18.0 // indirect
+	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
+	github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect
+	github.com/klauspost/compress v1.17.11 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/onsi/ginkgo/v2 v2.9.5 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/rogpeppe/go-internal v1.10.0 // indirect
 	github.com/weppos/publicsuffix-go v0.40.3-0.20250127173806-e489a31678ca // indirect
 	github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248 // indirect
-	google.golang.org/protobuf v1.36.5 // indirect
+	go.uber.org/mock v0.5.0 // indirect
+	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
+	golang.org/x/mod v0.18.0 // indirect
+	golang.org/x/sync v0.12.0 // indirect
+	golang.org/x/tools v0.22.0 // indirect
+	google.golang.org/protobuf v1.36.3 // indirect
 )