forked from networklore/netrasp
/
reader.go
84 lines (70 loc) · 1.62 KB
/
reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
package netrasp
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"regexp"
)
var errRead = errors.New("reader error")
var minReadSize = 10000
var maxReadSize = 10 * 1000 * 1000
type contextReader struct {
ctx context.Context
r io.Reader
}
type readResult struct {
n int
err error
}
func (c *contextReader) Read(p []byte) (n int, err error) {
ctx, cancel := context.WithCancel(c.ctx)
defer cancel()
rrCh := make(chan *readResult)
go func() {
select {
case rrCh <- c.read(p):
case <-ctx.Done():
}
}()
select {
case <-ctx.Done():
return 0, ctx.Err()
case rr := <-rrCh:
return rr.n, rr.err
}
}
func (c *contextReader) read(p []byte) *readResult {
n, err := c.r.Read(p)
return &readResult{n, err}
}
func newContextReader(ctx context.Context, r io.Reader) io.Reader {
return &contextReader{
ctx: ctx,
r: r,
}
}
// readUntilPrompt reads until the specified prompt is found and returns the read data.
func readUntilPrompt(ctx context.Context, r io.Reader, prompt *regexp.Regexp) (string, error) {
var output = new(bytes.Buffer)
rc := newContextReader(ctx, r)
readSize := minReadSize
for {
b := make([]byte, readSize)
n, err := rc.Read(b)
if err != nil {
return "", fmt.Errorf("error reading output from device %w: %v", errRead, err)
}
if readSize == n && readSize < maxReadSize {
readSize *= 2
}
output.Write(b[:n])
tempSlice := bytes.ReplaceAll(output.Bytes(), []byte("\r\n"), []byte("\n"))
tempSlice = bytes.ReplaceAll(tempSlice, []byte("\r"), []byte("\n"))
if prompt.Match(tempSlice[bytes.LastIndex(tempSlice, []byte("\n"))+1:]) {
break
}
}
return output.String(), nil
}