From 9d308546b73830d689585fcdae2add26d3626d78 Mon Sep 17 00:00:00 2001 From: Alexandre Cesaro Date: Wed, 8 Jul 2015 19:53:11 +0200 Subject: [PATCH] Made File struct more flexible The File can now be streamed to the SMTP server directly without being buffered into memory first. Fixes #4 --- message.go | 55 +++++++++++++++++++++++++---------------------- message_test.go | 57 +++++++++++++++++++++++++++---------------------- writeto.go | 43 ++++++++++++++++++++++++------------- 3 files changed, 88 insertions(+), 67 deletions(-) diff --git a/message.go b/message.go index a075186..c71fcee 100644 --- a/message.go +++ b/message.go @@ -8,6 +8,7 @@ import ( "io" "io/ioutil" "mime" + "os" "path/filepath" "sync" "time" @@ -220,38 +221,40 @@ func (msg *Message) AddAlternativeWriter(contentType string, f func(io.Writer) e // A File represents a file that can be attached or embedded in an email. type File struct { - Name string - MimeType string - Content []byte - ContentID string + // Name represents the base name of the file. If the file is attached to the + // message it is the name of the attachment. + Name string + // Header represents the MIME header of the message part that contains the + // file content. + Header map[string][]string + // Copier is a function run when the message is sent. It should copy the + // content of the file to w. + Copier func(w io.Writer) error } -// OpenFile opens a file on disk to create a gomail.File. -func OpenFile(filename string) (*File, error) { - content, err := readFile(filename) - if err != nil { - return nil, err - } - - f := CreateFile(filepath.Base(filename), content) - - return f, nil -} - -// CreateFile creates a gomail.File from the given name and content. -func CreateFile(name string, content []byte) *File { - mimeType := mime.TypeByExtension(filepath.Ext(name)) - if mimeType == "" { - mimeType = "application/octet-stream" - } - +// NewFile creates a File from the given filename. +func NewFile(filename string) *File { return &File{ - Name: name, - MimeType: mimeType, - Content: content, + Name: filepath.Base(filename), + Header: make(map[string][]string), + Copier: func(w io.Writer) error { + h, err := os.Open(filename) + if err != nil { + return err + } + if _, err := io.Copy(w, h); err != nil { + h.Close() + return err + } + return h.Close() + }, } } +func (f *File) setHeader(field string, value ...string) { + f.Header[field] = value +} + // Attach attaches the files to the email. func (msg *Message) Attach(f ...*File) { if msg.attachments == nil { diff --git a/message_test.go b/message_test.go index 7cb7ea2..db851ec 100644 --- a/message_test.go +++ b/message_test.go @@ -198,11 +198,7 @@ func TestAttachmentOnly(t *testing.T) { msg := NewMessage() msg.SetHeader("From", "from@example.com") msg.SetHeader("To", "to@example.com") - f, err := OpenFile("/tmp/test.pdf") - if err != nil { - panic(err) - } - msg.Attach(f) + msg.Attach(testFile("/tmp/test.pdf")) want := &message{ from: "from@example.com", @@ -224,7 +220,7 @@ func TestAttachment(t *testing.T) { msg.SetHeader("From", "from@example.com") msg.SetHeader("To", "to@example.com") msg.SetBody("text/plain", "Test") - msg.Attach(CreateFile("test.pdf", []byte("Content"))) + msg.Attach(testFile("/tmp/test.pdf")) want := &message{ from: "from@example.com", @@ -243,7 +239,7 @@ func TestAttachment(t *testing.T) { "Content-Disposition: attachment; filename=\"test.pdf\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.pdf")) + "\r\n" + "--_BOUNDARY_1_--\r\n", } @@ -254,8 +250,8 @@ func TestAttachmentsOnly(t *testing.T) { msg := NewMessage() msg.SetHeader("From", "from@example.com") msg.SetHeader("To", "to@example.com") - msg.Attach(CreateFile("test.pdf", []byte("Content 1"))) - msg.Attach(CreateFile("test.zip", []byte("Content 2"))) + msg.Attach(testFile("/tmp/test.pdf")) + msg.Attach(testFile("/tmp/test.zip")) want := &message{ from: "from@example.com", @@ -269,13 +265,13 @@ func TestAttachmentsOnly(t *testing.T) { "Content-Disposition: attachment; filename=\"test.pdf\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 1")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.pdf")) + "\r\n" + "--_BOUNDARY_1_\r\n" + "Content-Type: application/zip; name=\"test.zip\"\r\n" + "Content-Disposition: attachment; filename=\"test.zip\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 2")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.zip")) + "\r\n" + "--_BOUNDARY_1_--\r\n", } @@ -287,8 +283,8 @@ func TestAttachments(t *testing.T) { msg.SetHeader("From", "from@example.com") msg.SetHeader("To", "to@example.com") msg.SetBody("text/plain", "Test") - msg.Attach(CreateFile("test.pdf", []byte("Content 1"))) - msg.Attach(CreateFile("test.zip", []byte("Content 2"))) + msg.Attach(testFile("/tmp/test.pdf")) + msg.Attach(testFile("/tmp/test.zip")) want := &message{ from: "from@example.com", @@ -307,13 +303,13 @@ func TestAttachments(t *testing.T) { "Content-Disposition: attachment; filename=\"test.pdf\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 1")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.pdf")) + "\r\n" + "--_BOUNDARY_1_\r\n" + "Content-Type: application/zip; name=\"test.zip\"\r\n" + "Content-Disposition: attachment; filename=\"test.zip\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 2")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.zip")) + "\r\n" + "--_BOUNDARY_1_--\r\n", } @@ -324,10 +320,10 @@ func TestEmbedded(t *testing.T) { msg := NewMessage() msg.SetHeader("From", "from@example.com") msg.SetHeader("To", "to@example.com") - f := CreateFile("image1.jpg", []byte("Content 1")) - f.ContentID = "test-content-id" + f := testFile("image1.jpg") + f.Header["Content-ID"] = []string{""} msg.Embed(f) - msg.Embed(CreateFile("image2.jpg", []byte("Content 2"))) + msg.Embed(testFile("image2.jpg")) msg.SetBody("text/plain", "Test") want := &message{ @@ -348,14 +344,14 @@ func TestEmbedded(t *testing.T) { "Content-ID: \r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 1")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of image1.jpg")) + "\r\n" + "--_BOUNDARY_1_\r\n" + "Content-Type: image/jpeg; name=\"image2.jpg\"\r\n" + "Content-Disposition: inline; filename=\"image2.jpg\"\r\n" + "Content-ID: \r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 2")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of image2.jpg")) + "\r\n" + "--_BOUNDARY_1_--\r\n", } @@ -368,8 +364,8 @@ func TestFullMessage(t *testing.T) { msg.SetHeader("To", "to@example.com") msg.SetBody("text/plain", "¡Hola, señor!") msg.AddAlternative("text/html", "¡Hola, señor!") - msg.Attach(CreateFile("test.pdf", []byte("Content 1"))) - msg.Embed(CreateFile("image.jpg", []byte("Content 2"))) + msg.Attach(testFile("test.pdf")) + msg.Embed(testFile("image.jpg")) want := &message{ from: "from@example.com", @@ -402,7 +398,7 @@ func TestFullMessage(t *testing.T) { "Content-ID: \r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 2")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of image.jpg")) + "\r\n" + "--_BOUNDARY_2_--\r\n" + "\r\n" + "--_BOUNDARY_1_\r\n" + @@ -410,7 +406,7 @@ func TestFullMessage(t *testing.T) { "Content-Disposition: attachment; filename=\"test.pdf\"\r\n" + "Content-Transfer-Encoding: base64\r\n" + "\r\n" + - base64.StdEncoding.EncodeToString([]byte("Content 1")) + "\r\n" + + base64.StdEncoding.EncodeToString([]byte("Content of test.pdf")) + "\r\n" + "--_BOUNDARY_1_--\r\n", } @@ -582,6 +578,15 @@ func getBoundaries(t *testing.T, count int, msg string) []string { var boundaryRegExp = regexp.MustCompile("boundary=(\\w+)") +func testFile(name string) *File { + f := NewFile(name) + f.Copier = func(w io.Writer) error { + _, err := w.Write([]byte("Content of " + filepath.Base(f.Name))) + return err + } + return f +} + func BenchmarkFull(b *testing.B) { emptyFunc := func(from string, to []string, msg io.WriterTo) error { return nil @@ -598,8 +603,8 @@ func BenchmarkFull(b *testing.B) { }) msg.SetBody("text/plain", "¡Hola, señor!") msg.AddAlternative("text/html", "

