From c2823a933bd1369be1aa5814d2d0960fb77eafce Mon Sep 17 00:00:00 2001
From: Nia Weiss <nia.weiss@molecula.com>
Date: Wed, 26 Aug 2020 16:39:11 -0400
Subject: [PATCH] add a config option to require TLS on postgres

---
 pg/protocol.go   | 8 ++++++++
 pg/server.go     | 4 ++++
 server/config.go | 2 ++
 server/server.go | 1 +
 4 files changed, 15 insertions(+)

diff --git a/pg/protocol.go b/pg/protocol.go
index c55decbd7..2267b1ca7 100644
--- a/pg/protocol.go
+++ b/pg/protocol.go
@@ -100,6 +100,8 @@ func (s *Server) handle(ctx context.Context, conn net.Conn) (err error) {
 		}
 	}
 
+	var hasTLS bool
+
 startup:
 	// Read startup packet.
 	var buf [4]byte
@@ -151,6 +153,7 @@ startup:
 					return errors.Wrap(err, "transferring startup deadline to TLS connection")
 				}
 			}
+			hasTLS = true
 			goto startup
 		}
 
@@ -163,6 +166,11 @@ startup:
 		goto startup
 	}
 
+	if s.RequireTLS && !hasTLS {
+		// Reject the unsecured connection.
+		return errors.Errorf("client at %s attempted to initiate an unsecured postgres conenction", conn.RemoteAddr())
+	}
+
 	// Handle regular postgres.
 	return s.handleStandard(ctx, proto, conn, data)
 }
diff --git a/pg/server.go b/pg/server.go
index 55d374005..686c86ebe 100644
--- a/pg/server.go
+++ b/pg/server.go
@@ -35,6 +35,10 @@ type Server struct {
 	// TLSConfig is the TLS configuration to use to serve postgres TLS connections.
 	TLSConfig *tls.Config
 
+	// RequireTLS rejects unencrypted postgres connections.
+	// This allows for mTLS to be used as auth.
+	RequireTLS bool
+
 	// StartupTimeout is the timeout to use for connection startup.
 	// If a connection fails to set up a protocol before this completes, it will be terminated.
 	StartupTimeout time.Duration
diff --git a/server/config.go b/server/config.go
index 7437d65ed..5ad987dcd 100644
--- a/server/config.go
+++ b/server/config.go
@@ -174,6 +174,8 @@ type Config struct {
 		// TLS configuration for postgres connections.
 		TLS TLSConfig `toml:"tls"`
 
+		RequireTLS bool `toml:"require-tls"`
+
 		StartupTimeout toml.Duration `toml:"startup-timeout"`
 		ReadTimeout    toml.Duration `toml:"read-timeout"`
 		WriteTimeout   toml.Duration `toml:"write-timout"`
diff --git a/server/server.go b/server/server.go
index 90061a1b4..aa9919c2b 100644
--- a/server/server.go
+++ b/server/server.go
@@ -191,6 +191,7 @@ func (m *Command) Start() (err error) {
 		m.pgserver.s.WriteTimeout = time.Duration(m.Config.Postgres.WriteTimeout)
 		m.pgserver.s.MaxStartupSize = m.Config.Postgres.MaxStartupSize
 		m.pgserver.s.ConnectionLimit = m.Config.Postgres.ConnectionLimit
+		m.pgserver.s.RequireTLS = m.Config.Postgres.RequireTLS
 		err := m.pgserver.Start(m.Config.Postgres.Addr)
 		if err != nil {
 			return errors.Wrap(err, "starting postgres")