blob: 39deb19454eab1196155ca78377734b99ec0e181 [file] [log] [blame]
package runner
import (
"bufio"
"encoding/hex"
"errors"
"fmt"
"io"
"net"
"strconv"
"strings"
"sync"
)
type flowType int
const (
readFlow flowType = iota
writeFlow
specialFlow
)
type flow struct {
flowType flowType
message string
data []byte
}
// recordingConn is a net.Conn that records the traffic that passes through it.
// WriteTo can be used to produce output that can be later be loaded with
// ParseTestData.
type recordingConn struct {
net.Conn
sync.Mutex
flows []flow
isDatagram bool
local, peer string
}
func (r *recordingConn) appendFlow(flowType flowType, message string, data []byte) {
r.Lock()
defer r.Unlock()
if l := len(r.flows); flowType == specialFlow || r.isDatagram || l == 0 || r.flows[l-1].flowType != flowType {
buf := make([]byte, len(data))
copy(buf, data)
r.flows = append(r.flows, flow{flowType, message, buf})
} else {
r.flows[l-1].data = append(r.flows[l-1].data, data...)
}
}
func (r *recordingConn) Read(b []byte) (n int, err error) {
if n, err = r.Conn.Read(b); n == 0 {
return
}
r.appendFlow(readFlow, "", b[:n])
return
}
func (r *recordingConn) Write(b []byte) (n int, err error) {
if n, err = r.Conn.Write(b); n == 0 {
return
}
r.appendFlow(writeFlow, "", b[:n])
return
}
// LogSpecial appends an entry to the record of type 'special'.
func (r *recordingConn) LogSpecial(message string, data []byte) {
r.appendFlow(specialFlow, message, data)
}
// WriteTo writes hex dumps to w that contains the recorded traffic.
func (r *recordingConn) WriteTo(w io.Writer) {
fmt.Fprintf(w, ">>> runner is %s, shim is %s\n", r.local, r.peer)
for i, flow := range r.flows {
switch flow.flowType {
case readFlow:
fmt.Fprintf(w, ">>> Flow %d (%s to %s)\n", i+1, r.peer, r.local)
case writeFlow:
fmt.Fprintf(w, ">>> Flow %d (%s to %s)\n", i+1, r.local, r.peer)
case specialFlow:
fmt.Fprintf(w, ">>> Flow %d %q\n", i+1, flow.message)
}
if flow.data != nil {
dumper := hex.Dumper(w)
dumper.Write(flow.data)
dumper.Close()
}
}
}
func parseTestData(r io.Reader) (flows [][]byte, err error) {
var currentFlow []byte
scanner := bufio.NewScanner(r)
for scanner.Scan() {
line := scanner.Text()
// If the line starts with ">>> " then it marks the beginning
// of a new flow.
if strings.HasPrefix(line, ">>> ") {
if len(currentFlow) > 0 || len(flows) > 0 {
flows = append(flows, currentFlow)
currentFlow = nil
}
continue
}
// Otherwise the line is a line of hex dump that looks like:
// 00000170 fc f5 06 bf (...) |.....X{&?......!|
// (Some bytes have been omitted from the middle section.)
if i := strings.IndexByte(line, ' '); i >= 0 {
line = line[i:]
} else {
return nil, errors.New("invalid test data")
}
if i := strings.IndexByte(line, '|'); i >= 0 {
line = line[:i]
} else {
return nil, errors.New("invalid test data")
}
hexBytes := strings.Fields(line)
for _, hexByte := range hexBytes {
val, err := strconv.ParseUint(hexByte, 16, 8)
if err != nil {
return nil, errors.New("invalid hex byte in test data: " + err.Error())
}
currentFlow = append(currentFlow, byte(val))
}
}
if len(currentFlow) > 0 {
flows = append(flows, currentFlow)
}
return flows, nil
}