From e2f599a40d6e0e6fea6041d8be224d62c1f019d2 Mon Sep 17 00:00:00 2001 From: Ian Kent Date: Thu, 24 Apr 2014 00:22:50 +0100 Subject: [PATCH] Add in-memory storage, clean up config --- mailhog/{ => config}/config.go | 3 +- mailhog/data/message.go | 7 ++-- mailhog/http/api/v1.go | 60 +++++++++++++++++++++++++--------- mailhog/http/server.go | 11 +++---- mailhog/smtp/session.go | 23 +++++++++---- mailhog/storage/memory.go | 58 ++++++++++++++++++++++++++++++++ mailhog/storage/mongodb.go | 11 ++++--- main.go | 24 +++++++++----- main_test.go | 4 +-- mime_test.go | 4 +-- 10 files changed, 155 insertions(+), 50 deletions(-) rename mailhog/{ => config}/config.go (90%) create mode 100644 mailhog/storage/memory.go diff --git a/mailhog/config.go b/mailhog/config/config.go similarity index 90% rename from mailhog/config.go rename to mailhog/config/config.go index 4aeb093..8e291b1 100644 --- a/mailhog/config.go +++ b/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{} } diff --git a/mailhog/data/message.go b/mailhog/data/message.go index d045a14..a82bdf3 100644 --- a/mailhog/data/message.go +++ b/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["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["Message-ID"] = []string{msg.Id + "@" + hostname} + 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 } diff --git a/mailhog/http/api/v1.go b/mailhog/http/api/v1.go index 93de6f3..e8b4386 100644 --- a/mailhog/http/api/v1.go +++ b/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) - bytes, _ := json.Marshal(messages) - w.Header().Set("Content-Type", "text/json") - w.Write(bytes) + switch apiv1.config.Storage.(type) { + case *storage.MongoDB: + messages, _ := apiv1.config.Storage.(*storage.MongoDB).List(0, 1000) + 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) - bytes, _ := json.Marshal(message) - w.Header().Set("Content-Type", "text/json") - w.Write(bytes) + switch apiv1.config.Storage.(type) { + case *storage.MongoDB: + message, _ := apiv1.config.Storage.(*storage.MongoDB).Load(id) + 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) + } } diff --git a/mailhog/http/server.go b/mailhog/http/server.go index d08c232..428a54e 100644 --- a/mailhog/http/server.go +++ b/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() } diff --git a/mailhog/smtp/session.go b/mailhog/smtp/session.go index 39c1072..86aeacf 100644 --- a/mailhog/smtp/session.go +++ b/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) { - conv := &Session{conn, "", conf, ESTABLISH, &data.SMTPMessage{}, mongo, false} +func StartSession(conn *net.TCPConn, conf *config.Config) { + 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) diff --git a/mailhog/storage/memory.go b/mailhog/storage/memory.go new file mode 100644 index 0000000..dc65e01 --- /dev/null +++ b/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; +} diff --git a/mailhog/storage/mongodb.go b/mailhog/storage/mongodb.go index 2953664..1559870 100644 --- a/mailhog/storage/mongodb.go +++ b/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) diff --git a/main.go b/main.go index 774f155..27f872f 100644 --- a/main.go +++ b/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 mongo *storage.MongoDB +var conf *config.Config 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) } } diff --git a/main_test.go b/main_test.go index 3ac4cd3..4863223 100644 --- a/main_test.go +++ b/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) diff --git a/mime_test.go b/mime_test.go index 0ef8e55..942236a 100644 --- a/mime_test.go +++ b/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)