feat: enter and exit handlers

pkg/message/parser.go

@@ -31,7 +31,7 @@ import (
 	"github.com/jaytaylor/html2text"
 )
 
-func Parse(r io.Reader, key, keyName string) (m *pmapi.Message, mime, plain string, atts []io.Reader, err error) {
+func Parse(r io.Reader, key, keyName string) (m *pmapi.Message, mimeMessage, plainBody string, attReaders []io.Reader, err error) {
 	p, err := parser.New(r)
 	if err != nil {
 		return
@@ -43,27 +43,29 @@ func Parse(r io.Reader, key, keyName string) (m *pmapi.Message, mime, plain stri
 		return
 	}
 
-	if m.Attachments, atts, err = collectAttachments(p); err != nil {
+	atts, attReaders, err := collectAttachments(p)
+	if err != nil {
 		return
 	}
+	m.Attachments = atts
 
-	var isHTML bool
+	richBody, plainBody, err := collectBodyParts(p)
-
+	if err != nil {
-	if m.Body, plain, isHTML, err = collectBodyParts(p); err != nil {
 		return
 	}
+	m.Body = richBody
 
-	if isHTML {
+	mimeType, err := determineMIMEType(p)
-		m.MIMEType = "text/html"
+	if err != nil {
-	} else {
+		return
-		m.MIMEType = "text/plain"
 	}
+	m.MIMEType = mimeType
 
 	if key != "" {
 		attachPublicKey(p.Root(), key, keyName)
 	}
 
-	if mime, err = writeMIMEMessage(p); err != nil {
+	if mimeMessage, err = writeMIMEMessage(p); err != nil {
 		return
 	}
 
@@ -71,9 +73,10 @@ func Parse(r io.Reader, key, keyName string) (m *pmapi.Message, mime, plain stri
 }
 
 func collectAttachments(p *parser.Parser) (atts []*pmapi.Attachment, data []io.Reader, err error) {
-	w := p.
+	w := p.NewWalker()
-		NewWalker().
+
-		WithContentDispositionHandler("attachment", func(p *parser.Part, _ parser.PartHandler) (err error) {
+	w.RegisterContentDispositionHandler("attachment").
+		OnEnter(func(p *parser.Part, _ parser.PartHandlerFunc) (err error) {
 			att, err := parseAttachment(p.Header)
 			if err != nil {
 				return
@@ -92,34 +95,80 @@ func collectAttachments(p *parser.Parser) (atts []*pmapi.Attachment, data []io.R
 	return
 }
 
-func collectBodyParts(p *parser.Parser) (body, plain string, isHTML bool, err error) {
+// collectBodyParts returns a richtext body (used for normal sending)
-	var parts, plainParts []string
+// and a plaintext body (used for sending to recipients that prefer plaintext).
+func collectBodyParts(p *parser.Parser) (richBody, plainBody string, err error) {
+	var richParts, plainParts []string
 
-	w := p.
+	w := p.NewWalker()
-		NewWalker().
+
-		WithContentTypeHandler("text/plain", func(p *parser.Part) (err error) {
+	w.RegisterContentTypeHandler("text/plain").
-			parts = append(parts, string(p.Body))
+		OnEnter(func(p *parser.Part) error {
 			plainParts = append(plainParts, string(p.Body))
-			return
-		}).
-		WithContentTypeHandler("text/html", func(p *parser.Part) (err error) {
-			parts = append(parts, string(p.Body))
-			isHTML = true
 
-			text, err := html2text.FromString(string(p.Body))
+			if !isAlternative(p) {
-			if err != nil {
+				richParts = append(richParts, string(p.Body))
-				text = string(p.Body)
 			}
-			plainParts = append(plainParts, text)
 
-			return
+			return nil
+		})
+
+	w.RegisterContentTypeHandler("text/html").
+		OnEnter(func(p *parser.Part) error {
+			richParts = append(richParts, string(p.Body))
+
+			if !isAlternative(p) {
+				plain, htmlErr := html2text.FromString(string(p.Body))
+				if htmlErr != nil {
+					plain = string(p.Body)
+				}
+				plainParts = append(plainParts, plain)
+			}
+
+			return nil
 		})
 
 	if err = w.Walk(); err != nil {
 		return
 	}
 
-	return strings.Join(parts, "\r\n"), strings.Join(plainParts, "\r\n"), isHTML, nil
+	return strings.Join(richParts, "\r\n"), strings.Join(plainParts, "\r\n"), nil
+}
+
+func isAlternative(p *parser.Part) bool {
+	parent := p.Parent()
+	if parent == nil {
+		return false
+	}
+
+	t, _, err := parent.Header.ContentType()
+	if err != nil {
+		return false
+	}
+
+	return t == "multipart/alternative"
+}
+
+func determineMIMEType(p *parser.Parser) (string, error) {
+	w := p.NewWalker()
+
+	var isHTML bool
+
+	w.RegisterContentTypeHandler("text/html").
+		OnEnter(func(p *parser.Part) (err error) {
+			isHTML = true
+			return
+		})
+
+	if err := w.Walk(); err != nil {
+		return "", err
+	}
+
+	if isHTML {
+		return "text/html", nil
+	}
+
+	return "text/plain", nil
 }
 
 func writeMIMEMessage(p *parser.Parser) (mime string, err error) {

pkg/message/parser/handlers.go

@@ -0,0 +1,72 @@
+package parser
+
+type PartHandlerFunc func(*Part) error
+type DispHandlerFunc func(*Part, PartHandlerFunc) error
+
+type PartHandler struct {
+	enter, exit PartHandlerFunc
+}
+
+func NewPartHandler() *PartHandler {
+	return &PartHandler{
+		enter: partNoop,
+		exit:  partNoop,
+	}
+}
+
+func (h *PartHandler) OnEnter(fn PartHandlerFunc) *PartHandler {
+	h.enter = fn
+	return h
+}
+
+func (h *PartHandler) OnExit(fn PartHandlerFunc) *PartHandler {
+	h.exit = fn
+	return h
+}
+
+func (h *PartHandler) handleEnter(_ *Walker, p *Part) error {
+	return h.enter(p)
+}
+
+func (h *PartHandler) handleExit(_ *Walker, p *Part) error {
+	return h.exit(p)
+}
+
+type DispHandler struct {
+	enter, exit DispHandlerFunc
+}
+
+func NewDispHandler() *DispHandler {
+	return &DispHandler{
+		enter: dispNoop,
+		exit:  dispNoop,
+	}
+}
+
+func (h *DispHandler) OnEnter(fn DispHandlerFunc) *DispHandler {
+	h.enter = fn
+	return h
+}
+
+func (h *DispHandler) OnExit(fn DispHandlerFunc) *DispHandler {
+	h.exit = fn
+	return h
+}
+
+func (h *DispHandler) handleEnter(w *Walker, p *Part) error {
+	// NOTE: This is hacky -- is there a better solution?
+	return h.enter(p, func(p *Part) error {
+		return w.getTypeHandler(p).handleEnter(w, p)
+	})
+}
+
+func (h *DispHandler) handleExit(w *Walker, p *Part) error {
+	// NOTE: This is hacky -- is there a better solution?
+	return h.exit(p, func(p *Part) error {
+		return w.getTypeHandler(p).handleExit(w, p)
+	})
+}
+
+func partNoop(*Part) error { return nil }
+
+func dispNoop(*Part, PartHandlerFunc) error { return nil }

pkg/message/parser/parser.go

@@ -63,7 +63,7 @@ func (p *Parser) parse(r io.Reader) (err error) {
 }
 
 func (p *Parser) enter() {
-	p.stack = append(p.stack, &Part{})
+	p.stack = append(p.stack, &Part{parent: p.top()})
 }
 
 func (p *Parser) exit() {
@@ -79,6 +79,10 @@ func (p *Parser) exit() {
 }
 
 func (p *Parser) top() *Part {
+	if len(p.stack) == 0 {
+		return nil
+	}
+
 	return p.stack[len(p.stack)-1]
 }
 

pkg/message/parser/part.go

@@ -9,6 +9,7 @@ import (
 type Part struct {
 	Header   message.Header
 	Body     []byte
+	parent   *Part
 	children []*Part
 }
 
@@ -24,12 +25,34 @@ func (p *Part) Parts() (n int) {
 	return len(p.children)
 }
 
+func (p *Part) Parent() *Part {
+	return p.parent
+}
+
+func (p *Part) Siblings() []*Part {
+	if p.parent == nil {
+		return nil
+	}
+
+	siblings := []*Part{}
+
+	for _, sibling := range p.parent.children {
+		if sibling != p {
+			siblings = append(siblings, sibling)
+		}
+	}
+
+	return siblings
+}
+
 func (p *Part) AddChild(child *Part) {
 	p.children = append(p.children, child)
 }
 
 func (p *Part) visit(w *Walker) (err error) {
-	if err = p.handle(w); err != nil {
+	hdl := p.getHandler(w)
+
+	if err = hdl.handleEnter(w, p); err != nil {
 		return
 	}
 
@@ -39,45 +62,15 @@ func (p *Part) visit(w *Walker) (err error) {
 		}
 	}
 
-	return
+	return hdl.handleExit(w, p)
 }
 
-func (p *Part) getTypeHandler(w *Walker) (hdl PartHandler) {
+func (p *Part) getHandler(w *Walker) handler {
-	t, _, err := p.Header.ContentType()
+	if dispHandler := w.getDispHandler(p); dispHandler != nil {
-	if err != nil {
+		return dispHandler
-		return
 	}
 
-	return w.typeHandlers[t]
+	return w.getTypeHandler(p)
-}
-
-func (p *Part) getDispHandler(w *Walker) (hdl DispHandler) {
-	t, _, err := p.Header.ContentDisposition()
-	if err != nil {
-		return
-	}
-
-	return w.dispHandlers[t]
-}
-
-func (p *Part) handle(w *Walker) (err error) {
-	typeHandler := p.getTypeHandler(w)
-	dispHandler := p.getDispHandler(w)
-	defaultHandler := w.defaultHandler
-
-	switch {
-	case dispHandler != nil && typeHandler != nil:
-		return dispHandler(p, typeHandler)
-
-	case dispHandler != nil && typeHandler == nil:
-		return dispHandler(p, defaultHandler)
-
-	case dispHandler == nil && typeHandler != nil:
-		return typeHandler(p)
-
-	default:
-		return defaultHandler(p)
-	}
 }
 
 func (p *Part) write(writer *message.Writer, w *Writer) (err error) {

pkg/message/parser/walker.go

@@ -3,20 +3,22 @@ package parser
 type Walker struct {
 	root *Part
 
-	defaultHandler PartHandler
+	defaultHandler handler
-	typeHandlers   map[string]PartHandler
+	typeHandlers   map[string]handler
-	dispHandlers   map[string]DispHandler
+	dispHandlers   map[string]handler
 }
 
-type PartHandler func(*Part) error
+type handler interface {
-type DispHandler func(*Part, PartHandler) error
+	handleEnter(*Walker, *Part) error
+	handleExit(*Walker, *Part) error
+}
 
 func newWalker(root *Part) *Walker {
 	return &Walker{
 		root:           root,
-		defaultHandler: func(*Part) (err error) { return },
+		defaultHandler: NewPartHandler(),
-		typeHandlers:   make(map[string]PartHandler),
+		typeHandlers:   make(map[string]handler),
-		dispHandlers:   make(map[string]DispHandler),
+		dispHandlers:   make(map[string]handler),
 	}
 }
 
@@ -24,16 +26,49 @@ func (w *Walker) Walk() (err error) {
 	return w.root.visit(w)
 }
 
-func (w *Walker) WithDefaultHandler(handler PartHandler) *Walker {
+func (w *Walker) WithDefaultHandler(handler handler) *Walker {
 	w.defaultHandler = handler
 	return w
 }
-func (w *Walker) WithContentTypeHandler(contType string, handler PartHandler) *Walker {
+func (w *Walker) RegisterContentTypeHandler(contType string) *PartHandler {
-	w.typeHandlers[contType] = handler
+	hdl := NewPartHandler()
-	return w
+
+	w.typeHandlers[contType] = hdl
+
+	return hdl
 }
 
-func (w *Walker) WithContentDispositionHandler(contDisp string, handler DispHandler) *Walker {
+func (w *Walker) RegisterContentDispositionHandler(contDisp string) *DispHandler {
-	w.dispHandlers[contDisp] = handler
+	hdl := NewDispHandler()
-	return w
+
+	w.dispHandlers[contDisp] = hdl
+
+	return hdl
+}
+
+// getTypeHandler returns the appropriate PartHandler to handle the given part.
+// If no specialised handler exists, it returns the default handler.
+func (w *Walker) getTypeHandler(p *Part) handler {
+	t, _, err := p.Header.ContentType()
+	if err != nil {
+		return w.defaultHandler
+	}
+
+	hdl, ok := w.typeHandlers[t]
+	if !ok {
+		return w.defaultHandler
+	}
+
+	return hdl
+}
+
+// getDispHandler returns the appropriate DispHandler to handle the given part.
+// If no specialised handler exists, it returns nil.
+func (w *Walker) getDispHandler(p *Part) handler {
+	t, _, err := p.Header.ContentDisposition()
+	if err != nil {
+		return nil
+	}
+
+	return w.dispHandlers[t]
 }

pkg/message/parser/walker_test.go

@@ -13,12 +13,12 @@ func TestWalker(t *testing.T) {
 
 	walker := p.
 		NewWalker().
-		WithDefaultHandler(func(p *Part) (err error) {
+		WithDefaultHandler(NewPartHandler().OnEnter(func(p *Part) (err error) {
 			if p.Body != nil {
 				allBodies = append(allBodies, p.Body)
 			}
 			return
-		})
+		}))
 
 	assert.NoError(t, walker.Walk())
 	assert.ElementsMatch(t, [][]byte{
@@ -32,9 +32,11 @@ func TestWalkerTypeHandler(t *testing.T) {
 
 	html := [][]byte{}
 
-	walker := p.
+	walker := p.NewWalker()
-		NewWalker().
+
-		WithContentTypeHandler("text/html", func(p *Part) (err error) {
+	walker.
+		RegisterContentTypeHandler("text/html").
+		OnEnter(func(p *Part) (err error) {
 			html = append(html, p.Body)
 			return
 		})
@@ -50,9 +52,11 @@ func TestWalkerDispositionHandler(t *testing.T) {
 
 	attachments := [][]byte{}
 
-	walker := p.
+	walker := p.NewWalker()
-		NewWalker().
+
-		WithContentDispositionHandler("attachment", func(p *Part, hdl PartHandler) (err error) {
+	walker.
+		RegisterContentDispositionHandler("attachment").
+		OnEnter(func(p *Part, hdl PartHandlerFunc) (err error) {
 			attachments = append(attachments, p.Body)
 			return
 		})
@@ -62,3 +66,25 @@ func TestWalkerDispositionHandler(t *testing.T) {
 		[]byte("if you are reading this, hi!"),
 	}, attachments)
 }
+
+func TestWalkerDispositionAndTypeHandler(t *testing.T) {
+	p := newTestParser(t, "text_html_octet_attachment.eml")
+
+	walker := p.NewWalker()
+
+	var enter, exit int
+
+	walker.
+		RegisterContentTypeHandler("application/octet-stream").
+		OnEnter(func(p *Part) (err error) { enter++; return }).
+		OnExit(func(p *Part) (err error) { exit--; return })
+
+	walker.
+		RegisterContentDispositionHandler("attachment").
+		OnEnter(func(p *Part, hdl PartHandlerFunc) (err error) { _ = hdl(p); _ = hdl(p); return }).
+		OnExit(func(p *Part, hdl PartHandlerFunc) (err error) { _ = hdl(p); _ = hdl(p); return })
+
+	assert.NoError(t, walker.Walk())
+	assert.Equal(t, 2, enter)
+	assert.Equal(t, -2, exit)
+}