feat: Implement byte matching in TCP query responses

Currently the exporter only supports lines, which breaks byte-oriented
protocols such as the PostgreSQL StartTLS handshake.

We also give a working example for Postgres in the sample configuration.

Signed-off-by: Stanislav Grozev <[email protected]>
Co-authored-by: Conrad Hoffmann <[email protected]>

Author: tacho

CONFIGURATION.md

@@ -196,13 +196,15 @@ regexp: <regex>,
 [ source_ip_address: <string> ]
 
 # The query sent in the TCP probe and the expected associated response.
-# "expect" matches a regular expression;
+# "expect" matches a regular expression; it is mutually exclusive with "expect_bytes".
+# "expect_bytes" does exact byte-by-byte match; it is mutually exclusive with "expect".
 # "labels" can define labels which will be exported on metric "probe_expect_info";
 # "send" sends some content;
 # "send" and "labels.value" can contain values matched by "expect" (such as "${1}");
 # "starttls" upgrades TCP connection to TLS.
 query_response:
   [ - [ [ expect: <string> ],
+        [ expect_bytes: <string> ],
         [ labels:
           - [ name: <string>
               value: <string>

blackbox.yml

@@ -67,3 +67,10 @@ modules:
       enable_http3: true
       enable_http2: false
       valid_http_versions: ["HTTP/3.0"]
+  postgresql:
+    prober: tcp
+    tcp:
+      query_response:
+      - send: !!binary AAAACATSFi8= # 0x00, 0x00, 0x00, 0x08, 0x04, 0xD2, 0x16, 0x2F - PostgreSQL SSLRequest
+      - expect_bytes: S # 0x53 - Reply will be 'S' if SSL is enabled, and 'N' if it is not.
+      - starttls: true

config/config.go

@@ -341,10 +341,11 @@ type Label struct {
 }
 
 type QueryResponse struct {
-	Expect   Regexp  `yaml:"expect,omitempty"`
-	Labels   []Label `yaml:"labels,omitempty"`
-	Send     string  `yaml:"send,omitempty"`
-	StartTLS bool    `yaml:"starttls,omitempty"`
+	Expect      Regexp  `yaml:"expect,omitempty"`
+	ExpectBytes string  `yaml:"expect_bytes,omitempty"`
+	Labels      []Label `yaml:"labels,omitempty"`
+	Send        string  `yaml:"send,omitempty"`
+	StartTLS    bool    `yaml:"starttls,omitempty"`
 }
 
 type TCPProbe struct {
@@ -568,7 +569,9 @@ func (s *QueryResponse) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if err := unmarshal((*plain)(s)); err != nil {
 		return err
 	}
-
+	if s.Expect.Regexp != nil && s.ExpectBytes != "" {
+		return errors.New("expect and expect_bytes are mutually exclusive")
+	}
 	return nil
 }
 

example.yml

@@ -222,3 +222,10 @@ modules:
       query_response:
         - send: "PING"
         - expect: "PONG"
+  postgresql:
+    prober: tcp
+    tcp:
+      query_response:
+      - send: !!binary AAAACATSFi8= # 0x00, 0x00, 0x00, 0x08, 0x04, 0xD2, 0x16, 0x2F - PostgreSQL SSLRequest
+      - expect_bytes: S # 0x53 - Reply will be 'S' if SSL is enabled, and 'N' if it is not.
+      - starttls: true

prober/query_response.go

@@ -15,6 +15,7 @@ package prober
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
@@ -63,7 +64,12 @@ func probeQueryResponses(ctx context.Context, target string, conn net.Conn, modu
 		Name: "probe_failed_due_to_regex",
 		Help: "Indicates if probe failed due to regex",
 	})
+	probeFailedDueToBytes := prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "probe_failed_due_to_bytes",
+		Help: "Indicates if probe failed due to bytes",
+	})
 	registry.MustRegister(probeFailedDueToRegex)
+	registry.MustRegister(probeFailedDueToBytes)
 
 	var queryResponses []config.QueryResponse
 	var tlsConfig *pconfig.TLSConfig
@@ -125,6 +131,32 @@ func probeQueryResponses(ctx context.Context, target string, conn net.Conn, modu
 				probeExpectInfo(registry, &qr, scanner.Bytes(), match)
 			}
 		}
+		if qr.ExpectBytes != "" {
+			expect_bytes := []byte(qr.ExpectBytes)
+
+			// Try to read same number of bytes as expected.
+			data := make([]byte, len(expect_bytes))
+			n, err := conn.Read(data)
+			if err != nil {
+				logger.Error("Error reading from connection", "err", err)
+				return false
+			}
+
+			logger.Debug("Read bytes", "bytes", data)
+
+			if n < len(expect_bytes) {
+				logger.Error("Read less data than expected", "expected", expect_bytes, "bytes", data)
+				return false
+			}
+
+			if !bytes.Equal(expect_bytes, data) {
+				probeFailedDueToRegex.Set(1)
+				logger.Error("Bytes did not match", "expected", expect_bytes, "bytes", data)
+				return false
+			}
+			logger.Debug("Bytes matched", "expected", expect_bytes, "bytes", data)
+			probeFailedDueToRegex.Set(0)
+		}
 		if send != "" {
 			logger.Debug("Sending line", "line", send)
 			if _, err := fmt.Fprintf(conn, "%s\n", send); err != nil {

prober/tcp_test.go

@@ -580,6 +580,49 @@ func TestTCPConnectionQueryResponseMatching(t *testing.T) {
 
 }
 
+func TestTCPConnectionQueryResponseByteMode(t *testing.T) {
+	ln, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("Error listening on socket: %s", err)
+	}
+	defer ln.Close()
+
+	testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	module := config.Module{
+		TCP: config.TCPProbe{
+			IPProtocolFallback: true,
+			QueryResponse: []config.QueryResponse{
+				{
+					ExpectBytes: "not-a-line",
+				},
+			},
+		},
+	}
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			panic(fmt.Sprintf("Error accepting on socket: %s", err))
+		}
+		conn.SetDeadline(time.Now().Add(1 * time.Second))
+		conn.Write([]byte("not-a-line"))
+		conn.Close()
+	}()
+	registry := prometheus.NewRegistry()
+	if !ProbeTCP(testCTX, ln.Addr().String(), module, registry, promslog.NewNopLogger()) {
+		t.Fatalf("TCP module failed, expected success.")
+	}
+	mfs, err := registry.Gather()
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedResults := map[string]float64{
+		"probe_failed_due_to_regex": 0,
+	}
+	checkRegistryResults(expectedResults, mfs, t)
+}
+
 func TestTCPConnectionProtocol(t *testing.T) {
 	if os.Getenv("CI") == "true" {
 		t.Skip("skipping; CI is failing on ipv6 dns requests")