From b13e2154adca8457eb902ae24ceb8e1f3c16a5e8 Mon Sep 17 00:00:00 2001
From: Patrick Easters <peasters@redhat.com>
Date: Wed, 11 Dec 2019 21:45:19 -0500
Subject: [PATCH] Add option to not fallback to in-memory storage from mongo

---
 config/config.go | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/config/config.go b/config/config.go
index 59cb706..a2d0227 100644
--- a/config/config.go
+++ b/config/config.go
@@ -15,18 +15,19 @@ import (
 // DefaultConfig is the default config
 func DefaultConfig() *Config {
 	return &Config{
-		SMTPBindAddr: "0.0.0.0:1025",
-		APIBindAddr:  "0.0.0.0:8025",
-		Hostname:     "mailhog.example",
-		MongoURI:     "127.0.0.1:27017",
-		MongoDb:      "mailhog",
-		MongoColl:    "messages",
-		MaildirPath:  "",
-		StorageType:  "memory",
-		CORSOrigin:   "",
-		WebPath:      "",
-		MessageChan:  make(chan *data.Message),
-		OutgoingSMTP: make(map[string]*OutgoingSMTP),
+		SMTPBindAddr:  "0.0.0.0:1025",
+		APIBindAddr:   "0.0.0.0:8025",
+		Hostname:      "mailhog.example",
+		MongoURI:      "127.0.0.1:27017",
+		MongoDb:       "mailhog",
+		MongoColl:     "messages",
+		MongoFallback: true,
+		MaildirPath:   "",
+		StorageType:   "memory",
+		CORSOrigin:    "",
+		WebPath:       "",
+		MessageChan:   make(chan *data.Message),
+		OutgoingSMTP:  make(map[string]*OutgoingSMTP),
 	}
 }
 
@@ -38,6 +39,7 @@ type Config struct {
 	MongoURI         string
 	MongoDb          string
 	MongoColl        string
+	MongoFallback    bool
 	StorageType      string
 	CORSOrigin       string
 	MaildirPath      string
@@ -78,6 +80,9 @@ func Configure() *Config {
 		log.Println("Using MongoDB message storage")
 		s := storage.CreateMongoDB(cfg.MongoURI, cfg.MongoDb, cfg.MongoColl)
 		if s == nil {
+			if !cfg.MongoFallback {
+				log.Fatal("MongoDB storage unavailable")
+			}
 			log.Println("MongoDB storage unavailable, reverting to in-memory storage")
 			cfg.Storage = storage.CreateInMemory()
 		} else {
@@ -124,6 +129,7 @@ func RegisterFlags() {
 	flag.StringVar(&cfg.MongoURI, "mongo-uri", envconf.FromEnvP("MH_MONGO_URI", "127.0.0.1:27017").(string), "MongoDB URI, e.g. 127.0.0.1:27017")
 	flag.StringVar(&cfg.MongoDb, "mongo-db", envconf.FromEnvP("MH_MONGO_DB", "mailhog").(string), "MongoDB database, e.g. mailhog")
 	flag.StringVar(&cfg.MongoColl, "mongo-coll", envconf.FromEnvP("MH_MONGO_COLLECTION", "messages").(string), "MongoDB collection, e.g. messages")
+	flag.BoolVar(&cfg.MongoFallback, "mongo-fallback", envconf.FromEnvP("MH_MONGO_FALLBACK", true).(bool), "Fallback to in-memory storage if MongoDB is unavailable")
 	flag.StringVar(&cfg.CORSOrigin, "cors-origin", envconf.FromEnvP("MH_CORS_ORIGIN", "").(string), "CORS Access-Control-Allow-Origin header for API endpoints")
 	flag.StringVar(&cfg.MaildirPath, "maildir-path", envconf.FromEnvP("MH_MAILDIR_PATH", "").(string), "Maildir path (if storage type is 'maildir')")
 	flag.BoolVar(&cfg.InviteJim, "invite-jim", envconf.FromEnvP("MH_INVITE_JIM", false).(bool), "Decide whether to invite Jim (beware, he causes trouble)")