Merge pull request #3 from Selly-Modules/add-tls-options

Add tls options

.gitignore

@@ -14,4 +14,5 @@
 # Dependency directories (remove the comment below to include it)
 # vendor/
 
-.idea
+.idea
+*.pem

base64.go

@@ -0,0 +1,12 @@
+package mongodb
+
+import "encoding/base64"
+
+func base64DecodeToBytes(text string) []byte {
+	s, _ := base64.StdEncoding.DecodeString(text)
+	return s
+}
+
+func base64DecodeToString(text string) string {
+	return string(base64DecodeToBytes(text))
+}

mongodb.go

@@ -7,37 +7,97 @@ import (
 	"github.com/logrusorgru/aurora"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/readpref"
 )
 
+// Config ...
+type Config struct {
+	Host   string
+	DBName string
+
+	TLS        *ConnectTLSOpts
+	Standalone *ConnectStandaloneOpts
+}
+
+// ConnectTLSOpts ...
+type ConnectTLSOpts struct {
+	ReplSet             string
+	CaFile              string
+	CertKeyFile         string
+	CertKeyFilePassword string
+}
+
+// ConnectStandaloneOpts ...
+type ConnectStandaloneOpts struct {
+	AuthMechanism string
+	AuthSource    string
+	Username      string
+	Password      string
+}
+
 var db *mongo.Database
 
 // Connect to mongo server
-func Connect(host, user, password, dbName, mechanism, source string) (*mongo.Database, error) {
+func Connect(cfg Config) (*mongo.Database, error) {
+	if cfg.TLS != nil && cfg.TLS.ReplSet != "" {
+		return connectWithTLS(cfg)
+	}
 	connectOptions := options.ClientOptions{}
+	opts := cfg.Standalone
 	// Set auth if existed
-	if user != "" && password != "" {
+	if opts.Username != "" && opts.Password != "" {
 		connectOptions.Auth = &options.Credential{
-			AuthMechanism: mechanism,
+			AuthMechanism: opts.AuthMechanism,
-			AuthSource:    source,
+			AuthSource:    opts.AuthSource,
-			Username:      user,
+			Username:      opts.Username,
-			Password:      password,
+			Password:      opts.Password,
 		}
 	}
 
 	// Connect
-	client, err := mongo.Connect(context.Background(), connectOptions.ApplyURI(host))
+	client, err := mongo.Connect(context.Background(), connectOptions.ApplyURI(cfg.Host))
 	if err != nil {
-		fmt.Println("Error when connect to MongoDB database", host, err)
+		fmt.Println("Error when connect to MongoDB database", cfg.Host, err)
 		return nil, err
 	}
 
-	fmt.Println(aurora.Green("*** CONNECTED TO MONGODB: " + host))
+	fmt.Println(aurora.Green("*** CONNECTED TO MONGODB: " + cfg.Host + " --- DB: " + cfg.DBName))
 
 	// Set data
-	db = client.Database(dbName)
+	db = client.Database(cfg.DBName)
 	return db, nil
 }
 
+func connectWithTLS(cfg Config) (*mongo.Database, error) {
+	ctx := context.Background()
+	opts := cfg.TLS
+
+	caFile, err := initFileFromBase64String("ca.pem", opts.CaFile)
+	if err != nil {
+		return nil, err
+	}
+	certFile, err := initFileFromBase64String("cert.pem", opts.CertKeyFile)
+	if err != nil {
+		return nil, err
+	}
+	pwd := base64DecodeToString(opts.CertKeyFilePassword)
+	s := "%s/?tls=true&tlsCAFile=./%s&tlsCertificateKeyFile=./%s&tlsCertificateKeyFilePassword=%s&authMechanism=MONGODB-X509"
+	uri := fmt.Sprintf(s, cfg.Host, caFile.Name(), certFile.Name(), pwd)
+	readPref := readpref.SecondaryPreferred()
+	clientOpts := options.Client().SetReadPreference(readPref).SetReplicaSet(opts.ReplSet).ApplyURI(uri)
+	client, err := mongo.Connect(ctx, clientOpts)
+	if err != nil {
+		return nil, err
+	}
+	if err := client.Ping(ctx, readpref.SecondaryPreferred()); err != nil {
+		return nil, err
+	}
+	db = client.Database(cfg.DBName)
+
+	fmt.Println(aurora.Green("*** CONNECTED TO MONGODB: " + cfg.Host + " --- DB: " + cfg.DBName))
+	return db, err
+}
+
 // GetInstance ...
 func GetInstance() *mongo.Database {
 	return db

util.go

@@ -0,0 +1,21 @@
+package mongodb
+
+import (
+	"fmt"
+	"os"
+)
+
+func initFileFromBase64String(filename, value string) (*os.File, error) {
+	f, err := os.Create(filename)
+	if err != nil {
+		fmt.Println("mongodb.initFileFromBase64String - err: ", err)
+		return nil, err
+	}
+	b := base64DecodeToBytes(value)
+	if _, err := f.Write(b); err != nil {
+		fmt.Println("mongodb.initFileFromBase64String - write file err: ", err)
+		return nil, err
+	}
+	f.Sync()
+	return f, nil
+}