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 }