L4 module for Postgres: matchers

modules/l4postgres/matcher.go

@@ -12,7 +12,33 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package l4postgres allows the L4 multiplexing of Postgres connections
+// Package l4postgres allows the L4 multiplexing of Postgres connections.
+// Connections can be required to have SSL disabled.
+// Non-SSL connections can also match on Message parameters.
+//
+// Example matcher configs:
+//
+//	{
+//		"postgres": {}
+//	}
+//
+//	{
+//		"postgres": {
+//			"user": {
+//				"*": ["public_db"],
+//				"alice": ["planets_db", "stars_db"]
+//			}
+//		}
+//	}
+//
+//	{
+//		"postgres_client": ["psql", "TablePlus"]
+//	}
+//
+//	{
+//		"postgres_ssl": {
+//			disabled: false
+//	}
 //
 // With thanks to docs and code published at these links:
 // ref: https://github.com/mholt/caddy-l4/blob/master/modules/l4ssh/matcher.go
@@ -27,59 +53,46 @@ import (
 	"encoding/binary"
 	"errors"
 	"io"
+	"slices"
 
 	"github.com/caddyserver/caddy/v2"
 	"github.com/mholt/caddy-l4/layer4"
 )
 
 func init() {
 	caddy.RegisterModule(MatchPostgres{})
+	caddy.RegisterModule(MatchPostgresClient{})
+	caddy.RegisterModule(MatchPostgresSSL{})
 }
 
 const (
 	// Magic number to identify a SSLRequest message
 	sslRequestCode = 80877103
-	// byte size of the message length field
-	initMessageSizeLength = 4
 )
 
