diff --git a/pkg/message/section.go b/pkg/message/section.go index 515b68ca..2e99927b 100644 --- a/pkg/message/section.go +++ b/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 } diff --git a/pkg/message/section_test.go b/pkg/message/section_test.go index baf5129f..69eaebee 100644 --- a/pkg/message/section_test.go +++ b/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 +To: Mary Smith +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