¡Hola, señor!

") - msg.Attach(CreateFile("benchmark.txt", []byte("Benchmark"))) - msg.Embed(CreateFile("benchmark.jpg", []byte("Benchmark"))) + msg.Attach(testFile("benchmark.txt")) + msg.Embed(testFile("benchmark.jpg")) if err := Send(SendFunc(emptyFunc), msg); err != nil { panic(err) diff --git a/writeto.go b/writeto.go index 15c200a..7a584bb 100644 --- a/writeto.go +++ b/writeto.go @@ -4,8 +4,10 @@ import ( "encoding/base64" "errors" "io" + "mime" "mime/multipart" "mime/quotedprintable" + "path/filepath" "time" ) @@ -109,24 +111,35 @@ func (w *messageWriter) closeMultipart() { func (w *messageWriter) addFiles(files []*File, isAttachment bool) { for _, f := range files { - h := make(map[string][]string) - h["Content-Type"] = []string{f.MimeType + "; name=\"" + f.Name + "\""} - h["Content-Transfer-Encoding"] = []string{string(Base64)} - if isAttachment { - h["Content-Disposition"] = []string{"attachment; filename=\"" + f.Name + "\""} - } else { - h["Content-Disposition"] = []string{"inline; filename=\"" + f.Name + "\""} - if f.ContentID != "" { - h["Content-ID"] = []string{"<" + f.ContentID + ">"} + if _, ok := f.Header["Content-Type"]; !ok { + mediaType := mime.TypeByExtension(filepath.Ext(f.Name)) + if mediaType == "" { + mediaType = "application/octet-stream" + } + f.setHeader("Content-Type", mediaType+`; name="`+f.Name+`"`) + } + + if _, ok := f.Header["Content-Transfer-Encoding"]; !ok { + f.setHeader("Content-Transfer-Encoding", string(Base64)) + } + + if _, ok := f.Header["Content-Disposition"]; !ok { + var disp string + if isAttachment { + disp = "attachment" } else { - h["Content-ID"] = []string{"<" + f.Name + ">"} + disp = "inline" + } + f.setHeader("Content-Disposition", disp+`; filename="`+f.Name+`"`) + } + + if !isAttachment { + if _, ok := f.Header["Content-ID"]; !ok { + f.setHeader("Content-ID", "<"+f.Name+">") } } - w.writeHeaders(h) - w.writeBody(func(w io.Writer) error { - _, err := w.Write(f.Content) - return err - }, Base64) + w.writeHeaders(f.Header) + w.writeBody(f.Copier, Base64) } }