github.com/hernad/nomad@v1.6.112/helper/escapingio/reader.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package escapingio
     5  
     6  import (
     7  	"bufio"
     8  	"io"
     9  )
    10  
    11  // Handler is a callback for handling an escaped char.  Reader would skip
    12  // the escape char and passed char if returns true; otherwise, it preserves them
    13  // in output
    14  type Handler func(c byte) bool
    15  
    16  // NewReader returns a reader that escapes the c character (following new lines),
    17  // in the same manner OpenSSH handling, which defaults to `~`.
    18  //
    19  // For illustrative purposes, we use `~` in documentation as a shorthand for escaping character.
    20  //
    21  // If following a new line, reader sees:
    22  //   - `~~`, only one is emitted
    23  //   - `~.` (or any character), the handler is invoked with the character.
    24  //     If handler returns true, `~.` will be skipped; otherwise, it's propagated.
    25  //   - `~` and it's the last character in stream, it's propagated
    26  //
    27  // Appearances of `~` when not preceded by a new line are propagated unmodified.
    28  func NewReader(r io.Reader, c byte, h Handler) io.Reader {
    29  	pr, pw := io.Pipe()
    30  	reader := &reader{
    31  		impl:       r,
    32  		escapeChar: c,
    33  		handler:    h,
    34  		pr:         pr,
    35  		pw:         pw,
    36  	}
    37  	go reader.pipe()
    38  	return reader
    39  }
    40  
    41  // lookState represents the state of reader for what character of `\n~.` sequence
    42  // reader is looking for
    43  type lookState int
    44  
    45  const (
    46  	// sLookNewLine indicates that reader is looking for new line
    47  	sLookNewLine lookState = iota
    48  
    49  	// sLookEscapeChar indicates that reader is looking for ~
    50  	sLookEscapeChar
    51  
    52  	// sLookChar indicates that reader just read `~` is waiting for next character
    53  	// before acting
    54  	sLookChar
    55  )
    56  
    57  // to ease comments, i'll assume escape character to be `~`
    58  type reader struct {
    59  	impl       io.Reader
    60  	escapeChar uint8
    61  	handler    Handler
    62  
    63  	// buffers
    64  	pw *io.PipeWriter
    65  	pr *io.PipeReader
    66  }
    67  
    68  func (r *reader) Read(buf []byte) (int, error) {
    69  	return r.pr.Read(buf)
    70  }
    71  
    72  func (r *reader) pipe() {
    73  	rb := make([]byte, 4096)
    74  	bw := bufio.NewWriter(r.pw)
    75  
    76  	state := sLookEscapeChar
    77  
    78  	for {
    79  		n, err := r.impl.Read(rb)
    80  
    81  		if n > 0 {
    82  			state = r.processBuf(bw, rb, n, state)
    83  			bw.Flush()
    84  			if state == sLookChar {
    85  				// terminated with ~ - let's read one more character
    86  				n, err = r.impl.Read(rb[:1])
    87  				if n == 1 {
    88  					state = sLookNewLine
    89  					if rb[0] == r.escapeChar {
    90  						// only emit escape character once
    91  						bw.WriteByte(rb[0])
    92  						bw.Flush()
    93  					} else if r.handler(rb[0]) {
    94  						// skip if handled
    95  					} else {
    96  						bw.WriteByte(r.escapeChar)
    97  						bw.WriteByte(rb[0])
    98  						bw.Flush()
    99  						if rb[0] == '\n' || rb[0] == '\r' {
   100  							state = sLookEscapeChar
   101  						}
   102  					}
   103  				}
   104  			}
   105  		}
   106  
   107  		if err != nil {
   108  			// write ~ if it's the last thing
   109  			if state == sLookChar {
   110  				bw.WriteByte(r.escapeChar)
   111  			}
   112  			bw.Flush()
   113  			r.pw.CloseWithError(err)
   114  			break
   115  		}
   116  	}
   117  }
   118  
   119  // processBuf process buffer and emits all output to writer
   120  // if the last part of buffer is a new line followed by sequnce, it writes
   121  // all output until the new line and returns sLookChar
   122  func (r *reader) processBuf(bw io.Writer, buf []byte, n int, s lookState) lookState {
   123  	i := 0
   124  
   125  	wi := 0
   126  
   127  START:
   128  	if s == sLookEscapeChar && buf[i] == r.escapeChar {
   129  		if i+1 >= n {
   130  			// buf terminates with ~ - write all before
   131  			bw.Write(buf[wi:i])
   132  			return sLookChar
   133  		}
   134  
   135  		nc := buf[i+1]
   136  		if nc == r.escapeChar {
   137  			// skip one escape char
   138  			bw.Write(buf[wi:i])
   139  			i++
   140  			wi = i
   141  		} else if r.handler(nc) {
   142  			// skip both characters
   143  			bw.Write(buf[wi:i])
   144  			i = i + 2
   145  			wi = i
   146  		} else if nc == '\n' || nc == '\r' {
   147  			i = i + 2
   148  			s = sLookEscapeChar
   149  			goto START
   150  		} else {
   151  			i = i + 2
   152  			// need to write everything keep going
   153  		}
   154  	}
   155  
   156  	// search until we get \n~, or buf terminates
   157  	for {
   158  		if i >= n {
   159  			// got to end without new line, write and return
   160  			bw.Write(buf[wi:n])
   161  			return sLookNewLine
   162  		}
   163  
   164  		if buf[i] == '\n' || buf[i] == '\r' {
   165  			// buf terminated at new line
   166  			if i+1 >= n {
   167  				bw.Write(buf[wi:n])
   168  				return sLookEscapeChar
   169  			}
   170  
   171  			// peek to see escape character go back to START if so
   172  			if buf[i+1] == r.escapeChar {
   173  				s = sLookEscapeChar
   174  				i++
   175  				goto START
   176  			}
   177  		}
   178  
   179  		i++
   180  	}
   181  }