yesBad/proton-bridge
1
0
Fork 0
forked from Silverfish/proton-bridge
Code Actions Activity

feat: enter and exit handlers

Browse Source
This commit is contained in:
James Houlahan
2020-07-29 13:59:52 +02:00
parent 8bd74c5edc
commit a7da66ccbc
6 changed files with 269 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
pkg/message/parser.go
View File

@ -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
if m.Body, plain, isHTML, err = collectBodyParts(p); err != nil {
richBody, plainBody, err := collectBodyParts(p)
if err != nil {
return
}
m.Body = richBody
if isHTML {
m.MIMEType = "text/html"
} else {
m.MIMEType = "text/plain"
mimeType, err := determineMIMEType(p)
if err != nil {
return
}
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.
NewWalker().
WithContentDispositionHandler("attachment", func(p *parser.Part, _ parser.PartHandler) (err error) {
w := p.NewWalker()
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) {
var parts, plainParts []string
// collectBodyParts returns a richtext body (used for normal sending)
// 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.
NewWalker().
WithContentTypeHandler("text/plain", func(p *parser.Part) (err error) {
parts = append(parts, string(p.Body))
w := p.NewWalker()
w.RegisterContentTypeHandler("text/plain").
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 err != nil {
text = string(p.Body)
if !isAlternative(p) {
richParts = append(richParts, 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) {

72
pkg/message/parser/handlers.go Normal file
View File

@ -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 }

6
pkg/message/parser/parser.go
View File

@ -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]
}

65
pkg/message/parser/part.go
View File

@ -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) {

65
pkg/message/parser/walker.go
View File

@ -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]
}

42
pkg/message/parser/walker_test.go
View File

@ -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)
}