forked from mastercactapus/proxyprotocol
/
headerv2_test.go
155 lines (137 loc) · 4.44 KB
/
headerv2_test.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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package proxyprotocol
import (
"bufio"
"bytes"
"io"
"net"
"testing"
"github.com/stretchr/testify/assert"
)
func TestHeaderV2(t *testing.T) {
type section struct {
name string
value []byte
}
check := func(name string, h HeaderV2, exp []section) {
t.Run(name+"_WriteTo", func(t *testing.T) {
var buf bytes.Buffer
_, err := h.WriteTo(&buf)
assert.NoError(t, err)
for _, s := range exp {
cmp := make([]byte, len(s.value))
_, err := io.ReadFull(&buf, cmp)
assert.NoError(t, err)
assert.Equal(t, s.value, cmp, s.name)
}
})
t.Run(name+"_Parse", func(t *testing.T) {
var buf bytes.Buffer
for _, s := range exp {
buf.Write(s.value)
}
hdr, err := Parse(bufio.NewReader(&buf))
if !assert.NoError(t, err) {
return
}
assert.IsType(t, &HeaderV2{}, hdr, "Header Type")
p := hdr.(*HeaderV2)
assert.Equal(t, h.Command, p.Command, "Command")
if h.Src != nil {
assert.NotNil(t, p.Src)
assert.Equal(t, h.Src.String(), p.Src.String(), "SrcAddr")
} else {
assert.Nil(t, p.Src)
}
if h.Dest != nil {
assert.NotNil(t, p.Dest)
assert.Equal(t, h.Dest.String(), p.Dest.String(), "DestAddr")
} else {
assert.Nil(t, p.Dest)
}
})
}
check("local", HeaderV2{},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x20}}, // v2, Local
{name: "Fam/Proto", value: []byte{0x00}}, // unspec, unspec
{name: "Length", value: []byte{0, 0}}, // zero length
},
)
check("tcp-ipv4", HeaderV2{
Command: CmdProxy,
Src: &net.TCPAddr{IP: net.ParseIP("192.168.0.1"), Port: 80},
Dest: &net.TCPAddr{IP: net.ParseIP("192.168.0.2"), Port: 90},
},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x21}}, // v2, Proxy
{name: "Fam/Proto", value: []byte{0x11}}, // INET, STREAM
{name: "Length", value: []byte{0, 12}}, // length=12
{name: "SrcAddr", value: []byte{192, 168, 0, 1}},
{name: "DestAddr", value: []byte{192, 168, 0, 2}},
{name: "SrcPort", value: []byte{0, 80}},
{name: "DstPort", value: []byte{0, 90}},
},
)
check("udp-ipv4", HeaderV2{
Command: CmdProxy,
Src: &net.UDPAddr{IP: net.ParseIP("192.168.0.1"), Port: 80},
Dest: &net.UDPAddr{IP: net.ParseIP("192.168.0.2"), Port: 90},
},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x21}}, // v2, Proxy
{name: "Fam/Proto", value: []byte{0x12}}, // INET, DGRAM
{name: "Length", value: []byte{0, 12}}, // length=12
{name: "SrcAddr", value: []byte{192, 168, 0, 1}},
{name: "DestAddr", value: []byte{192, 168, 0, 2}},
{name: "SrcPort", value: []byte{0, 80}},
{name: "DstPort", value: []byte{0, 90}},
},
)
check("udp-ipv6", HeaderV2{
Command: CmdProxy,
Src: &net.UDPAddr{IP: net.ParseIP("2001::1"), Port: 80},
Dest: &net.UDPAddr{IP: net.ParseIP("2002::2"), Port: 90},
},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x21}}, // v2, Proxy
{name: "Fam/Proto", value: []byte{0x22}}, // INET6, DGRAM
{name: "Length", value: []byte{0, 36}}, // length=36
{name: "SrcAddr", value: []byte{0x20, 0x01, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x01}},
{name: "DestAddr", value: []byte{0x20, 0x02, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x02}},
{name: "SrcPort", value: []byte{0, 80}},
{name: "DstPort", value: []byte{0, 90}},
},
)
check("unixstream", HeaderV2{
Command: CmdProxy,
Src: &net.UnixAddr{Net: "unix", Name: "foo"},
Dest: &net.UnixAddr{Net: "unix", Name: "bar"},
},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x21}}, // v2, Proxy
{name: "Fam/Proto", value: []byte{0x31}}, // UNIX, STREAM
{name: "Length", value: []byte{0, 216}}, // length=216
{name: "SrcAddr", value: append([]byte("foo"), make([]byte, 105)...)},
{name: "DestAddr", value: append([]byte("bar"), make([]byte, 105)...)},
},
)
check("unixgram", HeaderV2{
Command: CmdProxy,
Src: &net.UnixAddr{Net: "unixgram", Name: "foo"},
Dest: &net.UnixAddr{Net: "unixgram", Name: "bar"},
},
[]section{
{name: "Signature", value: sigV2},
{name: "Version", value: []byte{0x21}}, // v2, Proxy
{name: "Fam/Proto", value: []byte{0x32}}, // UNIX, DGRAM
{name: "Length", value: []byte{0, 216}}, // length=216
{name: "SrcAddr", value: append([]byte("foo"), make([]byte, 105)...)},
{name: "DestAddr", value: append([]byte("bar"), make([]byte, 105)...)},
},
)
}