storage: More refactoring for #69

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

pkg/message/message.go

@@ -2,10 +2,13 @@
 package message
 
 import (
+	"io"
+	"io/ioutil"
 	"net/mail"
 	"time"
 
 	"github.com/jhillyerd/enmime"
+	"github.com/jhillyerd/inbucket/pkg/storage"
 )
 
 // Metadata holds information about a message, but not the content.
@@ -24,3 +27,51 @@ type Message struct {
 	Metadata
 	Envelope *enmime.Envelope
 }
+
+// Delivery is used to add a message to storage.
+type Delivery struct {
+	Meta   Metadata
+	Reader io.Reader
+}
+
+var _ storage.StoreMessage = &Delivery{}
+
+// Mailbox getter.
+func (d *Delivery) Mailbox() string {
+	return d.Meta.Mailbox
+}
+
+// ID getter.
+func (d *Delivery) ID() string {
+	return d.Meta.ID
+}
+
+// From getter.
+func (d *Delivery) From() *mail.Address {
+	return d.Meta.From
+}
+
+// To getter.
+func (d *Delivery) To() []*mail.Address {
+	return d.Meta.To
+}
+
+// Date getter.
+func (d *Delivery) Date() time.Time {
+	return d.Meta.Date
+}
+
+// Subject getter.
+func (d *Delivery) Subject() string {
+	return d.Meta.Subject
+}
+
+// Size getter.
+func (d *Delivery) Size() int64 {
+	return d.Meta.Size
+}
+
+// RawReader contains the raw content of the message.
+func (d *Delivery) RawReader() (io.ReadCloser, error) {
+	return ioutil.NopCloser(d.Reader), nil
+}

pkg/server/smtp/handler.go

@@ -12,7 +12,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/jhillyerd/enmime"
 	"github.com/jhillyerd/inbucket/pkg/log"
+	"github.com/jhillyerd/inbucket/pkg/message"
 	"github.com/jhillyerd/inbucket/pkg/msghub"
 	"github.com/jhillyerd/inbucket/pkg/stringutil"
 )
