Add memory and CPU profiling (if ZGRAB_[MEM/CPU]PROFILE is set); add upper bound on data read from postgres

cmd/zgrab2/main.go

@@ -4,15 +4,89 @@ import (
 	"encoding/json"
 	"os"
 	"time"
+	"runtime/pprof"
 
 	flags "github.com/zmap/zflags"
 	log "github.com/sirupsen/logrus"
 	"github.com/zmap/zgrab2"
 	_ "github.com/zmap/zgrab2/modules"
+	"fmt"
+	"runtime"
+	"strings"
 )
 
+// Get the value of the ZGRAB2_MEMPROFILE variable (or the empty string).
+// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
+// getFormattedFile().
+func getMemProfileFile() string {
+	return os.Getenv("ZGRAB2_MEMPROFILE")
+}
+
+// Get the value of the ZGRAB2_CPUPROFILE variable (or the empty string).
+// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
+// getFormattedFile().
+func getCPUProfileFile() string {
+	return os.Getenv("ZGRAB2_CPUPROFILE")
+}
+
+// Replace instances in formatString of {TIMESTAMP} with when formatted as
+// YYYYMMDDhhmmss, and {NANOS} as the decimal nanosecond offset.
+func getFormattedFile(formatString string, when time.Time) string {
+	timestamp := when.Format("20060102150405")
+	nanos := fmt.Sprintf("%d", when.Nanosecond())
+	ret := strings.Replace(formatString, "{TIMESTAMP}", timestamp, -1)
+	ret = strings.Replace(ret, "{NANOS}", nanos, -1)
+	return ret
+}
+
+// If memory profiling is enabled (ZGRAB2_MEMPROFILE is not empty), perform a GC
+// then write the heap profile to the profile file.
+func dumpHeapProfile() {
+	if file := getMemProfileFile(); file != "" {
+		now := time.Now()
+		fullFile := getFormattedFile(file, now)
+		f, err := os.Create(fullFile)
+		if err != nil {
+			log.Fatal("could not create heap profile: ", err)
+		}
+		runtime.GC()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("could not write heap profile: ", err)
+		}
+		f.Close()
+	}
+}
+
+// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), start tracking
+// CPU profiling in the configured file. Caller is responsible for invoking
+// stopCPUProfile() when finished.
+func startCPUProfile() {
+	if file := getCPUProfileFile(); file != "" {
+		now := time.Now()
+		fullFile := getFormattedFile(file, now)
+		f, err := os.Create(fullFile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+	}
+}
+
+// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), stop profiling
+// CPU usage.
+func stopCPUProfile() {
+	if getCPUProfileFile() != "" {
+		pprof.StopCPUProfile()
+	}
+}
 func main() {
+	startCPUProfile()
+	defer stopCPUProfile()
+	defer dumpHeapProfile()
 	_, moduleType, flag, err := zgrab2.ParseCommandLine(os.Args[1:])
+
 	// Blanked arg is positional arguments
 	if err != nil {
 		// Outputting help is returned as an error. Exit successfuly on help output.
@@ -54,6 +128,9 @@ func main() {
 		zgrab2.RegisterScan(moduleType, s)
 	}
 	monitor := zgrab2.MakeMonitor()
+	monitor.Callback = func(_ string) {
+		dumpHeapProfile()
+	}
 	start := time.Now()
 	log.Infof("started grab at %s", start.Format(time.RFC3339))
 	zgrab2.Process(monitor)

modules/postgres/connection.go

@@ -14,8 +14,14 @@ import (
 	"github.com/zmap/zgrab2"
 )
 
+// Don't allow unbounded reads
+const maxPacketSize = 128 * 1024 * 1024
+
 // Connection wraps the state of a given connection to a server.
 type Connection struct {
+	// Target is the requested scan target.
+	Target *zgrab2.ScanTarget
+
 	// Connection is the underlying TCP (or TLS) stream.
 	Connection net.Conn
 
@@ -104,11 +110,20 @@ func (c *Connection) tryReadPacket(header byte) (*ServerPacket, *zgrab2.ScanErro
 		}
 		return nil, zgrab2.NewScanError(zgrab2.SCAN_PROTOCOL_ERROR, fmt.Errorf("Server returned too much data: length = 0x%x; first %d bytes = %s", ret.Length, n, hex.EncodeToString(buf[:n])))
 	}
-	ret.Body = make([]byte, ret.Length-4) // Length includes the length of the Length uint32
+	sizeToRead := ret.Length
+	if sizeToRead > maxPacketSize {
+		log.Debugf("postgres server %s reported packet size of %d bytes; only reading %d bytes.", c.Target.String(), ret.Length, maxPacketSize)
+		sizeToRead = maxPacketSize
+	}
+	ret.Body = make([]byte, sizeToRead) // Length includes the length of the Length uint32
 	_, err = io.ReadFull(c.Connection, ret.Body)
 	if err != nil && err != io.EOF {
 		return nil, zgrab2.DetectScanError(err)
 	}
+	if sizeToRead < ret.Length && len(ret.Body) >= maxPacketSize {
+		// Warn if we actually truncate (as opposed getting an huge length but only a few bytes are actually available)
+		log.Warnf("Truncated postgres packet from %s: advertised size = %d bytes, read size = %d bytes", c.Target.String(), ret.Length, len(ret.Body))
+	}
 	return &ret, nil
 }
 

modules/postgres/scanner.go

@@ -330,7 +330,7 @@ func (s *Scanner) newConnection(t *zgrab2.ScanTarget, mgr *connectionManager, no
 		return nil, zgrab2.DetectScanError(err)
 	}
 	mgr.addConnection(conn)
-	sql := Connection{Connection: conn, Config: s.Config}
+	sql := Connection{Target: t, Connection: conn, Config: s.Config}
 	sql.IsSSL = false
 	if !nossl && !s.Config.SkipSSL {
 		hasSSL, sslError := sql.RequestSSL()

monitor.go

@@ -5,6 +5,8 @@ package zgrab2
 type Monitor struct {
 	states       map[string]*State
 	statusesChan chan moduleStatus
+	// Callback is invoked after each scan.
+	Callback     func(string)
 }
 
 // State contains the respective number of successes and failures
@@ -43,6 +45,9 @@ func MakeMonitor() *Monitor {
 			if m.states[s.name] == nil {
 				m.states[s.name] = new(State)
 			}
+			if m.Callback != nil {
+				m.Callback(s.name)
+			}
 			switch s.st {
 			case statusSuccess:
 				m.states[s.name].Successes++