Added TLS configuration for acra-rollback/acra-rotate (#623)

Added TLS configuration for acra-rotate/acra-rollback
cossacklabs · Jan 20, 2023 · 153f3cd · 153f3cd
1 parent 9c642fc
commit 153f3cd
Show file tree

Hide file tree

Showing 12 changed files with 451 additions and 73 deletions.
diff --git a/benchmarks/common/db.go b/benchmarks/common/db.go
@@ -17,10 +17,11 @@ package common
 
 import (
 	"database/sql"
+	"os"
+
 	// import driver for connect function
 	_ "github.com/jackc/pgx/v4/stdlib"
 	"github.com/sirupsen/logrus"
-	"os"
 )
 
 func connect(connectionString string) *sql.DB {

diff --git a/cmd/acra-rollback/acra-rollback.go b/cmd/acra-rollback/acra-rollback.go
@@ -25,13 +25,20 @@ package main
 import (
 	"bufio"
 	"container/list"
+	"crypto/tls"
 	"database/sql"
 	"flag"
 	"fmt"
+	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
 
+	"github.com/go-sql-driver/mysql"
+	"github.com/jackc/pgx/v4"
+	"github.com/jackc/pgx/v4/stdlib"
+	log "github.com/sirupsen/logrus"
+
 	"github.com/cossacklabs/acra/acrastruct"
 	"github.com/cossacklabs/acra/cmd"
 	"github.com/cossacklabs/acra/keystore"
@@ -40,10 +47,8 @@ import (
 	keystoreV2 "github.com/cossacklabs/acra/keystore/v2/keystore"
 	filesystemV2 "github.com/cossacklabs/acra/keystore/v2/keystore/filesystem"
 	"github.com/cossacklabs/acra/logging"
+	"github.com/cossacklabs/acra/network"
 	"github.com/cossacklabs/acra/utils"
-	_ "github.com/go-sql-driver/mysql"
-	_ "github.com/lib/pq"
-	log "github.com/sirupsen/logrus"
 )
 
 // Constants used by AcraRollback
@@ -158,15 +163,18 @@ func (ex *WriteToFileExecutor) Close() {
 func main() {
 	keysDir := flag.String("keys_dir", keystore.DefaultKeyDirShort, "Folder from which the keys will be loaded")
 	clientID := flag.String("client_id", "", "Client ID should be name of file with private key")
-	connectionString := flag.String("connection_string", "", "Connection string for db")
+	connectionString := flag.String("connection_string", "", "Connection string for DB PostgreSQL(postgresql://{user}:{password}@{host}:{port}/{dbname}?sslmode={sslmode}), MySQL ({user}:{password}@tcp({host}:{port})/{dbname})")
 	sqlSelect := flag.String("select", "", "Query to fetch data for decryption")
 	sqlInsert := flag.String("insert", "", "Query for insert decrypted data with placeholders (pg: $n, mysql: ?)")
 	outputFile := flag.String("output_file", "decrypted.sql", "File for store inserts queries")
 	execute := flag.Bool("execute", false, "Execute inserts")
 	escapeFormat := flag.Bool("escape", false, "Escape bytea format")
 	useMysql := flag.Bool("mysql_enable", false, "Handle MySQL connections")
 	usePostgresql := flag.Bool("postgresql_enable", false, "Handle Postgresql connections")
+	dbTLSEnabled := flag.Bool("tls_database_enabled", false, "Enable TLS for DB")
 
+	network.RegisterTLSArgsForService(flag.CommandLine, true, "", network.DatabaseNameConstructorFunc())
+	network.RegisterTLSBaseArgs(flag.CommandLine)
 	keyloader.RegisterKeyStoreStrategyParameters()
 	logging.SetLogLevel(logging.LogVerbose)
 
@@ -192,21 +200,13 @@ func main() {
 		os.Exit(1)
 	}
 
-	dbDriverName := "postgres"
-	if *useMysql {
-		// https://github.com/ziutek/mymysql
-		//dbDriverName = "mymysql"
-		// https://github.com/go-sql-driver/mysql/
-		dbDriverName = "mysql"
-	}
-
-	cmd.ValidateClientID(*clientID)
-
 	if *connectionString == "" {
 		log.Errorln("Connection_string arg is missing")
 		os.Exit(1)
 	}
 
+	cmd.ValidateClientID(*clientID)
+
 	if *sqlSelect == "" {
 		log.Errorln("Sql_select arg is missing")
 		os.Exit(1)
@@ -215,6 +215,7 @@ func main() {
 		log.Errorln("Sql_insert arg is missing")
 		os.Exit(1)
 	}
+
 	if *outputFile == "" && !*execute {
 		log.Errorln("Output_file missing or execute flag")
 		os.Exit(1)
@@ -227,11 +228,61 @@ func main() {
 		keystorage = openKeyStoreV1(*keysDir)
 	}
 
-	db, err := sql.Open(dbDriverName, *connectionString)
-	if err != nil {
-		log.WithError(err).Errorln("Can't connect to db")
-		os.Exit(1)
+	var dbTLSConfig *tls.Config
+	if *dbTLSEnabled {
+		host, err := network.GetDriverConnectionStringHost(*connectionString, *useMysql)
+		if err != nil {
+			log.WithError(err).Errorln("Failed to get DB host from connection URL")
+			os.Exit(1)
+		}
+
+		dbTLSConfig, err = network.NewTLSConfigByName(flag.CommandLine, "", host, network.DatabaseNameConstructorFunc())
+		if err != nil {
+			log.WithError(err).WithField(logging.FieldKeyEventCode, logging.EventCodeErrorTransportConfiguration).
+				Errorln("Configuration error: can't create database TLS config")
+			os.Exit(1)
+		}
 	}
+
+	var db *sql.DB
+	if *useMysql {
+		if dbTLSConfig != nil {
+			connectionURL, err := url.Parse(*connectionString)
+			if err != nil {
+				log.WithError(err).Errorln("Failed to parse DB connection string")
+				os.Exit(1)
+			}
+
+			if err := mysql.RegisterTLSConfig("custom", dbTLSConfig); err != nil {
+				log.WithError(err).Errorln("Failed to register TLS config")
+				os.Exit(1)
+			}
+
+			connectioQueryParams := connectionURL.Query()
+			connectioQueryParams.Set("tls", "custom")
+			connectionURL.RawQuery = connectioQueryParams.Encode()
+			*connectionString = connectionURL.String()
+		}
+
+		db, err = sql.Open("mysql", *connectionString)
+		if err != nil {
+			log.WithError(err).Errorln("Can't connect to db")
+			os.Exit(1)
+		}
+	} else {
+		config, err := pgx.ParseConfig(*connectionString)
+		if err != nil {
+			log.WithError(err).Errorln("Can't parse config ")
+			os.Exit(1)
+		}
+
+		if dbTLSConfig != nil {
+			config.TLSConfig = dbTLSConfig
+		}
+
+		db = stdlib.OpenDB(*config)
+	}
+
 	defer db.Close()
 	err = db.Ping()
 	if err != nil {

diff --git a/cmd/acra-rotate/acra-rotate.go b/cmd/acra-rotate/acra-rotate.go
@@ -19,10 +19,16 @@ limitations under the License.
 package main
 
 import (
+	"crypto/tls"
 	"database/sql"
 	"flag"
+	"net/url"
 	"os"
 
+	"github.com/go-sql-driver/mysql"
+	"github.com/jackc/pgx/v4"
+	"github.com/jackc/pgx/v4/stdlib"
+
 	"github.com/cossacklabs/acra/cmd"
 	"github.com/cossacklabs/acra/crypto"
 	"github.com/cossacklabs/acra/keystore"
@@ -32,9 +38,8 @@ import (
 	filesystemV2 "github.com/cossacklabs/acra/keystore/v2/keystore/filesystem"
 	filesystemBackendV2 "github.com/cossacklabs/acra/keystore/v2/keystore/filesystem/backend"
 	"github.com/cossacklabs/acra/logging"
+	"github.com/cossacklabs/acra/network"
 	"github.com/cossacklabs/acra/utils"
-	_ "github.com/go-sql-driver/mysql"
-	_ "github.com/lib/pq"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -51,12 +56,16 @@ func main() {
 	fileMapConfig := flag.String("file_map_config", "", "Path to file with map of <ClientId>: <FilePaths> in json format {\"client_id1\": [\"filepath1\", \"filepath2\"], \"client_id2\": [\"filepath1\", \"filepath2\"]}")
 	sqlSelect := flag.String("sql_select", "", "Select query with ? as placeholders where last columns in result must be ClientId and AcraStruct. Other columns will be passed into insert/update query into placeholders")
 	sqlUpdate := flag.String("sql_update", "", "Insert/Update query with ? as placeholder where into first will be placed rotated AcraStruct")
-	connectionString := flag.String("db_connection_string", "", "Connection string to db")
+	connectionString := flag.String("db_connection_string", "", "Connection string for DB PostgreSQL(postgresql://{user}:{password}@{host}:{port}/{dbname}?sslmode={sslmode}), MySQL ({user}:{password}@tcp({host}:{port})/{dbname})")
 	useMysql := flag.Bool("mysql_enable", false, "Handle MySQL connections")
 	_ = flag.Bool("postgresql_enable", false, "Handle Postgresql connections")
 	dryRun := flag.Bool("dry-run", false, "perform rotation without saving rotated AcraStructs and keys")
+	dbTLSEnabled := flag.Bool("tls_database_enabled", false, "Enable TLS for DB")
+
 	logging.SetLogLevel(logging.LogVerbose)
 
+	network.RegisterTLSArgsForService(flag.CommandLine, true, "", network.DatabaseNameConstructorFunc())
+	network.RegisterTLSBaseArgs(flag.CommandLine)
 	cmd.RegisterRedisKeystoreParameters()
 	keyloader.RegisterKeyStoreStrategyParameters()
 
@@ -88,14 +97,58 @@ func main() {
 			log.Errorln("sql_select and sql_update must be set both")
 			os.Exit(1)
 		}
+
+		var dbTLSConfig *tls.Config
+		if *dbTLSEnabled {
+			host, err := network.GetDriverConnectionStringHost(*connectionString, *useMysql)
+			if err != nil {
+				log.WithError(err).Errorln("Failed to get DB host from connection URL")
+				os.Exit(1)
+			}
+
+			dbTLSConfig, err = network.NewTLSConfigByName(flag.CommandLine, "", host, network.DatabaseNameConstructorFunc())
+			if err != nil {
+				log.WithError(err).WithField(logging.FieldKeyEventCode, logging.EventCodeErrorTransportConfiguration).
+					Errorln("Configuration error: can't create database TLS config")
+				os.Exit(1)
+			}
+		}
+
 		var db *sql.DB
 		var encoder utils.BinaryEncoder
 		if *useMysql {
+			if dbTLSConfig != nil {
+				connectionURL, err := url.Parse(*connectionString)
+				if err != nil {
+					log.WithError(err).Errorln("Failed to parse DB connection string")
+					os.Exit(1)
+				}
+
+				if err := mysql.RegisterTLSConfig("custom", dbTLSConfig); err != nil {
+					log.WithError(err).Errorln("Failed to register TLS config")
+					os.Exit(1)
+				}
+
+				connectioQueryParams := connectionURL.Query()
+				connectioQueryParams.Set("tls", "custom")
+				connectionURL.RawQuery = connectioQueryParams.Encode()
+				*connectionString = connectionURL.String()
+			}
+
 			db, err = sql.Open("mysql", *connectionString)
 			encoder = &utils.HexEncoder{}
 		} else {
-			db, err = sql.Open("postgres", *connectionString)
+			config, err := pgx.ParseConfig(*connectionString)
+			if err != nil {
+				log.WithError(err).Errorln("Can't parse config")
+				os.Exit(1)
+			}
+
+			if dbTLSConfig != nil {
+				config.TLSConfig = dbTLSConfig
+			}
 
+			db = stdlib.OpenDB(*config)
 			encoder = &utils.MysqlEncoder{}
 		}
 

diff --git a/cmd/utils.go b/cmd/utils.go
@@ -19,28 +19,30 @@ package cmd
 
 import (
 	"context"
+	"encoding/base64"
 	"errors"
 	flag_ "flag"
 	"fmt"
-	"github.com/cossacklabs/acra/logging"
 	"io"
 	"io/ioutil"
+	"math/rand"
 	"net"
 	"os"
 	"os/signal"
 	"path/filepath"
 	"reflect"
+	"strconv"
+	"strings"
 	"sync"
+	"time"
+
+	"github.com/cossacklabs/acra/logging"
 
-	"encoding/base64"
-	"github.com/cossacklabs/acra/keystore"
-	"github.com/cossacklabs/acra/utils"
 	log "github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v2"
-	"math/rand"
-	"strconv"
-	"strings"
-	"time"
+
+	"github.com/cossacklabs/acra/keystore"
+	"github.com/cossacklabs/acra/utils"
 )
 
 var (