writer.go

     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 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 Reader's 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  			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  
   103  func (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
View as plain text