1// Copyright 2014 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package main
6
7import (
8	"encoding/binary"
9	"fmt"
10	"io"
11	"net"
12	"time"
13)
14
15// opcodePacket signals a packet, encoded with a 32-bit length prefix, followed
16// by the payload.
17const opcodePacket = byte('P')
18
19// opcodeTimeout signals a read timeout, encoded by a 64-bit number of
20// nanoseconds. On receipt, the peer should reply with
21// opcodeTimeoutAck. opcodeTimeout may only be sent by the Go side.
22const opcodeTimeout = byte('T')
23
24// opcodeTimeoutAck acknowledges a read timeout. This opcode has no payload and
25// may only be sent by the C side. Timeout ACKs act as a synchronization point
26// at the timeout, to bracket one flight of messages from C.
27const opcodeTimeoutAck = byte('t')
28
29type packetAdaptor struct {
30	net.Conn
31}
32
33// newPacketAdaptor wraps a reliable streaming net.Conn into a reliable
34// packet-based net.Conn. The stream contains packets and control commands,
35// distinguished by a one byte opcode.
36func newPacketAdaptor(conn net.Conn) *packetAdaptor {
37	return &packetAdaptor{conn}
38}
39
40func (p *packetAdaptor) readOpcode() (byte, error) {
41	out := make([]byte, 1)
42	if _, err := io.ReadFull(p.Conn, out); err != nil {
43		return 0, err
44	}
45	return out[0], nil
46}
47
48func (p *packetAdaptor) readPacketBody() ([]byte, error) {
49	var length uint32
50	if err := binary.Read(p.Conn, binary.BigEndian, &length); err != nil {
51		return nil, err
52	}
53	out := make([]byte, length)
54	if _, err := io.ReadFull(p.Conn, out); err != nil {
55		return nil, err
56	}
57	return out, nil
58}
59
60func (p *packetAdaptor) Read(b []byte) (int, error) {
61	opcode, err := p.readOpcode()
62	if err != nil {
63		return 0, err
64	}
65	if opcode != opcodePacket {
66		return 0, fmt.Errorf("unexpected opcode '%d'", opcode)
67	}
68	out, err := p.readPacketBody()
69	if err != nil {
70		return 0, err
71	}
72	return copy(b, out), nil
73}
74
75func (p *packetAdaptor) Write(b []byte) (int, error) {
76	payload := make([]byte, 1+4+len(b))
77	payload[0] = opcodePacket
78	binary.BigEndian.PutUint32(payload[1:5], uint32(len(b)))
79	copy(payload[5:], b)
80	if _, err := p.Conn.Write(payload); err != nil {
81		return 0, err
82	}
83	return len(b), nil
84}
85
86// SendReadTimeout instructs the peer to simulate a read timeout. It then waits
87// for acknowledgement of the timeout, buffering any packets received since
88// then. The packets are then returned.
89func (p *packetAdaptor) SendReadTimeout(d time.Duration) ([][]byte, error) {
90	payload := make([]byte, 1+8)
91	payload[0] = opcodeTimeout
92	binary.BigEndian.PutUint64(payload[1:], uint64(d.Nanoseconds()))
93	if _, err := p.Conn.Write(payload); err != nil {
94		return nil, err
95	}
96
97	var packets [][]byte
98	for {
99		opcode, err := p.readOpcode()
100		if err != nil {
101			return nil, err
102		}
103		switch opcode {
104		case opcodeTimeoutAck:
105			// Done! Return the packets buffered and continue.
106			return packets, nil
107		case opcodePacket:
108			// Buffer the packet for the caller to process.
109			packet, err := p.readPacketBody()
110			if err != nil {
111				return nil, err
112			}
113			packets = append(packets, packet)
114		default:
115			return nil, fmt.Errorf("unexpected opcode '%d'", opcode)
116		}
117	}
118}
119
120type replayAdaptor struct {
121	net.Conn
122	prevWrite []byte
123}
124
125// newReplayAdaptor wraps a packeted net.Conn. It transforms it into
126// one which, after writing a packet, always replays the previous
127// write.
128func newReplayAdaptor(conn net.Conn) net.Conn {
129	return &replayAdaptor{Conn: conn}
130}
131
132func (r *replayAdaptor) Write(b []byte) (int, error) {
133	n, err := r.Conn.Write(b)
134
135	// Replay the previous packet and save the current one to
136	// replay next.
137	if r.prevWrite != nil {
138		r.Conn.Write(r.prevWrite)
139	}
140	r.prevWrite = append(r.prevWrite[:0], b...)
141
142	return n, err
143}
144
145type damageAdaptor struct {
146	net.Conn
147	damage bool
148}
149
150// newDamageAdaptor wraps a packeted net.Conn. It transforms it into one which
151// optionally damages the final byte of every Write() call.
152func newDamageAdaptor(conn net.Conn) *damageAdaptor {
153	return &damageAdaptor{Conn: conn}
154}
155
156func (d *damageAdaptor) setDamage(damage bool) {
157	d.damage = damage
158}
159
160func (d *damageAdaptor) Write(b []byte) (int, error) {
161	if d.damage && len(b) > 0 {
162		b = append([]byte{}, b...)
163		b[len(b)-1]++
164	}
165	return d.Conn.Write(b)
166}
167