refactor: Better MIMEType usage

ProtonMail · Oct 6, 2022 · 634f8b9 · 634f8b9
1 parent 44257cf
commit 634f8b9
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 18 deletions.
diff --git a/imap/structure.go b/imap/structure.go
@@ -2,7 +2,6 @@ package imap
 
 import (
 	"bytes"
-	"errors"
 	"mime"
 	"strings"
 
@@ -47,7 +46,7 @@ func structure(section *rfc822.Section, fields *paramList, writer *dualParListWr
 		return err
 	}
 
-	_, mimeSubType, mimeParams, err := getMIMEInfo(header)
+	_, mimeSubType, mimeParams, err := getMIMEInfo(section)
 	if err != nil {
 		return err
 	}
@@ -69,7 +68,7 @@ func singlePartStructure(section *rfc822.Section, fields *paramList, writer *dua
 		return err
 	}
 
-	mimeType, mimeSubType, mimeParams, err := getMIMEInfo(header)
+	mimeType, mimeSubType, mimeParams, err := getMIMEInfo(section)
 	if err != nil {
 		return err
 	}
@@ -138,18 +137,13 @@ func childStructures(section *rfc822.Section, c *paramList, writer *dualParListW
 	return nil
 }
 
-func getMIMEInfo(header *rfc822.Header) (string, string, map[string]string, error) {
-	contentType, contentTypeParams, err := rfc822.ParseContentType(header.Get("Content-Type"))
+func getMIMEInfo(section *rfc822.Section) (string, string, map[string]string, error) {
+	mimeType, mimeParams, err := section.ContentType()
 	if err != nil {
 		return "", "", nil, err
 	}
 
-	split := strings.Split(contentType, "/")
-	if len(split) != 2 {
-		return "", "", nil, errors.New("malformed MIME type")
-	}
-
-	return split[0], split[1], contentTypeParams, nil
+	return mimeType.Type(), mimeType.SubType(), mimeParams, nil
 }
 
 func addDispInfo(c *paramList, writer parListWriter, header *rfc822.Header) {

diff --git a/rfc822/mime.go b/rfc822/mime.go
@@ -1,6 +1,9 @@
 package rfc822
 
-import "mime"
+import (
+	"mime"
+	"strings"
+)
 
 type MIMEType string
 
@@ -12,10 +15,35 @@ const (
 	MessageRFC822    MIMEType = "message/rfc822"
 )
 
-func ParseContentType(val string) (string, map[string]string, error) {
+func (mimeType MIMEType) IsMultiPart() bool {
+	return strings.HasPrefix(string(mimeType), "multipart/")
+}
+
+func (mimeType MIMEType) Type() string {
+	if split := strings.SplitN(string(mimeType), "/", 2); len(split) == 2 {
+		return split[0]
+	}
+
+	return ""
+}
+
+func (mimeType MIMEType) SubType() string {
+	if split := strings.SplitN(string(mimeType), "/", 2); len(split) == 2 {
+		return split[1]
+	}
+
+	return ""
+}
+
+func parseMIMEType(val string) (MIMEType, map[string]string, error) {
 	if val == "" {
 		val = string(TextPlain)
 	}
 
-	return mime.ParseMediaType(val)
+	mimeType, mimeParams, err := mime.ParseMediaType(val)
+	if err != nil {
+		return "", nil, err
+	}
+
+	return MIMEType(mimeType), mimeParams, nil
 }
diff --git a/rfc822/parser.go b/rfc822/parser.go
@@ -3,7 +3,6 @@ package rfc822
 import (
 	"bytes"
 	"fmt"
-	"strings"
 )
 
 type Section struct {
@@ -24,13 +23,13 @@ func (section *Section) Identifier() []int {
 	return section.identifier
 }
 
-func (section *Section) ContentType() (string, map[string]string, error) {
+func (section *Section) ContentType() (MIMEType, map[string]string, error) {
 	header, err := section.ParseHeader()
 	if err != nil {
 		return "", nil, err
 	}
 
-	return ParseContentType(header.Get("Content-Type"))
+	return parseMIMEType(header.Get("Content-Type"))
 }
 
 func (section *Section) Header() []byte {
@@ -106,7 +105,7 @@ func (section *Section) load() error {
 		}
 
 		section.children = append(section.children, child.children...)
-	} else if strings.HasPrefix(contentType, "multipart/") {
+	} else if contentType.IsMultiPart() {
 		scanner, err := NewByteScanner(section.literal[section.body:section.end], []byte(contentParams["boundary"]))
 		if err != nil {
 			return err