github.com/slayercat/go@v0.0.0-20170428012452-c51559813f61/src/net/textproto/writer.go (about)

     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 ...interface{}) 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  	wstateBeginLine = iota // beginning of line; initial state; must be zero
    62  	wstateCR               // wrote \r (possibly at end of line)
    63  	wstateData             // writing data in middle of line
    64  )
    65  
    66  func (d *dotWriter) Write(b []byte) (n int, err error) {
    67  	bw := d.w.W
    68  	for n < len(b) {
    69  		c := b[n]
    70  		switch d.state {
    71  		case wstateBeginLine:
    72  			d.state = wstateData
    73  			if c == '.' {
    74  				// escape leading dot
    75  				bw.WriteByte('.')
    76  			}
    77  			fallthrough
    78  
    79  		case wstateData:
    80  			if c == '\r' {
    81  				d.state = wstateCR
    82  			}
    83  			if c == '\n' {
    84  				bw.WriteByte('\r')
    85  				d.state = wstateBeginLine
    86  			}
    87  
    88  		case wstateCR:
    89  			d.state = wstateData
    90  			if c == '\n' {
    91  				d.state = wstateBeginLine
    92  			}
    93  		}
    94  		if err = bw.WriteByte(c); err != nil {
    95  			break
    96  		}
    97  		n++
    98  	}
    99  	return
   100  }
   101  
   102  func (d *dotWriter) Close() error {
   103  	if d.w.dot == d {
   104  		d.w.dot = nil
   105  	}
   106  	bw := d.w.W
   107  	switch d.state {
   108  	default:
   109  		bw.WriteByte('\r')
   110  		fallthrough
   111  	case wstateCR:
   112  		bw.WriteByte('\n')
   113  		fallthrough
   114  	case wstateBeginLine:
   115  		bw.Write(dotcrnl)
   116  	}
   117  	return bw.Flush()
   118  }