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
|
package handlers
import (
"context"
"encoding/base64"
"runtime"
"testing"
"time"
"github.com/mimecast/dtail/internal"
"github.com/mimecast/dtail/internal/lcontext"
"github.com/mimecast/dtail/internal/protocol"
)
// TestHandleCommandCancelsContextAfterCommandFinished verifies that
// baseHandler.handleCommand no longer discards the cancel func returned by
// newCommandContext. Pre-fix the cancel was dropped, so the per-command
// context (and the watcher goroutine spawned by newCommandContext) leaked
// for the lifetime of the SSH session. The cancel must fire when
// commandFinished is invoked.
func TestHandleCommandCancelsContextAfterCommandFinished(t *testing.T) {
resetServerLogger(t)
handler := newSessionTestHandler("handle-command-cancel-user")
readServerMessage(t, handler.serverMessages)
handler.handleCommandCb = handler.handleUserCommand
type captured struct {
ctx context.Context
finish func()
}
ch := make(chan captured, 1)
handler.commands = map[string]commandHandler{
// AUTHKEY is a side-effect command so commandFinished does not
// trigger handler shutdown, keeping the test focused on the
// per-command cancel contract.
"AUTHKEY": func(ctx context.Context, _ lcontext.LContext, _ int, _ []string, commandFinished func()) {
ch <- captured{ctx: ctx, finish: commandFinished}
},
}
encoded := base64.StdEncoding.EncodeToString([]byte("AUTHKEY dummy"))
handler.handleCommand("protocol " + protocol.ProtocolCompat + " base64 " + encoded)
var got captured
select {
case got = <-ch:
case <-time.After(500 * time.Millisecond):
t.Fatal("AUTHKEY command was not dispatched")
}
select {
case <-got.ctx.Done():
t.Fatal("per-command context cancelled before commandFinished ran")
default:
}
got.finish()
select {
case <-got.ctx.Done():
case <-time.After(500 * time.Millisecond):
t.Fatal("per-command context was not cancelled after commandFinished ran; cancel was discarded")
}
}
// TestNewCommandContextReleasesWatcherGoroutine ensures the watcher
// goroutine spawned by newCommandContext exits promptly once either the
// per-command cancel fires or the handler is shut down. This is the
// defensive safety net that keeps a leak from accumulating even if a
// future caller forgets to invoke cancel.
func TestNewCommandContextReleasesWatcherGoroutine(t *testing.T) {
h := &baseHandler{done: internal.NewDone()}
t.Cleanup(h.done.Shutdown)
// Warm up so any lazily-started runtime goroutines are already up.
_, cancel := h.newCommandContext(context.Background())
cancel()
time.Sleep(20 * time.Millisecond)
baseline := runtime.NumGoroutine()
const N = 100
for i := 0; i < N; i++ {
_, cancel := h.newCommandContext(context.Background())
cancel()
}
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if delta := runtime.NumGoroutine() - baseline; delta <= 4 {
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("watcher goroutines leaked: delta=%d (expected <= 4)", runtime.NumGoroutine()-baseline)
}
// TestNewCommandContextHandlerShutdownReleasesWatcher verifies the
// defensive safety net: if a caller forgets to cancel a per-command
// context, shutting down the handler still drains the watcher goroutine
// rather than leaving it blocked until process exit.
func TestNewCommandContextHandlerShutdownReleasesWatcher(t *testing.T) {
h := &baseHandler{done: internal.NewDone()}
// Warm up.
_, cancel := h.newCommandContext(context.Background())
cancel()
time.Sleep(20 * time.Millisecond)
baseline := runtime.NumGoroutine()
const N = 50
ctxs := make([]context.Context, 0, N)
for i := 0; i < N; i++ {
ctx, _ := h.newCommandContext(context.Background())
ctxs = append(ctxs, ctx)
}
if delta := runtime.NumGoroutine() - baseline; delta < N/2 {
t.Fatalf("expected goroutines to accumulate before shutdown, delta=%d", delta)
}
h.done.Shutdown()
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if delta := runtime.NumGoroutine() - baseline; delta <= 4 {
// All watcher goroutines should have observed the contexts
// being cancelled via the defensive done.Done() branch.
for _, ctx := range ctxs {
select {
case <-ctx.Done():
default:
t.Fatalf("context not cancelled by handler shutdown")
}
}
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("watcher goroutines leaked past shutdown: delta=%d", runtime.NumGoroutine()-baseline)
}
|