-// Message provides readers for various types and
-// updates the offset after each read
-type message struct {
-	data   []byte
-	offset uint32
-}
-
-func (b *message) ReadUint32() (r uint32) {
-	r = binary.BigEndian.Uint32(b.data[b.offset : b.offset+4])
-	b.offset += 4
-	return r
-}
+// NewMessageFromConn create a message from the Connection
+func newMessageFromConn(cx *layer4.Connection) (*message, error) {
+	// Get bytes containing the message length
+	head := make([]byte, lengthFieldSize)
+	if _, err := io.ReadFull(cx, head); err != nil {
+		return nil, err
+	}
 
-func (b *message) ReadString() (r string) {
-	end := b.offset
-	max := uint32(len(b.data))
-	for ; end != max && b.data[end] != 0; end++ {
+	// Get actual message length
+	data := make([]byte, binary.BigEndian.Uint32(head)-lengthFieldSize)
+	if _, err := io.ReadFull(cx, data); err != nil {
+		return nil, err
 	}
-	r = string(b.data[b.offset:end])
-	b.offset = end + 1
-	return r
-}
 
-// NewMessageFromBytes wraps the raw bytes of a message to enable processing
-func newMessageFromBytes(b []byte) *message {
-	return &message{data: b}
+	return newMessageFromBytes(data), nil
 }
 
-// StartupMessage contains the values parsed from the startup message
-type startupMessage struct {
-	ProtocolVersion uint32
-	Parameters      map[string]string
+// MatchPostgres is able to match Postgres connections
+type MatchPostgres struct {
+	User    map[string][]string
+	startup *startupMessage
 }
 
-// MatchPostgres is able to match Postgres connections.
-type MatchPostgres struct{}
-
 // CaddyModule returns the Caddy module information.
 func (MatchPostgres) CaddyModule() caddy.ModuleInfo {
 	return caddy.ModuleInfo{
@@ -88,46 +101,139 @@ func (MatchPostgres) CaddyModule() caddy.ModuleInfo {
 	}
 }
 
-// Match returns true if the connection looks like the Postgres protocol.
+// Match returns true if the connection looks like the Postgres protocol, and
+// can match `user` and `database` parameters
 func (m MatchPostgres) Match(cx *layer4.Connection) (bool, error) {
-	// Get bytes containing the message length
-	head := make([]byte, initMessageSizeLength)
-	if _, err := io.ReadFull(cx, head); err != nil {
+	b, err := newMessageFromConn(cx)
+	if err != nil {
 		return false, err
 	}
 
-	// Get actual message length
-	data := make([]byte, binary.BigEndian.Uint32(head)-initMessageSizeLength)
-	if _, err := io.ReadFull(cx, data); err != nil {
-		return false, err
+	m.startup = newStartupMessage(b)
+	hasConfig := len(m.User) > 0
+
+	// Finish if this is a SSLRequest and there are no other matchers
+	if m.startup.IsSSL() && !hasConfig {
+		return true, nil
 	}
 
-	b := newMessageFromBytes(data)
+	// Check supported protocol
+	if !m.startup.IsSupported() {
+		return false, errors.New("pg protocol < 3.0 is not supported")
+	}
 
-	// Check if it is a SSLRequest
-	code := b.ReadUint32()
-	if code == sslRequestCode {
+	// Finish if no more matchers are configured
+	if !hasConfig {
 		return true, nil
 	}
 
+	// Is there a user to check?
+	user, ok := m.startup.Parameters["user"]
+	if !ok {
+		// Are there public databases to check?
+		if databases, ok := m.User["*"]; ok {
+			if db, ok := m.startup.Parameters["database"]; ok {
+				return slices.Contains(databases, db), nil
+			}
+		}
+		return false, nil
+	}
+
+	databases, ok := m.User[user]
+	if !ok {
+		return false, nil
+	}
+
+	// Are there databases to check?
+	if len(databases) > 0 {
+		if db, ok := m.startup.Parameters["database"]; ok {
+			return slices.Contains(databases, db), nil
+		}
+	}
+
+	return true, nil
+}
+
+// MatchPostgresClient is able to match Postgres connections that
+// contain an `application_name` field
+type MatchPostgresClient struct {
+	Client  []string
+	startup *startupMessage
+}
+
+// CaddyModule returns the Caddy module information.
+func (MatchPostgresClient) CaddyModule() caddy.ModuleInfo {
+	return caddy.ModuleInfo{
+		ID:  "layer4.matchers.postgres_client",
+		New: func() caddy.Module { return new(MatchPostgresClient) },
+	}
+}
+
+// Match returns true if the connection looks like the Postgres protocol and
+// passes any `application_name` parameter matchers
+func (m MatchPostgresClient) Match(cx *layer4.Connection) (bool, error) {
+	b, err := newMessageFromConn(cx)
+	if err != nil {
+		return false, err
+	}
+
+	m.startup = newStartupMessage(b)
+
+	// Reject if this is a SSLRequest as it has no params
+	if m.startup.IsSSL() {
+		return false, nil
+	}
+
 	// Check supported protocol
-	if majorVersion := code >> 16; majorVersion < 3 {
+	if !m.startup.IsSupported() {
 		return false, errors.New("pg protocol < 3.0 is not supported")
 	}
 
-	// Try parsing Postgres Params
-	startup := &startupMessage{ProtocolVersion: code, Parameters: make(map[string]string)}
-	for {
-		k := b.ReadString()
-		if k == "" {
-			break
-		}
-		startup.Parameters[k] = b.ReadString()
+	// Is there a application_name to check?
+	name, ok := m.startup.Parameters["application_name"]
+	if !ok {
+		return false, nil
+	}
+
+	// Check clients list
+	return slices.Contains(m.Client, name), nil
+}
+
+// MatchPostgresSSL is able to require/reject Postgres SSL connections.
+type MatchPostgresSSL struct {
+	Disabled bool
+}
+
+// CaddyModule returns the Caddy module information.
+func (MatchPostgresSSL) CaddyModule() caddy.ModuleInfo {
+	return caddy.ModuleInfo{
+		ID:  "layer4.matchers.postgres_ssl",
+		New: func() caddy.Module { return new(MatchPostgresSSL) },
 	}
-	// TODO(metafeather): match on param values: user, database, options, etc
+}
 
-	return len(startup.Parameters) > 0, nil
+// Match checks whether the connection is a Postgres SSL request.
+func (m MatchPostgresSSL) Match(cx *layer4.Connection) (bool, error) {
+	b, err := newMessageFromConn(cx)
+	if err != nil {
+		return false, err
+	}
+
+	code := b.ReadUint32()
+	disabled := !isSSLRequest(code)
+
+	// Enforce SSL enabled
+	if !m.Disabled && !disabled {
+		return true, nil
+	}
+	// Enforce SSL disabled
+	if m.Disabled && disabled {
+		return true, nil
+	}
+	return false, nil
 }
 
 // Interface guard
 var _ layer4.ConnMatcher = (*MatchPostgres)(nil)
+var _ layer4.ConnMatcher = (*MatchPostgresClient)(nil)
+var _ layer4.ConnMatcher = (*MatchPostgresSSL)(nil)