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 5 package textproto
6 7 import (
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.
15 type Writer struct {
16 W *bufio.Writer
17 dot *dotWriter
18 }
19 20 // NewWriter returns a new [Writer] writing to w.
21 func NewWriter(w *bufio.Writer) *Writer {
22 return &Writer{W: w}
23 }
24 25 var crnl = []byte{'\r', '\n'}
26 var dotcrnl = []byte{'.', '\r', '\n'}
27 28 // PrintfLine writes the formatted output followed by \r\n.
29 func (w *Writer) PrintfLine(format []byte, 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.
43 func (w *Writer) DotWriter() io.WriteCloser {
44 w.closeDot()
45 w.dot = &dotWriter{w: w}
46 return w.dot
47 }
48 49 func (w *Writer) closeDot() {
50 if w.dot != nil {
51 w.dot.Close() // sets w.dot = nil
52 }
53 }
54 55 type dotWriter struct {
56 w *Writer
57 state int
58 }
59 60 const (
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 67 func (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 if c == '\r' {
79 d.state = wstateCR
80 }
81 if c == '\n' {
82 bw.WriteByte('\r')
83 d.state = wstateBeginLine
84 }
85 86 case wstateData:
87 if c == '\r' {
88 d.state = wstateCR
89 }
90 if c == '\n' {
91 bw.WriteByte('\r')
92 d.state = wstateBeginLine
93 }
94 95 case wstateCR:
96 d.state = wstateData
97 if c == '\n' {
98 d.state = wstateBeginLine
99 }
100 }
101 if err = bw.WriteByte(c); err != nil {
102 break
103 }
104 n++
105 }
106 return
107 }
108 109 func (d *dotWriter) Close() error {
110 if d.w.dot == d {
111 d.w.dot = nil
112 }
113 bw := d.w.W
114 switch d.state {
115 default:
116 bw.WriteByte('\r')
117 bw.WriteByte('\n')
118 bw.Write(dotcrnl)
119 case wstateCR:
120 bw.WriteByte('\n')
121 bw.Write(dotcrnl)
122 case wstateBeginLine:
123 bw.Write(dotcrnl)
124 }
125 return bw.Flush()
126 }
127