1// Copyright (c) 2016, Google Inc.
2//
3// Permission to use, copy, modify, and/or distribute this software for any
4// purpose with or without fee is hereby granted, provided that the above
5// copyright notice and this permission notice appear in all copies.
6//
7// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
8// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
9// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
10// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
11// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
12// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
13// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
14
15package runner
16
17import (
18	"bufio"
19	"encoding/hex"
20	"errors"
21	"fmt"
22	"io"
23	"net"
24	"strconv"
25	"strings"
26	"sync"
27)
28
29type flowType int
30
31const (
32	readFlow flowType = iota
33	writeFlow
34	specialFlow
35)
36
37type flow struct {
38	flowType flowType
39	message  string
40	data     []byte
41}
42
43// recordingConn is a net.Conn that records the traffic that passes through it.
44// WriteTo can be used to produce output that can be later be loaded with
45// ParseTestData.
46type recordingConn struct {
47	net.Conn
48	sync.Mutex
49	flows       []flow
50	isDatagram  bool
51	local, peer string
52}
53
54func (r *recordingConn) appendFlow(flowType flowType, message string, data []byte) {
55	r.Lock()
56	defer r.Unlock()
57
58	if l := len(r.flows); flowType == specialFlow || r.isDatagram || l == 0 || r.flows[l-1].flowType != flowType {
59		buf := make([]byte, len(data))
60		copy(buf, data)
61		r.flows = append(r.flows, flow{flowType, message, buf})
62	} else {
63		r.flows[l-1].data = append(r.flows[l-1].data, data...)
64	}
65}
66
67func (r *recordingConn) Read(b []byte) (n int, err error) {
68	if n, err = r.Conn.Read(b); n == 0 {
69		return
70	}
71	r.appendFlow(readFlow, "", b[:n])
72	return
73}
74
75func (r *recordingConn) Write(b []byte) (n int, err error) {
76	if n, err = r.Conn.Write(b); n == 0 {
77		return
78	}
79	r.appendFlow(writeFlow, "", b[:n])
80	return
81}
82
83// LogSpecial appends an entry to the record of type 'special'.
84func (r *recordingConn) LogSpecial(message string, data []byte) {
85	r.appendFlow(specialFlow, message, data)
86}
87
88// WriteTo writes hex dumps to w that contains the recorded traffic.
89func (r *recordingConn) WriteTo(w io.Writer) {
90	fmt.Fprintf(w, ">>> runner is %s, shim is %s\n", r.local, r.peer)
91	for i, flow := range r.flows {
92		switch flow.flowType {
93		case readFlow:
94			fmt.Fprintf(w, ">>> Flow %d (%s to %s)\n", i+1, r.peer, r.local)
95		case writeFlow:
96			fmt.Fprintf(w, ">>> Flow %d (%s to %s)\n", i+1, r.local, r.peer)
97		case specialFlow:
98			fmt.Fprintf(w, ">>> Flow %d %q\n", i+1, flow.message)
99		}
100
101		if flow.data != nil {
102			dumper := hex.Dumper(w)
103			dumper.Write(flow.data)
104			dumper.Close()
105		}
106	}
107}
108
109func (r *recordingConn) Transcript() []byte {
110	var ret []byte
111	for _, flow := range r.flows {
112		if flow.flowType != writeFlow {
113			continue
114		}
115		if r.isDatagram {
116			// Prepend a length prefix to preserve packet boundaries.
117			ret = append(ret, byte(len(flow.data)>>16), byte(len(flow.data)>>8), byte(len(flow.data)))
118		}
119		ret = append(ret, flow.data...)
120	}
121	return ret
122}
123
124func parseTestData(r io.Reader) (flows [][]byte, err error) {
125	var currentFlow []byte
126
127	scanner := bufio.NewScanner(r)
128	for scanner.Scan() {
129		line := scanner.Text()
130		// If the line starts with ">>> " then it marks the beginning
131		// of a new flow.
132		if strings.HasPrefix(line, ">>> ") {
133			if len(currentFlow) > 0 || len(flows) > 0 {
134				flows = append(flows, currentFlow)
135				currentFlow = nil
136			}
137			continue
138		}
139
140		// Otherwise the line is a line of hex dump that looks like:
141		// 00000170  fc f5 06 bf (...)  |.....X{&?......!|
142		// (Some bytes have been omitted from the middle section.)
143
144		if i := strings.IndexByte(line, ' '); i >= 0 {
145			line = line[i:]
146		} else {
147			return nil, errors.New("invalid test data")
148		}
149
150		if i := strings.IndexByte(line, '|'); i >= 0 {
151			line = line[:i]
152		} else {
153			return nil, errors.New("invalid test data")
154		}
155
156		hexBytes := strings.Fields(line)
157		for _, hexByte := range hexBytes {
158			val, err := strconv.ParseUint(hexByte, 16, 8)
159			if err != nil {
160				return nil, errors.New("invalid hex byte in test data: " + err.Error())
161			}
162			currentFlow = append(currentFlow, byte(val))
163		}
164	}
165
166	if len(currentFlow) > 0 {
167		flows = append(flows, currentFlow)
168	}
169
170	return flows, nil
171}
172