feat: enter and exit handlers

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) {
-	t, _, err := p.Header.ContentType()
-	if err != nil {
-		return
+func (p *Part) getHandler(w *Walker) handler {
+	if dispHandler := w.getDispHandler(p); dispHandler != nil {
+		return dispHandler
 	}
 
-	return w.typeHandlers[t]
-}
-
-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)
-	}
+	return w.getTypeHandler(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
-	typeHandlers   map[string]PartHandler
-	dispHandlers   map[string]DispHandler
+	defaultHandler handler
+	typeHandlers   map[string]handler
+	dispHandlers   map[string]handler
 }
 
-type PartHandler func(*Part) error
-type DispHandler func(*Part, PartHandler) error
+type handler interface {
+	handleEnter(*Walker, *Part) error
+	handleExit(*Walker, *Part) error
+}
 
 func newWalker(root *Part) *Walker {
 	return &Walker{
 		root:           root,
-		defaultHandler: func(*Part) (err error) { return },
-		typeHandlers:   make(map[string]PartHandler),
-		dispHandlers:   make(map[string]DispHandler),
+		defaultHandler: NewPartHandler(),
+		typeHandlers:   make(map[string]handler),
+		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 {
-	w.typeHandlers[contType] = handler
-	return w
+func (w *Walker) RegisterContentTypeHandler(contType string) *PartHandler {
+	hdl := NewPartHandler()
+
+	w.typeHandlers[contType] = hdl
+
+	return hdl
 }
 
-func (w *Walker) WithContentDispositionHandler(contDisp string, handler DispHandler) *Walker {
-	w.dispHandlers[contDisp] = handler
-	return w
+func (w *Walker) RegisterContentDispositionHandler(contDisp string) *DispHandler {
+	hdl := NewDispHandler()
+
+	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.
-		NewWalker().
-		WithContentTypeHandler("text/html", func(p *Part) (err error) {
+	walker := p.NewWalker()
+
+	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.
-		NewWalker().
-		WithContentDispositionHandler("attachment", func(p *Part, hdl PartHandler) (err error) {
+	walker := p.NewWalker()
+
+	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)
+}