From ae71996632dca2438a9ab906ad90bb0a4c3557c7 Mon Sep 17 00:00:00 2001
From: Andres Montalban <andres@labelbox.com>
Date: Tue, 12 Oct 2021 17:07:04 -0300
Subject: [PATCH] Allow configuration of sslmode for DB connections

---
 config.example.yml | 10 +++++++++-
 config/config.go   |  5 ++++-
 core/core.go       |  6 +++---
 core/db.go         | 12 +++++++++---
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/config.example.yml b/config.example.yml
index 8c1b20a..22a19bd 100644
--- a/config.example.yml
+++ b/config.example.yml
@@ -6,6 +6,10 @@ Monitor:
   Password: 123456
   DBName: pg_auto_failover
 
+  # sslmode for connecting to the monitor service (Defaults to 'disable')
+  # Reference: https://www.postgresql.org/docs/current/libpq-ssl.html#LIBPQ-SSL-SSLMODE-STATEMENTS
+  SSLMode: disable
+
 # A list of coordinator nodes that you want to be monitored for changes in their worker nodes.
 Coordinators:
   - DBName: postgres
@@ -15,6 +19,10 @@ Coordinators:
     # Formation that the node can be found with in monitor.
     Formation: default
 
+    # sslmode for connecting to the coordinator nodes (Defaults to 'disable')
+    # Reference: https://www.postgresql.org/docs/current/libpq-ssl.html#LIBPQ-SSL-SSLMODE-STATEMENTS
+    SSLMode: disable
+
 # Service settings
 Settings:
   # Check interval for changes (in ms).
@@ -31,4 +39,4 @@ API:
 
   # The secret key that will be used to authorize requests.
   # All requests require this string in their header as SECRET.
-  Secret: SECRET_STRING
\ No newline at end of file
+  Secret: SECRET_STRING
diff --git a/config/config.go b/config/config.go
index a29cfd8..895f38d 100644
--- a/config/config.go
+++ b/config/config.go
@@ -2,11 +2,12 @@ package config
 
 import (
 	"fmt"
-	"github.com/spf13/viper"
 	"log"
 	"os"
 	"path/filepath"
 	"strings"
+
+	"github.com/spf13/viper"
 )
 
 // ServiceConfig represents the config data for the application.
@@ -17,12 +18,14 @@ type ServiceConfig struct {
 		User     string
 		Password string
 		DBName   string
+		SSLMode  string `default:"disable"`
 	}
 	Coordinators []struct {
 		Formation string
 		Username  string
 		Password  string
 		DBName    string
+		SSLMode   string `default:"disable"`
 	}
 	Settings struct {
 		CheckInterval int
diff --git a/core/core.go b/core/core.go
index 3283b67..bd90b1d 100644
--- a/core/core.go
+++ b/core/core.go
@@ -144,19 +144,19 @@ func (d *database) connect(coordinatorNode *Coordinator) error {
 	if d.db == nil {
 		d.host = coordinatorNode.Host
 		d.port = coordinatorNode.Port
-		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port)
+		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port, d.sslmode)
 		return err
 	}
 	if d.host != coordinatorNode.Host || d.port != coordinatorNode.Port {
 		logger.CoordinatorChanged(coordinatorNode.Host, d.host, d.dbname, coordinatorNode.Port, d.port)
 		d.host = coordinatorNode.Host
 		d.port = coordinatorNode.Port
-		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port)
+		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port, d.sslmode)
 		return err
 	}
 	if d.db.Ping() != nil {
 		logger.CoordinatorConnectionLost(d.host, d.dbname, d.username, d.port)
-		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port)
+		d.db, err = openDBConnection(d.host, d.username, d.dbname, d.password, d.port, d.sslmode)
 		return err
 	}
 	return nil
diff --git a/core/db.go b/core/db.go
index af0ce08..a02e002 100644
--- a/core/db.go
+++ b/core/db.go
@@ -2,6 +2,7 @@ package core
 
 import (
 	"fmt"
+
 	"github.com/Navid2zp/citus-failover/config"
 	"github.com/jmoiron/sqlx"
 )
@@ -13,11 +14,13 @@ type database struct {
 	username  string
 	password  string
 	dbname    string
+	sslmode   string `default:"disable"`
 	db        *sqlx.DB
 }
 
 // databases holds the list of all databases to be monitored
 var databases []*database
+
 // monitorDB is database instance for monitor
 var monitorDB *sqlx.DB
 
@@ -28,12 +31,13 @@ func openMonitoringConnection() {
 	monitorDB, err = sqlx.Connect(
 		"postgres",
 		fmt.Sprintf(
-			"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+			"host=%s port=%d user=%s password=%s dbname=%s sslmode=%s",
 			config.Config.Monitor.Host,
 			config.Config.Monitor.Port,
 			config.Config.Monitor.User,
 			config.Config.Monitor.Password,
 			config.Config.Monitor.DBName,
+			config.Config.Monitor.SSLMode,
 		),
 	)
 	if err != nil {
@@ -51,6 +55,7 @@ func setupDatabases() {
 			username:  db.Username,
 			password:  db.Password,
 			dbname:    db.DBName,
+			sslmode:   db.SSLMode,
 			db:        nil,
 		}
 		databases = append(databases, &coordinator)
@@ -58,16 +63,17 @@ func setupDatabases() {
 }
 
 // openDBConnection opens a database connection.
-func openDBConnection(host, username, dbname, password string, port int) (*sqlx.DB, error) {
+func openDBConnection(host, username, dbname, password string, port int, sslmode string) (*sqlx.DB, error) {
 	return sqlx.Connect(
 		"postgres",
 		fmt.Sprintf(
-			"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+			"host=%s port=%d user=%s password=%s dbname=%s sslmode=%s",
 			host,
 			port,
 			username,
 			password,
 			dbname,
+			sslmode,
 		),
 	)
 }