Add in-memory storage, clean up config

mailhog/config/config.go

@@ -1,4 +1,4 @@
-package mailhog
+package config
 
 func DefaultConfig() *Config {
 	return &Config{
@@ -18,4 +18,5 @@ type Config struct {
 	MongoUri string
 	MongoDb string
 	MongoColl string
+	Storage interface{}
 }

mailhog/data/message.go

@@ -6,7 +6,6 @@ import (
 	"time"
 	"regexp"
     "labix.org/v2/mgo/bson"
-    "github.com/ian-kent/MailHog/mailhog"
 )
 
 type Messages []Message
@@ -44,7 +43,7 @@ type MIMEBody struct {
 	Parts []*Content
 }
 
-func ParseSMTPMessage(c *mailhog.Config, m *SMTPMessage) *Message {
+func ParseSMTPMessage(m *SMTPMessage, hostname string) *Message {
 	arr := make([]*Path, 0)
 	for _, path := range m.To {
 		arr = append(arr, PathFromString(path))
@@ -62,8 +61,8 @@ func ParseSMTPMessage(c *mailhog.Config, m *SMTPMessage) *Message {
 		msg.MIME = msg.Content.ParseMIMEBody()
 	}
 	
-	msg.Content.Headers["Message-ID"] = []string{msg.Id + "@" + c.Hostname}
+	msg.Content.Headers["Message-ID"] = []string{msg.Id + "@" + hostname}
-	msg.Content.Headers["Received"] = []string{"from " + m.Helo + " by " + c.Hostname + " (Go-MailHog)\r\n          id " + msg.Id + "@" + c.Hostname + "; " + time.Now().Format(time.RFC1123Z)}
+	msg.Content.Headers["Received"] = []string{"from " + m.Helo + " by " + hostname + " (Go-MailHog)\r\n          id " + msg.Id + "@" + hostname + "; " + time.Now().Format(time.RFC1123Z)}
 	msg.Content.Headers["Return-Path"] = []string{"<" + m.From + ">"}
 	return msg
 }

mailhog/http/api/v1.go

@@ -5,25 +5,23 @@ import (
 	"encoding/json"
 	"net/http"
 	"regexp"
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/storage"
 	"github.com/ian-kent/MailHog/mailhog/http/handler"
 )
 
 type APIv1 struct {
-	config *mailhog.Config
+	config *config.Config
 	exitChannel chan int
 	server *http.Server
-	mongo *storage.MongoDB
 }
 
-func CreateAPIv1(exitCh chan int, conf *mailhog.Config, server *http.Server, mongo *storage.MongoDB) *APIv1 {
+func CreateAPIv1(exitCh chan int, conf *config.Config, server *http.Server) *APIv1 {
 	log.Println("Creating API v1")
 	apiv1 := &APIv1{
 		config: conf,
 		exitChannel: exitCh,
 		server: server,
-		mongo: mongo,
 	}
 
 	server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/api/v1/messages/?$"), apiv1.messages)
@@ -38,10 +36,21 @@ func (apiv1 *APIv1) messages(w http.ResponseWriter, r *http.Request, route *hand
 	log.Println("[APIv1] GET /api/v1/messages")
 
 	// TODO start, limit
-	messages, _ := apiv1.mongo.List(0, 1000)
+	switch apiv1.config.Storage.(type) {
-	bytes, _ := json.Marshal(messages)
+		case *storage.MongoDB:
-	w.Header().Set("Content-Type", "text/json")
+			messages, _ := apiv1.config.Storage.(*storage.MongoDB).List(0, 1000)
-	w.Write(bytes)
+			bytes, _ := json.Marshal(messages)
+			w.Header().Set("Content-Type", "text/json")
+			w.Write(bytes)
+		case *storage.Memory:
+			messages, _ := apiv1.config.Storage.(*storage.Memory).List(0, 1000)
+			bytes, _ := json.Marshal(messages)
+			w.Header().Set("Content-Type", "text/json")
+			w.Write(bytes)
+		default:
+			w.Header().Set("Content-Type", "text/json")
+			w.Write([]byte("[]"))
+	}
 }
 
 func (apiv1 *APIv1) message(w http.ResponseWriter, r *http.Request, route *handler.Route) {
@@ -49,17 +58,33 @@ func (apiv1 *APIv1) message(w http.ResponseWriter, r *http.Request, route *handl
 	id := match[1]
 	log.Printf("[APIv1] GET /api/v1/messages/%s\n", id)
 
-	message, _ := apiv1.mongo.Load(id)
+	switch apiv1.config.Storage.(type) {
-	bytes, _ := json.Marshal(message)
+		case *storage.MongoDB:
-	w.Header().Set("Content-Type", "text/json")
+			message, _ := apiv1.config.Storage.(*storage.MongoDB).Load(id)
-	w.Write(bytes)
+			bytes, _ := json.Marshal(message)
+			w.Header().Set("Content-Type", "text/json")
+			w.Write(bytes)
+		case *storage.Memory:
+			message, _ := apiv1.config.Storage.(*storage.Memory).Load(id)
+			bytes, _ := json.Marshal(message)
+			w.Header().Set("Content-Type", "text/json")
+			w.Write(bytes)
+		default:
+			w.Header().Set("Content-Type", "text/json")
+			w.Write([]byte("[]"))
+	}
 }
 
 func (apiv1 *APIv1) delete_all(w http.ResponseWriter, r *http.Request, route *handler.Route) {
 	log.Println("[APIv1] POST /api/v1/messages/delete")
 
 	w.Header().Set("Content-Type", "text/json")
-	apiv1.mongo.DeleteAll()
+	switch apiv1.config.Storage.(type) {
+		case *storage.MongoDB:
+			apiv1.config.Storage.(*storage.MongoDB).DeleteAll()
+		case *storage.Memory:
+			apiv1.config.Storage.(*storage.Memory).DeleteAll()
+	}
 }
 
 func (apiv1 *APIv1) delete_one(w http.ResponseWriter, r *http.Request, route *handler.Route) {
@@ -68,5 +93,10 @@ func (apiv1 *APIv1) delete_one(w http.ResponseWriter, r *http.Request, route *ha
 	log.Printf("[APIv1] POST /api/v1/messages/%s/delete\n", id)
 
 	w.Header().Set("Content-Type", "text/json")
-	apiv1.mongo.DeleteOne(id)
+	switch apiv1.config.Storage.(type) {
+		case *storage.MongoDB:
+			apiv1.config.Storage.(*storage.MongoDB).DeleteOne(id)
+		case *storage.Memory:
+			apiv1.config.Storage.(*storage.Memory).DeleteOne(id)
+	}
 }

mailhog/http/server.go

@@ -3,9 +3,8 @@ package http
 import (
 	"regexp"
 	"net/http"
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/templates"
-	"github.com/ian-kent/MailHog/mailhog/storage"
 	"github.com/ian-kent/MailHog/mailhog/templates/images"
 	"github.com/ian-kent/MailHog/mailhog/templates/js"
 	"github.com/ian-kent/MailHog/mailhog/http/api"
@@ -13,7 +12,7 @@ import (
 )
 
 var exitChannel chan int
-var config *mailhog.Config
+var cfg *config.Config
 
 func web_exit(w http.ResponseWriter, r *http.Request, route *handler.Route) {
 	web_headers(w)
@@ -44,9 +43,9 @@ func web_headers(w http.ResponseWriter) {
 	w.Header().Set("Content-Type", "text/html")
 }
 
-func Start(exitCh chan int, conf *mailhog.Config, mongo *storage.MongoDB) {
+func Start(exitCh chan int, conf *config.Config) {
 	exitChannel = exitCh
-	config = conf
+	cfg = conf
 
 	server := &http.Server{
 		Addr: conf.HTTPBindAddr,
@@ -58,7 +57,7 @@ func Start(exitCh chan int, conf *mailhog.Config, mongo *storage.MongoDB) {
 	server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/images/hog.png$"), web_imgcontroller)
 	server.Handler.(*handler.RegexpHandler).HandleFunc(regexp.MustCompile("^/$"), web_index)
 
-	api.CreateAPIv1(exitCh, conf, server, mongo)
+	api.CreateAPIv1(exitCh, conf, server)
 
 	server.ListenAndServe()
 }

mailhog/smtp/session.go

@@ -7,7 +7,7 @@ import (
 	"net"
 	"strings"
 	"regexp"
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/storage"
 	"github.com/ian-kent/MailHog/mailhog/data"
 )
@@ -15,10 +15,9 @@ import (
 type Session struct {
 	conn *net.TCPConn
 	line string
-	conf *mailhog.Config
+	conf *config.Config
 	state int
 	message *data.SMTPMessage
-	mongo *storage.MongoDB
 	isTLS bool
 }
 
@@ -34,8 +33,8 @@ const (
 
 // TODO replace ".." lines with . in data
 
-func StartSession(conn *net.TCPConn, conf *mailhog.Config, mongo *storage.MongoDB) {
+func StartSession(conn *net.TCPConn, conf *config.Config) {
-	conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}, mongo, false}
+	conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}, false}
 	conv.log("Starting session")
 	conv.Write("220", conv.conf.Hostname + " ESMTP Go-MailHog")
 	conv.Read()
@@ -82,7 +81,19 @@ func (c *Session) Parse() {
 				c.log("Got EOF, storing message and switching to MAIL state")
 				//c.log("Full message data: %s", c.message.Data)
 				c.message.Data = strings.TrimSuffix(c.message.Data, "\r\n.\r\n")
-				id, err := c.mongo.Store(c.message)
+				var id string
+				var err error
+				switch c.conf.Storage.(type) {
+					case *storage.MongoDB:
+						c.log("Storing message using MongoDB")
+						id, err = c.conf.Storage.(*storage.MongoDB).Store(c.message)
+					case *storage.Memory:
+						c.log("Storing message using Memory")
+						id, err = c.conf.Storage.(*storage.Memory).Store(c.message)
+					default:
+						c.log("Unknown storage type")
+						// TODO send error reply
+				}
 				c.state = MAIL
 				if err != nil {
 					c.log("Error storing message: %s", err)

mailhog/storage/memory.go

@@ -0,0 +1,58 @@
+package storage
+
+import (
+	"github.com/ian-kent/MailHog/mailhog/config"
+    "github.com/ian-kent/MailHog/mailhog/data"
+)
+
+type Memory struct {
+	Config *config.Config
+	Messages map[string]*data.Message
+	MessageIndex []string
+	MessageRIndex map[string]int
+}
+
+func CreateMemory(c *config.Config) *Memory {
+	return &Memory{
+		Config: c,
+		Messages: make(map[string]*data.Message, 0),
+		MessageIndex: make([]string, 0),
+		MessageRIndex: make(map[string]int, 0),
+	}
+}
+
+func (memory *Memory) Store(m *data.SMTPMessage) (string, error) {
+	msg := data.ParseSMTPMessage(m, memory.Config.Hostname)
+	memory.Messages[msg.Id] = msg
+	memory.MessageIndex = append(memory.MessageIndex, msg.Id)
+	memory.MessageRIndex[msg.Id] = len(memory.MessageIndex)
+	return msg.Id, nil
+}
+
+func (memory *Memory) List(start int, limit int) ([]*data.Message, error) {
+	if limit > len(memory.MessageIndex) { limit = len(memory.MessageIndex) }
+	messages := make([]*data.Message, 0)
+	for _, m := range memory.MessageIndex[start:limit] {
+		messages = append(messages, memory.Messages[m])
+	}
+	return messages, nil;
+}
+
+func (memory *Memory) DeleteOne(id string) error {
+	index := memory.MessageRIndex[id];
+	delete(memory.Messages, id)
+	memory.MessageIndex = append(memory.MessageIndex[:index], memory.MessageIndex[index+1:]...)
+	delete(memory.MessageRIndex, id)
+	return nil
+}
+
+func (memory *Memory) DeleteAll() error {
+	memory.Messages = make(map[string]*data.Message, 0)
+	memory.MessageIndex = make([]string, 0)
+	memory.MessageRIndex = make(map[string]int, 0)
+	return nil
+}
+
+func (memory *Memory) Load(id string) (*data.Message, error) {
+	return memory.Messages[id], nil;
+}

mailhog/storage/mongodb.go

@@ -5,19 +5,20 @@ import (
 	"labix.org/v2/mgo"
 	"labix.org/v2/mgo/bson"
     "github.com/ian-kent/MailHog/mailhog/data"
-    "github.com/ian-kent/MailHog/mailhog"
+    "github.com/ian-kent/MailHog/mailhog/config"
 )
 
 type MongoDB struct {
 	Session *mgo.Session
-	Config *mailhog.Config
+	Config *config.Config
 	Collection *mgo.Collection
 }
 
-func CreateMongoDB(c *mailhog.Config) *MongoDB {
+func CreateMongoDB(c *config.Config) *MongoDB {
+	log.Printf("Connecting to MongoDB: %s\n", c.MongoUri)
 	session, err := mgo.Dial(c.MongoUri)
 	if(err != nil) {
-		log.Fatalf("Error connecting to MongoDB: %s", err)
+		log.Printf("Error connecting to MongoDB: %s", err)
 		return nil
 	}
 	return &MongoDB{
@@ -28,7 +29,7 @@ func CreateMongoDB(c *mailhog.Config) *MongoDB {
 }
 
 func (mongo *MongoDB) Store(m *data.SMTPMessage) (string, error) {
-	msg := data.ParseSMTPMessage(mongo.Config, m)
+	msg := data.ParseSMTPMessage(m, mongo.Config.Hostname)
 	err := mongo.Collection.Insert(msg)
 	if err != nil {
 		log.Printf("Error inserting message: %s", err)

main.go

@@ -2,7 +2,7 @@ package main
 
 import (
 	"flag"
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/http"
 	"github.com/ian-kent/MailHog/mailhog/smtp"
 	"github.com/ian-kent/MailHog/mailhog/storage"
@@ -11,11 +11,10 @@ import (
 	"os"
 )
 
-var conf *mailhog.Config
+var conf *config.Config
-var mongo *storage.MongoDB
 var exitCh chan int
 
-func config() {
+func configure() {
 	var smtpbindaddr, httpbindaddr, hostname, mongouri, mongodb, mongocoll string
 
 	flag.StringVar(&smtpbindaddr, "smtpbindaddr", "0.0.0.0:1025", "SMTP bind interface and port, e.g. 0.0.0.0:1025 or just :1025")
@@ -27,7 +26,7 @@ func config() {
 
 	flag.Parse()
 
-	conf = &mailhog.Config{
+	conf = &config.Config{
 		SMTPBindAddr: smtpbindaddr,
 		HTTPBindAddr: httpbindaddr,
 		Hostname:     hostname,
@@ -36,11 +35,18 @@ func config() {
 		MongoColl:    mongocoll,
 	}
 
-	mongo = storage.CreateMongoDB(conf)
+	s := storage.CreateMongoDB(conf)
+	if s == nil {
+		log.Println("MongoDB storage unavailable, using in-memory storage")
+		conf.Storage = storage.CreateMemory(conf)
+	} else {
+		log.Println("Connected to MongoDB")
+		conf.Storage = s
+	}
 }
 
 func main() {
-	config()
+	configure()
 
 	exitCh = make(chan int)
 	go web_listen()
@@ -57,7 +63,7 @@ func main() {
 
 func web_listen() {
 	log.Printf("[HTTP] Binding to address: %s\n", conf.HTTPBindAddr)
-	http.Start(exitCh, conf, mongo)
+	http.Start(exitCh, conf)
 }
 
 func smtp_listen() *net.TCPListener {
@@ -76,6 +82,6 @@ func smtp_listen() *net.TCPListener {
 		}
 		defer conn.Close()
 
-		go smtp.StartSession(conn.(*net.TCPConn), conf, mongo)
+		go smtp.StartSession(conn.(*net.TCPConn), conf)
 	}
 }

main_test.go

@@ -1,7 +1,7 @@
 package main
 
 import (
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/storage"
 	"github.com/stretchr/testify/assert"
 	"net"
@@ -94,7 +94,7 @@ func TestBasicHappyPath(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, string(buf[0:n]), "221 Bye\n")
 
-	s := storage.CreateMongoDB(mailhog.DefaultConfig())
+	s := storage.CreateMongoDB(config.DefaultConfig())
 	message, err := s.Load(match[1])
 	assert.Nil(t, err)
 	assert.NotNil(t, message)

mime_test.go

@@ -1,7 +1,7 @@
 package main
 
 import (
-	"github.com/ian-kent/MailHog/mailhog"
+	"github.com/ian-kent/MailHog/mailhog/config"
 	"github.com/ian-kent/MailHog/mailhog/storage"
 	"github.com/stretchr/testify/assert"
 	"net"
@@ -112,7 +112,7 @@ func TestBasicMIMEHappyPath(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, string(buf[0:n]), "221 Bye\n")
 
-	s := storage.CreateMongoDB(mailhog.DefaultConfig())
+	s := storage.CreateMongoDB(config.DefaultConfig())
 	message, err := s.Load(match[1])
 	assert.Nil(t, err)
 	assert.NotNil(t, message)