@@ -442,48 +444,61 @@ func (ss *Session) dataHandler() {
 
 // deliverMessage creates and populates a new Message for the specified recipient
 func (ss *Session) deliverMessage(r recipientDetails, msgBuf [][]byte) (ok bool) {
-	msg, err := ss.server.dataStore.NewMessage(r.localPart)
-	if err != nil {
-		ss.logError("Failed to create message for %q: %s", r.localPart, err)
-		return false
-	}
-
-	// Generate Received header
-	stamp := time.Now().Format(timeStampFormat)
-	recd := fmt.Sprintf("Received: from %s ([%s]) by %s\r\n  for <%s>; %s\r\n",
-		ss.remoteDomain, ss.remoteHost, ss.server.domain, r.address, stamp)
-	if err := msg.Append([]byte(recd)); err != nil {
-		ss.logError("Failed to write received header for %q: %s", r.localPart, err)
-		return false
-	}
-
-	// Append lines from msgBuf
-	for _, line := range msgBuf {
-		if err := msg.Append(line); err != nil {
-			ss.logError("Failed to append to mailbox %v: %v", r.localPart, err)
-			// Should really cleanup the crap on filesystem
-			return false
-		}
-	}
-	if err := msg.Close(); err != nil {
-		ss.logError("Error while closing message for %v: %v", r.localPart, err)
-		return false
-	}
 	name, err := stringutil.ParseMailboxName(r.localPart)
 	if err != nil {
 		// This parse already succeeded when MailboxFor was called, shouldn't fail here.
 		return false
 	}
-
+	buf := bytes.Buffer{}
+	// Generate Received header
+	stamp := time.Now().Format(timeStampFormat)
+	recd := fmt.Sprintf("Received: from %s ([%s]) by %s\r\n  for <%s>; %s\r\n",
+		ss.remoteDomain, ss.remoteHost, ss.server.domain, r.address, stamp)
+	buf.WriteString(recd)
+	// Append lines from msgBuf
+	for _, line := range msgBuf {
+		buf.Write(line)
+	}
+	// TODO replace with something that only reads header?
+	env, err := enmime.ReadEnvelope(bytes.NewReader(buf.Bytes()))
+	if err != nil {
+		ss.logError("Failed to parse message for %q: %v", r.localPart, err)
+		return false
+	}
+	from, err := env.AddressList("From")
+	if err != nil {
+		ss.logError("Failed to get From address: %v", err)
+		return false
+	}
+	to, err := env.AddressList("To")
+	if err != nil {
+		ss.logError("Failed to get To addresses: %v", err)
+		return false
+	}
+	delivery := &message.Delivery{
+		Meta: message.Metadata{
+			Mailbox: name,
+			From:    from[0],
+			To:      to,
+			Date:    time.Now(),
+			Subject: env.GetHeader("Subject"),
+		},
+		Reader: bytes.NewReader(buf.Bytes()),
+	}
+	id, err := ss.server.dataStore.AddMessage(delivery)
+	if err != nil {
+		ss.logError("Failed to store message for %q: %s", r.localPart, err)
+		return false
+	}
 	// Broadcast message information
 	broadcast := msghub.Message{
 		Mailbox: name,
-		ID:      msg.ID(),
-		From:    msg.From().String(),
-		To:      stringutil.StringAddressList(msg.To()),
-		Subject: msg.Subject(),
-		Date:    msg.Date(),
-		Size:    msg.Size(),
+		ID:      id,
+		From:    delivery.From().String(),
+		To:      stringutil.StringAddressList(delivery.To()),
+		Subject: delivery.Subject(),
+		Date:    delivery.Date(),
+		Size:    delivery.Size(),
 	}
 	ss.server.msgHub.Dispatch(broadcast)
 

pkg/server/smtp/handler_test.go

@@ -8,7 +8,6 @@ import (
 
 	"log"
 	"net"
-	"net/mail"
 	"net/textproto"
 	"os"
 	"testing"
@@ -141,18 +140,7 @@ func TestReadyState(t *testing.T) {
 
 // Test commands in MAIL state
 func TestMailState(t *testing.T) {
-	// Setup mock objects
-	mds := &storage.MockDataStore{}
-	msg1 := &storage.MockMessage{}
-	mds.On("NewMessage", "u1").Return(msg1, nil)
-	msg1.On("ID").Return("")
-	msg1.On("From").Return(&mail.Address{})
-	msg1.On("To").Return(make([]*mail.Address, 0))
-	msg1.On("Date").Return(time.Time{})
-	msg1.On("Subject").Return("")
-	msg1.On("Size").Return(0)
-	msg1.On("Close").Return(nil)
-
+	mds := test.NewStore()
 	server, logbuf, teardown := setupSMTPServer(mds)
 	defer teardown()
 
@@ -214,7 +202,7 @@ func TestMailState(t *testing.T) {
 		{"MAIL FROM:<john@gmail.com>", 250},
 		{"RCPT TO:<u1@gmail.com>", 250},
 		{"DATA", 354},
-		{".", 250},
+		{".", 451},
 	}
 	if err := playSession(t, server, script); err != nil {
 		t.Error(err)
@@ -253,18 +241,7 @@ func TestMailState(t *testing.T) {
 
 // Test commands in DATA state
 func TestDataState(t *testing.T) {
-	// Setup mock objects
-	mds := &storage.MockDataStore{}
-	msg1 := &storage.MockMessage{}
-	mds.On("NewMessage", "u1").Return(msg1, nil)
-	msg1.On("ID").Return("")
-	msg1.On("From").Return(&mail.Address{})
-	msg1.On("To").Return(make([]*mail.Address, 0))
-	msg1.On("Date").Return(time.Time{})
-	msg1.On("Subject").Return("")
-	msg1.On("Size").Return(0)
-	msg1.On("Close").Return(nil)
-
+	mds := test.NewStore()
 	server, logbuf, teardown := setupSMTPServer(mds)
 	defer teardown()
 

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)

pkg/test/storage.go

@@ -23,9 +23,11 @@ func NewStore() *StoreStub {
 }
 
 // AddMessage adds a message to the specified mailbox.
-func (s *StoreStub) AddMessage(mailbox string, m storage.StoreMessage) {
-	msgs := s.mailboxes[mailbox]
-	s.mailboxes[mailbox] = append(msgs, m)
+func (s *StoreStub) AddMessage(m storage.StoreMessage) (id string, err error) {
+	mb := m.Mailbox()
+	msgs := s.mailboxes[mb]
+	s.mailboxes[mb] = append(msgs, m)
+	return m.ID(), nil
 }
 
 // GetMessage gets a message by ID from the specified mailbox.
@@ -80,11 +82,6 @@ func (s *StoreStub) VisitMailboxes(f func([]storage.StoreMessage) (cont bool)) e
 	return nil
 }
 
-// NewMessage is temproary until #69 MessageData refactor
-func (s *StoreStub) NewMessage(mailbox string) (storage.StoreMessage, error) {
-	return nil, nil
-}
-
 // LockFor mock function returns a new RWMutex, never errors.
 // TODO(#69) remove
 func (s *StoreStub) LockFor(name string) (*sync.RWMutex, error) {