GODT-1162: Fix wrong section 1 error when email has no MIME parts

pkg/message/section.go

@@ -212,6 +212,13 @@ func (bs *BodyStructure) hasInfo(sectionPath []int) bool {
 	return err == nil
 }
 
+func (bs *BodyStructure) getInfoCheckSection(sectionPath []int) (sectionInfo *SectionInfo, err error) {
+	if len(*bs) == 1 && len(sectionPath) == 1 && sectionPath[0] == 1 {
+		sectionPath = []int{}
+	}
+	return bs.getInfo(sectionPath)
+}
+
 func (bs *BodyStructure) getInfo(sectionPath []int) (sectionInfo *SectionInfo, err error) {
 	path := stringPathFromInts(sectionPath)
 	sectionInfo, ok := (*bs)[path]
@@ -223,7 +230,7 @@ func (bs *BodyStructure) getInfo(sectionPath []int) (sectionInfo *SectionInfo, e
 
 // GetSection returns bytes of section including MIME header.
 func (bs *BodyStructure) GetSection(wholeMail io.ReadSeeker, sectionPath []int) (section []byte, err error) {
-	info, err := bs.getInfo(sectionPath)
+	info, err := bs.getInfoCheckSection(sectionPath)
 	if err != nil {
 		return
 	}
@@ -232,7 +239,7 @@ func (bs *BodyStructure) GetSection(wholeMail io.ReadSeeker, sectionPath []int)
 
 // GetSectionContent returns bytes of section content (excluding MIME header).
 func (bs *BodyStructure) GetSectionContent(wholeMail io.ReadSeeker, sectionPath []int) (section []byte, err error) {
-	info, err := bs.getInfo(sectionPath)
+	info, err := bs.getInfoCheckSection(sectionPath)
 	if err != nil {
 		return
 	}
@@ -251,8 +258,11 @@ func (bs *BodyStructure) GetMailHeaderBytes(wholeMail io.ReadSeeker) (header []b
 }
 
 func goToOffsetAndReadNBytes(wholeMail io.ReadSeeker, offset, length int) ([]byte, error) {
-	if length < 1 {
-		return nil, errors.New("requested non positive length")
+	if length == 0 {
+		return []byte{}, nil
+	}
+	if length < 0 {
+		return nil, errors.New("requested negative length")
 	}
 	if offset > 0 {
 		if _, err := wholeMail.Seek(int64(offset), io.SeekStart); err != nil {
@@ -266,7 +276,7 @@ func goToOffsetAndReadNBytes(wholeMail io.ReadSeeker, offset, length int) ([]byt
 
 // GetSectionHeader returns the mime header of specified section.
 func (bs *BodyStructure) GetSectionHeader(sectionPath []int) (header textproto.MIMEHeader, err error) {
-	info, err := bs.getInfo(sectionPath)
+	info, err := bs.getInfoCheckSection(sectionPath)
 	if err != nil {
 		return
 	}
@@ -275,7 +285,7 @@ func (bs *BodyStructure) GetSectionHeader(sectionPath []int) (header textproto.M
 }
 
 func (bs *BodyStructure) GetSectionHeaderBytes(wholeMail io.ReadSeeker, sectionPath []int) (header []byte, err error) {
-	info, err := bs.getInfo(sectionPath)
+	info, err := bs.getInfoCheckSection(sectionPath)
 	if err != nil {
 		return
 	}

pkg/message/section_test.go

@@ -82,6 +82,14 @@ func TestGetSection(t *testing.T) {
 	structReader := strings.NewReader(sampleMail)
 	bs, err := NewBodyStructure(structReader)
 	require.NoError(t, err)
+
+	// Bad paths
+	wantPaths := [][]int{{0}, {-1}, {3, 2, 3}}
+	for _, wantPath := range wantPaths {
+		_, err = bs.getInfo(wantPath)
+		require.Error(t, err, "path %v", wantPath)
+	}
+
 	// Whole section.
 	for _, try := range testPaths {
 		mailReader := strings.NewReader(sampleMail)
@@ -108,6 +116,60 @@ func TestGetSection(t *testing.T) {
 	}
 }
 
+func TestGetSecionNoMIMEParts(t *testing.T) {
+	wantBody := "This is just a simple mail with no multipart structure.\n"
+	wantHeader := `Subject: Sample mail
+From: John Doe <jdoe@machine.example>
+To: Mary Smith <mary@example.net>
+Date: Fri, 21 Nov 1997 09:55:06 -0600
+Content-Type: plain/text
+
+`
+	wantMail := wantHeader + wantBody
+
+	r := require.New(t)
+	bs, err := NewBodyStructure(strings.NewReader(wantMail))
+	r.NoError(err)
+
+	// Bad parts
+	wantPaths := [][]int{{0}, {2}, {1, 2, 3}}
+	for _, wantPath := range wantPaths {
+		_, err = bs.getInfoCheckSection(wantPath)
+		r.Error(err, "path %v: %d %d\n__\n%s\n", wantPath)
+	}
+
+	debug := func(wantPath []int, info *SectionInfo, section []byte) string {
+		if info == nil {
+			info = &SectionInfo{}
+		}
+		return fmt.Sprintf("path %v %q: %d %d\n___\n%s\n‾‾‾\n",
+			wantPath, stringPathFromInts(wantPath), info.Start, info.Size,
+			string(section),
+		)
+	}
+
+	// Ok Parts
+	wantPaths = [][]int{{}, {1}}
+	for _, p := range wantPaths {
+		wantPath := append([]int{}, p...)
+
+		info, err := bs.getInfoCheckSection(wantPath)
+		r.NoError(err, debug(wantPath, info, []byte{}))
+
+		section, err := bs.GetSection(strings.NewReader(wantMail), wantPath)
+		r.NoError(err, debug(wantPath, info, section))
+		r.Equal(wantMail, string(section), debug(wantPath, info, section))
+
+		haveBody, err := bs.GetSectionContent(strings.NewReader(wantMail), wantPath)
+		r.NoError(err, debug(wantPath, info, haveBody))
+		r.Equal(wantBody, string(haveBody), debug(wantPath, info, haveBody))
+
+		haveHeader, err := bs.GetSectionHeaderBytes(strings.NewReader(wantMail), wantPath)
+		r.NoError(err, debug(wantPath, info, haveHeader))
+		r.Equal(wantHeader, string(haveHeader), debug(wantPath, info, haveHeader))
+	}
+}
+
 func TestGetMainHeaderBytes(t *testing.T) {
 	wantHeader := []byte(`Subject: Sample mail
 From: John Doe <jdoe@machine.example>