1package main
2
3import (
4	"bufio"
5	"encoding/hex"
6	"errors"
7	"fmt"
8	"io"
9	"net"
10	"strconv"
11	"strings"
12	"sync"
13)
14
15// recordingConn is a net.Conn that records the traffic that passes through it.
16// WriteTo can be used to produce output that can be later be loaded with
17// ParseTestData.
18type recordingConn struct {
19	net.Conn
20	sync.Mutex
21	flows   [][]byte
22	reading bool
23}
24
25func (r *recordingConn) Read(b []byte) (n int, err error) {
26	if n, err = r.Conn.Read(b); n == 0 {
27		return
28	}
29	b = b[:n]
30
31	r.Lock()
32	defer r.Unlock()
33
34	if l := len(r.flows); l == 0 || !r.reading {
35		buf := make([]byte, len(b))
36		copy(buf, b)
37		r.flows = append(r.flows, buf)
38	} else {
39		r.flows[l-1] = append(r.flows[l-1], b[:n]...)
40	}
41	r.reading = true
42	return
43}
44
45func (r *recordingConn) Write(b []byte) (n int, err error) {
46	if n, err = r.Conn.Write(b); n == 0 {
47		return
48	}
49	b = b[:n]
50
51	r.Lock()
52	defer r.Unlock()
53
54	if l := len(r.flows); l == 0 || r.reading {
55		buf := make([]byte, len(b))
56		copy(buf, b)
57		r.flows = append(r.flows, buf)
58	} else {
59		r.flows[l-1] = append(r.flows[l-1], b[:n]...)
60	}
61	r.reading = false
62	return
63}
64
65// WriteTo writes hex dumps to w that contains the recorded traffic.
66func (r *recordingConn) WriteTo(w io.Writer) {
67	// TLS always starts with a client to server flow.
68	clientToServer := true
69
70	for i, flow := range r.flows {
71		source, dest := "client", "server"
72		if !clientToServer {
73			source, dest = dest, source
74		}
75		fmt.Fprintf(w, ">>> Flow %d (%s to %s)\n", i+1, source, dest)
76		dumper := hex.Dumper(w)
77		dumper.Write(flow)
78		dumper.Close()
79		clientToServer = !clientToServer
80	}
81}
82
83func parseTestData(r io.Reader) (flows [][]byte, err error) {
84	var currentFlow []byte
85
86	scanner := bufio.NewScanner(r)
87	for scanner.Scan() {
88		line := scanner.Text()
89		// If the line starts with ">>> " then it marks the beginning
90		// of a new flow.
91		if strings.HasPrefix(line, ">>> ") {
92			if len(currentFlow) > 0 || len(flows) > 0 {
93				flows = append(flows, currentFlow)
94				currentFlow = nil
95			}
96			continue
97		}
98
99		// Otherwise the line is a line of hex dump that looks like:
100		// 00000170  fc f5 06 bf (...)  |.....X{&?......!|
101		// (Some bytes have been omitted from the middle section.)
102
103		if i := strings.IndexByte(line, ' '); i >= 0 {
104			line = line[i:]
105		} else {
106			return nil, errors.New("invalid test data")
107		}
108
109		if i := strings.IndexByte(line, '|'); i >= 0 {
110			line = line[:i]
111		} else {
112			return nil, errors.New("invalid test data")
113		}
114
115		hexBytes := strings.Fields(line)
116		for _, hexByte := range hexBytes {
117			val, err := strconv.ParseUint(hexByte, 16, 8)
118			if err != nil {
119				return nil, errors.New("invalid hex byte in test data: " + err.Error())
120			}
121			currentFlow = append(currentFlow, byte(val))
122		}
123	}
124
125	if len(currentFlow) > 0 {
126		flows = append(flows, currentFlow)
127	}
128
129	return flows, nil
130}
131