storage: More refactoring for #69

- impl Store.AddMessage
- file: Use AddMessage() in tests
- smtp: Switch to AddMessage
- storage: Remove NewMessage, Append, Close methods

pkg/storage/file/fmessage.go

@@ -2,16 +2,13 @@ package file
 
 import (
 	"bufio"
-	"fmt"
 	"io"
 	"net/mail"
 	"os"
 	"path/filepath"
 	"time"
 
-	"github.com/jhillyerd/enmime"
 	"github.com/jhillyerd/inbucket/pkg/log"
-	"github.com/jhillyerd/inbucket/pkg/storage"
 )
 
 // Message implements Message and contains a little bit of data about a
@@ -33,7 +30,7 @@ type Message struct {
 
 // newMessage creates a new FileMessage object and sets the Date and ID fields.
 // It will also delete messages over messageCap if configured.
-func (mb *mbox) newMessage() (storage.StoreMessage, error) {
+func (mb *mbox) newMessage() (*Message, error) {
 	// Load index
 	if !mb.indexLoaded {
 		if err := mb.readIndex(); err != nil {
@@ -84,11 +81,6 @@ func (m *Message) Subject() string {
 	return m.Fsubject
 }
 
-// String returns a string in the form: "Subject()" from From()
-func (m *Message) String() string {
-	return fmt.Sprintf("\"%v\" from %v", m.Fsubject, m.Ffrom)
-}
-
 // Size returns the size of the Message on disk in bytes
 func (m *Message) Size() int64 {
 	return m.Fsize
@@ -106,89 +98,3 @@ func (m *Message) RawReader() (reader io.ReadCloser, err error) {
 	}
 	return file, nil
 }
-
-// Append data to a newly opened Message, this will fail on a pre-existing Message and
-// after Close() is called.
-func (m *Message) Append(data []byte) error {
-	// Prevent Appending to a pre-existing Message
-	if !m.writable {
-		return storage.ErrNotWritable
-	}
-	// Open file for writing if we haven't yet
-	if m.writer == nil {
-		// Ensure mailbox directory exists
-		if err := m.mailbox.createDir(); err != nil {
-			return err
-		}
-		file, err := os.Create(m.rawPath())
-		if err != nil {
-			// Set writable false just in case something calls me a million times
-			m.writable = false
-			return err
-		}
-		m.writerFile = file
-		m.writer = bufio.NewWriter(file)
-	}
-	_, err := m.writer.Write(data)
-	m.Fsize += int64(len(data))
-	return err
-}
-
-// Close this Message for writing - no more data may be Appended.  Close() will also
-// trigger the creation of the .gob file.
-func (m *Message) Close() error {
-	// nil out the writer fields so they can't be used
-	writer := m.writer
-	writerFile := m.writerFile
-	m.writer = nil
-	m.writerFile = nil
-
-	if writer != nil {
-		if err := writer.Flush(); err != nil {
-			return err
-		}
-	}
-	if writerFile != nil {
-		if err := writerFile.Close(); err != nil {
-			return err
-		}
-	}
-
-	// Fetch envelope.
-	// TODO should happen outside of datastore.
-	r, err := m.RawReader()
-	if err != nil {
-		return err
-	}
-	env, err := enmime.ReadEnvelope(r)
-	_ = r.Close()
-	if err != nil {
-		return err
-	}
-
-	// Only public fields are stored in gob, hence starting with capital F
-	// Parse From address
-	if address, err := mail.ParseAddress(env.GetHeader("From")); err == nil {
-		m.Ffrom = address
-	} else {
-		m.Ffrom = &mail.Address{Address: env.GetHeader("From")}
-	}
-	m.Fsubject = env.GetHeader("Subject")
-
-	// Turn the To header into a slice
-	if addresses, err := env.AddressList("To"); err == nil {
-		m.Fto = addresses
-	} else {
-		m.Fto = []*mail.Address{{Address: env.GetHeader("To")}}
-	}
-
-	// Refresh the index before adding our message
-	err = m.mailbox.readIndex()
-	if err != nil {
-		return err
-	}
-
-	// Made it this far without errors, add it to the index
-	m.mailbox.messages = append(m.mailbox.messages, m)
-	return m.mailbox.writeIndex()
-}

pkg/storage/file/fstore.go

@@ -74,6 +74,64 @@ func New(cfg config.DataStoreConfig) storage.Store {
 	return &Store{path: path, mailPath: mailPath, messageCap: cfg.MailboxMsgCap}
 }
 
+// AddMessage adds a message to the specified mailbox.
+func (fs *Store) AddMessage(m storage.StoreMessage) (id string, err error) {
+	r, err := m.RawReader()
+	if err != nil {
+		return "", err
+	}
+	mb, err := fs.mbox(m.Mailbox())
+	if err != nil {
+		return "", err
+	}
+	// Create a new message.
+	fm, err := mb.newMessage()
+	if err != nil {
+		return "", err
+	}
+	// Ensure mailbox directory exists.
+	if err := mb.createDir(); err != nil {
+		return "", err
+	}
+	// Write the message content
+	file, err := os.Create(fm.rawPath())
+	if err != nil {
+		return "", err
+	}
+	w := bufio.NewWriter(file)
+	size, err := io.Copy(w, r)
+	if err != nil {
+		// Try to remove the file
+		_ = file.Close()
+		_ = os.Remove(fm.rawPath())
+		return "", err
+	}
+	_ = r.Close()
+	if err := w.Flush(); err != nil {
+		// Try to remove the file
+		_ = file.Close()
+		_ = os.Remove(fm.rawPath())
+		return "", err
+	}
+	if err := file.Close(); err != nil {
+		// Try to remove the file
+		_ = os.Remove(fm.rawPath())
+		return "", err
+	}
+	// Update the index.
+	fm.Fdate = m.Date()
+	fm.Ffrom = m.From()
+	fm.Fsize = size
+	fm.Fsubject = m.Subject()
+	mb.messages = append(mb.messages, fm)
+	if err := mb.writeIndex(); err != nil {
+		// Try to remove the file
+		_ = os.Remove(fm.rawPath())
+		return "", err
+	}
+	return fm.Fid, nil
+}
+
 // GetMessage returns the messages in the named mailbox, or an error.
 func (fs *Store) GetMessage(mailbox, id string) (storage.StoreMessage, error) {
 	mb, err := fs.mbox(mailbox)

pkg/storage/file/fstore_test.go

@@ -6,12 +6,15 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net/mail"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/jhillyerd/inbucket/pkg/config"
+	"github.com/jhillyerd/inbucket/pkg/message"
 	"github.com/jhillyerd/inbucket/pkg/storage"
 	"github.com/stretchr/testify/assert"
 )
@@ -480,32 +483,25 @@ func setupDataStore(cfg config.DataStoreConfig) (*Store, *bytes.Buffer) {
 
 // deliverMessage creates and delivers a message to the specific mailbox, returning
 // the size of the generated message.
-func deliverMessage(ds *Store, mbName string, subject string,
-	date time.Time) (id string, size int64) {
-	// Build fake SMTP message for delivery
-	testMsg := make([]byte, 0, 300)
-	testMsg = append(testMsg, []byte("To: somebody@host\r\n")...)
-	testMsg = append(testMsg, []byte("From: somebodyelse@host\r\n")...)
-	testMsg = append(testMsg, []byte(fmt.Sprintf("Subject: %s\r\n", subject))...)
-	testMsg = append(testMsg, []byte("\r\n")...)
-	testMsg = append(testMsg, []byte("Test Body\r\n")...)
-
-	// Create message object
-	id = generateID(date)
-	msg, err := ds.NewMessage(mbName)
+func deliverMessage(ds *Store, mbName string, subject string, date time.Time) (string, int64) {
+	// Build message for delivery
+	meta := message.Metadata{
+		Mailbox: mbName,
+		To:      []*mail.Address{{Name: "", Address: "somebody@host"}},
+		From:    &mail.Address{Name: "", Address: "somebodyelse@host"},
+		Subject: subject,
+		Date:    date,
+	}
+	testMsg := fmt.Sprintf("To: %s\r\nFrom: %s\r\nSubject: %s\r\n\r\nTest Body\r\n",
+		meta.To[0].Address, meta.From.Address, subject)
+	delivery := &message.Delivery{
+		Meta:   meta,
+		Reader: ioutil.NopCloser(strings.NewReader(testMsg)),
+	}
+	id, err := ds.AddMessage(delivery)
 	if err != nil {
 		panic(err)
 	}
-	fmsg := msg.(*Message)
-	fmsg.Fdate = date
-	fmsg.Fid = id
-	if err = msg.Append(testMsg); err != nil {
-		panic(err)
-	}
-	if err = msg.Close(); err != nil {
-		panic(err)
-	}
-
 	return id, int64(len(testMsg))
 }
 

pkg/storage/retention_test.go

@@ -13,24 +13,18 @@ import (
 func TestDoRetentionScan(t *testing.T) {
 	ds := test.NewStore()
 	// Mockup some different aged messages (num is in hours)
-	new1 := mockMessage(0)
-	new2 := mockMessage(1)
-	new3 := mockMessage(2)
-	old1 := mockMessage(4)
-	old2 := mockMessage(12)
-	old3 := mockMessage(24)
-	ds.AddMessage("mb1", new1)
-	new1.On("Mailbox").Return("mb1")
-	ds.AddMessage("mb1", old1)
-	old1.On("Mailbox").Return("mb1")
-	ds.AddMessage("mb1", old2)
-	old2.On("Mailbox").Return("mb1")
-	ds.AddMessage("mb2", old3)
-	old3.On("Mailbox").Return("mb2")
-	ds.AddMessage("mb2", new2)
-	new2.On("Mailbox").Return("mb2")
-	ds.AddMessage("mb3", new3)
-	new3.On("Mailbox").Return("mb3")
+	new1 := mockMessage("mb1", 0)
+	new2 := mockMessage("mb2", 1)
+	new3 := mockMessage("mb3", 2)
+	old1 := mockMessage("mb1", 4)
+	old2 := mockMessage("mb1", 12)
+	old3 := mockMessage("mb2", 24)
+	ds.AddMessage(new1)
+	ds.AddMessage(old1)
+	ds.AddMessage(old2)
+	ds.AddMessage(old3)
+	ds.AddMessage(new2)
+	ds.AddMessage(new3)
 	// Test 4 hour retention
 	cfg := config.DataStoreConfig{
 		RetentionMinutes: 239,
@@ -56,8 +50,9 @@ func TestDoRetentionScan(t *testing.T) {
 }
 
 // Make a MockMessage of a specific age
-func mockMessage(ageHours int) *storage.MockMessage {
+func mockMessage(mailbox string, ageHours int) *storage.MockMessage {
 	msg := &storage.MockMessage{}
+	msg.On("Mailbox").Return(mailbox)
 	msg.On("ID").Return(fmt.Sprintf("MSG[age=%vh]", ageHours))
 	msg.On("Date").Return(time.Now().Add(time.Duration(ageHours*-1) * time.Hour))
 	msg.On("Delete").Return(nil)

pkg/storage/storage.go

@@ -19,6 +19,8 @@ var (
 
 // Store is the interface Inbucket uses to interact with storage implementations.
 type Store interface {
+	// AddMessage stores the message, message ID and Size will be ignored.
+	AddMessage(message StoreMessage) (id string, err error)
 	GetMessage(mailbox, id string) (StoreMessage, error)
 	GetMessages(mailbox string) ([]StoreMessage, error)
 	PurgeMessages(mailbox string) error
@@ -39,8 +41,5 @@ type StoreMessage interface {
 	Date() time.Time
 	Subject() string
 	RawReader() (reader io.ReadCloser, err error)
-	Append(data []byte) error
-	Close() error
-	String() string
 	Size() int64
 }

pkg/storage/testing.go

@@ -15,6 +15,12 @@ type MockDataStore struct {
 	mock.Mock
 }
 
+// AddMessage mock function
+func (m *MockDataStore) AddMessage(message StoreMessage) (string, error) {
+	args := m.Called(message)
+	return args.String(0), args.Error(1)
+}
+
 // GetMessage mock function
 func (m *MockDataStore) GetMessage(name, id string) (StoreMessage, error) {
 	args := m.Called(name, id)