Add support for postgresql protocol (#77)

With postgresql to initiate SSL-encrypted connection specific combination
of bytes must be sent to the server.

Message flow is described on following page
https://www.postgresql.org/docs/13/protocol-flow.html#id-1.10.5.7.11

And SSLRequest message format is described on
https://www.postgresql.org/docs/13/protocol-message-formats.html

The value of SSLRequest message becomes to bytes that is used in the code

README.md

@@ -300,7 +300,7 @@ prober: <prober_string>
 ### <tcp_probe>
 
 ```
-# Use the STARTTLS command before starting TLS for those protocols that support it (smtp, ftp, imap)
+# Use the STARTTLS command before starting TLS for those protocols that support it (smtp, ftp, imap, postgres)
 [ starttls: <string> ]
 ```
 

prober/tcp.go

@@ -2,9 +2,11 @@ package prober
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"net"
 	"regexp"
 
@@ -49,6 +51,8 @@ func ProbeTCP(ctx context.Context, logger log.Logger, target string, module conf
 type queryResponse struct {
 	expect      string
 	send        string
+	sendBytes   []byte
+	expectBytes []byte
 }
 
 var (
@@ -107,6 +111,14 @@ var (
 				expect: "OK",
 			},
 		},
+		"postgres": []queryResponse{
+			queryResponse{
+				sendBytes: []byte{0x00, 0x00, 0x00, 0x08, 0x04, 0xd2, 0x16, 0x2f},
+			},
+			queryResponse{
+				expectBytes: []byte{0x53},
+			},
+		},
 	}
 )
 
@@ -141,12 +153,31 @@ func startTLS(logger log.Logger, conn net.Conn, proto string) error {
 				return fmt.Errorf("regex: %s didn't match: %s", qr.expect, scanner.Text())
 			}
 		}
+		if len(qr.expectBytes) > 0 {
+			buffer := make([]byte, len(qr.expectBytes))
+			_, err = io.ReadFull(conn, buffer)
+			if err != nil {
+				return nil
+			}
+			level.Debug(logger).Log("msg", fmt.Sprintf("read bytes: %x", buffer))
+			if bytes.Compare(buffer, qr.expectBytes) != 0 {
+				return fmt.Errorf("read bytes %x didn't match with expected bytes %x", buffer, qr.expectBytes)
+			} else {
+				level.Debug(logger).Log("msg", fmt.Sprintf("expected bytes %x matched with read bytes %x", qr.expectBytes, buffer))
+			}
+		}
 		if qr.send != "" {
 			level.Debug(logger).Log("msg", fmt.Sprintf("sending line: %s", qr.send))
 			if _, err := fmt.Fprintf(conn, "%s\r\n", qr.send); err != nil {
 				return err
 			}
 		}
+		if len(qr.sendBytes) > 0 {
+			level.Debug(logger).Log("msg", fmt.Sprintf("sending bytes: %x", qr.sendBytes))
+			if _, err = conn.Write(qr.sendBytes); err != nil {
+				return err
+			}
+		}
 	}
 	return nil
 }

prober/tcp_test.go

@@ -325,6 +325,45 @@ func TestProbeTCPStartTLSIMAP(t *testing.T) {
 	checkTLSVersionMetrics("TLS 1.3", registry, t)
 }
 
+// TestProbeTCPStartTLSPostgreSQL tests STARTTLS against a mock PostgreSQL server
+func TestProbeTCPStartTLSPostgreSQL(t *testing.T) {
+	server, certPEM, _, caFile, teardown, err := test.SetupTCPServer()
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	defer teardown()
+
+	server.StartPostgreSQL()
+	defer server.Close()
+
+	module := config.Module{
+		TCP: config.TCPProbe{
+			StartTLS: "postgres",
+		},
+		TLSConfig: pconfig.TLSConfig{
+			CAFile:             caFile,
+			InsecureSkipVerify: false,
+		},
+	}
+
+	registry := prometheus.NewRegistry()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	if err := ProbeTCP(ctx, newTestLogger(), server.Listener.Addr().String(), module, registry); err != nil {
+		t.Fatalf("error: %s", err)
+	}
+
+	cert, err := newCertificate(certPEM)
+	if err != nil {
+		t.Fatal(err)
+	}
+	checkCertificateMetrics(cert, registry, t)
+	checkOCSPMetrics([]byte{}, registry, t)
+	checkTLSVersionMetrics("TLS 1.3", registry, t)
+}
+
 // TestProbeTCPTimeout tests that the TCP probe respects the timeout in the
 // context
 func TestProbeTCPTimeout(t *testing.T) {

test/tcp.go

@@ -1,8 +1,10 @@
 package test
 
 import (
+	"bytes"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"time"
@@ -162,6 +164,45 @@ func (t *TCPServer) StartIMAP() {
 	}()
 }
 
+// StartPostgreSQL starts a listener that negotiates a TLS connection with an postgresql
+// client using STARTTLS
+func (t *TCPServer) StartPostgreSQL() {
+	go func() {
+		conn, err := t.Listener.Accept()
+		if err != nil {
+			panic(fmt.Sprintf("Error accepting on socket: %s", err))
+		}
+		defer conn.Close()
+
+		sslRequestMessage := []byte{0x00, 0x00, 0x00, 0x08, 0x04, 0xd2, 0x16, 0x2f}
+
+		buffer := make([]byte, len(sslRequestMessage))
+
+		_, err = io.ReadFull(conn, buffer)
+		if err != nil {
+			panic("Error reading input from client")
+		}
+
+		if bytes.Compare(buffer, sslRequestMessage) != 0 {
+			panic(fmt.Sprintf("Error in dialog. No %x received", buffer))
+		}
+
+		sslRequestResponse := []byte{0x53}
+
+		if _, err := conn.Write(sslRequestResponse); err != nil {
+			panic("Error writing response to client")
+		}
+
+		tlsConn := tls.Server(conn, t.TLS)
+		if err := tlsConn.Handshake(); err != nil {
+			level.Error(t.logger).Log("msg", err)
+		}
+		defer tlsConn.Close()
+
+		t.stopCh <- struct{}{}
+	}()
+}
+
 // Close stops the server and closes the listener
 func (t *TCPServer) Close() {
 	<-t.stopCh