1// Copyright 2010 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package textproto
6
7import (
8	"bufio"
9	"fmt"
10	"io"
11)
12
13// A Writer implements convenience methods for writing
14// requests or responses to a text protocol network connection.
15type Writer struct {
16	W   *bufio.Writer
17	dot *dotWriter
18}
19
20// NewWriter returns a new [Writer] writing to w.
21func NewWriter(w *bufio.Writer) *Writer {
22	return &Writer{W: w}
23}
24
25var crnl = []byte{'\r', '\n'}
26var dotcrnl = []byte{'.', '\r', '\n'}
27
28// PrintfLine writes the formatted output followed by \r\n.
29func (w *Writer) PrintfLine(format string, args ...any) error {
30	w.closeDot()
31	fmt.Fprintf(w.W, format, args...)
32	w.W.Write(crnl)
33	return w.W.Flush()
34}
35
36// DotWriter returns a writer that can be used to write a dot-encoding to w.
37// It takes care of inserting leading dots when necessary,
38// translating line-ending \n into \r\n, and adding the final .\r\n line
39// when the DotWriter is closed. The caller should close the
40// DotWriter before the next call to a method on w.
41//
42// See the documentation for the [Reader.DotReader] method for details about dot-encoding.
43func (w *Writer) DotWriter() io.WriteCloser {
44	w.closeDot()
45	w.dot = &dotWriter{w: w}
46	return w.dot
47}
48
49func (w *Writer) closeDot() {
50	if w.dot != nil {
51		w.dot.Close() // sets w.dot = nil
52	}
53}
54
55type dotWriter struct {
56	w     *Writer
57	state int
58}
59
60const (
61	wstateBegin     = iota // initial state; must be zero
62	wstateBeginLine        // beginning of line
63	wstateCR               // wrote \r (possibly at end of line)
64	wstateData             // writing data in middle of line
65)
66
67func (d *dotWriter) Write(b []byte) (n int, err error) {
68	bw := d.w.W
69	for n < len(b) {
70		c := b[n]
71		switch d.state {
72		case wstateBegin, wstateBeginLine:
73			d.state = wstateData
74			if c == '.' {
75				// escape leading dot
76				bw.WriteByte('.')
77			}
78			fallthrough
79
80		case wstateData:
81			if c == '\r' {
82				d.state = wstateCR
83			}
84			if c == '\n' {
85				bw.WriteByte('\r')
86				d.state = wstateBeginLine
87			}
88
89		case wstateCR:
90			d.state = wstateData
91			if c == '\n' {
92				d.state = wstateBeginLine
93			}
94		}
95		if err = bw.WriteByte(c); err != nil {
96			break
97		}
98		n++
99	}
100	return
101}
102
103func (d *dotWriter) Close() error {
104	if d.w.dot == d {
105		d.w.dot = nil
106	}
107	bw := d.w.W
108	switch d.state {
109	default:
110		bw.WriteByte('\r')
111		fallthrough
112	case wstateCR:
113		bw.WriteByte('\n')
114		fallthrough
115	case wstateBeginLine:
116		bw.Write(dotcrnl)
117	}
118	return bw.Flush()
119}
120