WIP: message: correctly wrap input lines for 7bit/8bit

encoding.go

@@ -7,8 +7,9 @@ import (
 	"io"
 	"mime/quotedprintable"
 	"strings"
+	"bytes"
 
-	"github.com/emersion/go-textwrapper"
+	"log"
 )
 
 type UnknownEncodingError struct {
@@ -56,13 +57,97 @@ func encodingWriter(enc string, w io.Writer) (io.WriteCloser, error) {
 	case "quoted-printable":
 		wc = quotedprintable.NewWriter(w)
 	case "base64":
-		wc = base64.NewEncoder(base64.StdEncoding, textwrapper.NewRFC822(w))
+		wc = base64.NewEncoder(base64.StdEncoding, &wrapWriter{w: w, max: 76})
 	case "7bit", "8bit":
-		wc = nopCloser{textwrapper.New(w, "\r\n", 1000)}
+		wc = nopCloser{&wrapWriter{w: w, max: 1000}}
 	case "binary", "":
 		wc = nopCloser{w}
 	default:
 		return nil, fmt.Errorf("unhandled encoding %q", enc)
 	}
 	return wc, nil
 }
+
+// wrapWriter is an io.Writer that wraps long text lines to a specified length.
+type wrapWriter struct {
+	w   io.Writer
+	max int // including CRLF
+
+	n  int  // current line length
+	cr bool // previous byte was \r
+	crlf bool // previous bytes were \r\n
+}
+
+func (w *wrapWriter) Write(b []byte) (int, error) {
+	N := 0
+	for len(b) > 0 {
+		i := bytes.IndexByte(b, '\n')
+
+		to := i + 1
+		if i < 0 || to > w.max - w.n + 2 {
+			to = w.max - w.n + 2
+			if to > len(b) {
+				to = len(b)
+			} else if b[to-2] == '\n' {
+				to--
+			} else if b[to-1] != '\n' || b[to-2] != '\r' {
+				to -= 2
+			}
+		}
+
+		n, err := w.writeChunk(b[:to])
+		N += n
+		if err != nil {
+			return N, err
+		}
+
+		b = b[to:]
+	}
+
+	return N, nil
+}
+
+func (w *wrapWriter) writeChunk(b []byte) (int, error) {
+	lf := bytes.HasSuffix(b, []byte{'\n'})
+	crlf := (w.cr && bytes.HasPrefix(b, []byte{'\n'})) || bytes.HasSuffix(b, []byte{'\r', '\n'})
+
+	log.Printf("%q crlf=%v n=%v", string(b), w.crlf, w.n)
+	if !w.crlf && w.n >= w.max {
+		// If the previous line didn't end with a CRLF, write one
+		if _, err := w.w.Write([]byte{'\r', '\n'}); err != nil {
+			return 0, err
+		}
+		w.n = 0
+	}
+
+	var (
+		n   int
+		err error
+	)
+	if lf && !crlf {
+		// Need to convert lone LF to CRLF
+		n, err = w.w.Write(b[:len(b)-1])
+		if err != nil {
+			return n, err
+		}
+		if _, err := w.w.Write([]byte{'\r', '\n'}); err != nil {
+			return n, err
+		}
+		n++
+		w.crlf = true
+	} else {
+		n, err = w.w.Write(b)
+		if err != nil {
+			return n, err
+		}
+		w.crlf = crlf
+	}
+
+	w.cr = bytes.HasSuffix(b, []byte{'\r'})
+	if lf || crlf {
+		w.n = 0
+	} else {
+		w.n += n
+	}
+	return n, nil
+}

encoding_test.go

@@ -71,3 +71,117 @@ func TestEncode(t *testing.T) {
 		}
 	}
 }
+
+var wrapWriterTests = []struct{
+	name   string
+	writes []string
+	out    string
+}{
+	{
+		name:   "singleLine",
+		writes: []string{"Hey"},
+		out:    "Hey",
+	},
+	{
+		name: "twoLines",
+		writes: []string{"Hey\r\nYo"},
+		out: "Hey\r\nYo",
+	},
+	{
+		name: "finalCRLF",
+		writes: []string{"Hey\r\n"},
+		out: "Hey\r\n",
+	},
+	{
+		name: "singleLineSplit",
+		writes: []string{"He", "y"},
+		out: "Hey",
+	},
+	{
+		name: "twoLinesSplit",
+		writes: []string{"He", "y", "\r\nY", "o"},
+		out: "Hey\r\nYo",
+	},
+	{
+		name: "longLine",
+		writes: []string{"How are you today?"},
+		out: "How are yo\r\nu today?",
+	},
+	{
+		name: "longLineSplit",
+		writes: []string{"How are ", "you today?"},
+		out: "How are yo\r\nu today?",
+	},
+	{
+		name: "lf",
+		writes: []string{"Hey\nYo"},
+		out: "Hey\r\nYo",
+	},
+	{
+		name: "max",
+		writes: []string{"Hey there!"},
+		out: "Hey there!",
+	},
+	{
+		name: "maxCRLF",
+		writes: []string{"Hey there!\r\n"},
+		out: "Hey there!\r\n",
+	},
+	{
+		name: "maxLF",
+		writes: []string{"Hey there!\n"},
+		out: "Hey there!\r\n",
+	},
+	{
+		name: "maxMinusOne",
+		writes: []string{"Hey there"},
+		out: "Hey there",
+	},
+	{
+		name: "maxMinusOneCRLF",
+		writes: []string{"Hey there\r\n"},
+		out: "Hey there\r\n",
+	},
+	{
+		name: "maxMinusOneLF",
+		writes: []string{"Hey there\n"},
+		out: "Hey there\r\n",
+	},
+	{
+		name: "maxPlusOne",
+		writes: []string{"Hey there!!"},
+		out: "Hey there!\r\n!",
+	},
+	{
+		name: "maxPlusTwo",
+		writes: []string{"Hey there!!!"},
+		out: "Hey there!\r\n!!",
+	},
+	{
+		name: "maxSplit",
+		writes: []string{"Hey ", "there!", "\r", "\n"},
+		out: "Hey there!\r\n",
+	},
+}
+
+func TestWrapWriter(t *testing.T) {
+	for _, test := range wrapWriterTests {
+		t.Run(test.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			ww := wrapWriter{
+				w:   &buf,
+				max: 10,
+			}
+
+			for _, s := range test.writes {
+				if _, err := ww.Write([]byte(s)); err != nil {
+					t.Fatalf("wrapWriter.Write() = %v", err)
+				}
+			}
+
+			if buf.String() != test.out {
+				t.Errorf("got %q, want %q", buf.String(), test.out)
+			}
+		})
+	}
+}