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
package proxy
import (
"errors"
"sync/atomic"
"testing"
"time"
)
func TestProxyBidirectionalClosesBothSidesOnCopyError(t *testing.T) {
t.Parallel()
client := newControlledNetConn(errors.New("client read failed"))
upstream := newControlledNetConn(nil)
done := make(chan struct{})
go func() {
proxyBidirectional(client, upstream, 0)
close(done)
}()
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("proxyBidirectional did not close the blocked peer after copy error")
}
select {
case <-upstream.closed:
default:
t.Fatal("expected upstream connection to be closed")
}
}
func TestTunnelWrappersForwardHalfCloseAndReleaseOnClose(t *testing.T) {
t.Parallel()
tests := []struct {
name string
conn func(*fakeHalfCloseConn, *atomic.Int32) halfCloseTestConn
}{
{
name: "pooled",
conn: func(underlying *fakeHalfCloseConn, releaseCount *atomic.Int32) halfCloseTestConn {
return &pooledTunnelConn{
Conn: underlying,
release: func() {
releaseCount.Add(1)
},
}
},
},
{
name: "managed",
conn: func(underlying *fakeHalfCloseConn, releaseCount *atomic.Int32) halfCloseTestConn {
return &managedTunnelConn{
Conn: underlying,
release: func() {
releaseCount.Add(1)
},
}
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
underlying := &fakeHalfCloseConn{}
var releaseCount atomic.Int32
conn := tc.conn(underlying, &releaseCount)
assertHalfCloseForwarding(t, conn, underlying, &releaseCount)
})
}
}
func assertHalfCloseForwarding(t *testing.T, conn halfCloseTestConn, underlying *fakeHalfCloseConn, releaseCount *atomic.Int32) {
t.Helper()
if err := conn.CloseWrite(); err != nil {
t.Fatalf("CloseWrite returned error: %v", err)
}
if got := underlying.closeWriteCount.Load(); got != 1 {
t.Fatalf("expected underlying CloseWrite once, got %d", got)
}
if err := conn.CloseRead(); err != nil {
t.Fatalf("CloseRead returned error: %v", err)
}
if got := underlying.closeReadCount.Load(); got != 1 {
t.Fatalf("expected underlying CloseRead once, got %d", got)
}
if err := conn.Close(); err != nil {
t.Fatalf("Close returned error: %v", err)
}
if err := conn.Close(); err != nil {
t.Fatalf("second Close returned error: %v", err)
}
if got := releaseCount.Load(); got != 1 {
t.Fatalf("expected pooled session release once, got %d", got)
}
}
arblant