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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
|
package handlers
import (
"bytes"
"io"
"testing"
"time"
"github.com/mimecast/dtail/internal"
"github.com/mimecast/dtail/internal/config"
"github.com/mimecast/dtail/internal/io/line"
sshserver "github.com/mimecast/dtail/internal/ssh/server"
userserver "github.com/mimecast/dtail/internal/user/server"
)
// frameSizeTestServerHandler builds a minimal ServerHandler with a very small
// maxCommandFrameSize so tests can exercise the limit without allocating a
// full megabyte of data.
func frameSizeTestServerHandler(maxFrameBytes int) *ServerHandler {
u := &userserver.User{Name: "frame-size-test-user"}
h := &ServerHandler{
baseHandler: baseHandler{
done: internal.NewDone(),
lines: make(chan *line.Line, 4),
serverMessages: make(chan string, 8),
maprMessages: make(chan string, 4),
ackCloseReceived: make(chan struct{}),
user: u,
codec: newProtocolCodec(u),
maxCommandFrameSize: maxFrameBytes,
},
serverCfg: &config.ServerConfig{
AuthKeyEnabled: true,
},
authKeyStore: sshserver.NewAuthKeyStore(time.Hour, 5),
}
// Use the real user-command handler so normal protocol paths work.
h.handleCommandCb = h.handleUserCommand
h.commands = h.newCommandRegistry()
return h
}
// frameSizeTestHealthHandler builds a minimal HealthHandler with a small limit.
func frameSizeTestHealthHandler(maxFrameBytes int) *HealthHandler {
u := &userserver.User{Name: "frame-size-health-test-user"}
return &HealthHandler{
baseHandler: baseHandler{
done: internal.NewDone(),
lines: make(chan *line.Line, 4),
serverMessages: make(chan string, 8),
maprMessages: make(chan string, 4),
ackCloseReceived: make(chan struct{}),
user: u,
codec: newProtocolCodec(u),
maxCommandFrameSize: maxFrameBytes,
},
}
}
// TestWriteOversizeFrameClosesSession verifies that sending a byte stream that
// never emits a ';' delimiter and grows beyond maxCommandFrameSize causes Write
// to return io.ErrClosedPipe and shuts down the done channel so the SSH layer
// can tear down the connection.
func TestWriteOversizeFrameClosesSession(t *testing.T) {
resetServerLogger(t)
const limit = 64 // tiny limit to keep the test fast
h := frameSizeTestServerHandler(limit)
// Build a payload larger than the limit that contains no ';' delimiter.
oversizeFrame := bytes.Repeat([]byte("x"), limit+1)
_, err := h.Write(oversizeFrame)
if err != io.ErrClosedPipe {
t.Fatalf("expected io.ErrClosedPipe from Write on oversize frame, got %v", err)
}
// The done channel must be closed so callers can observe session termination.
select {
case <-h.done.Done():
// expected
default:
t.Fatalf("expected handler done channel to be closed after oversize frame")
}
}
// TestWriteOversizeFrameHealthHandlerClosesSession exercises the same limit
// through the HealthHandler, which does not hold a ServerConfig and instead
// uses the default limit (or the one injected in tests via the struct field).
func TestWriteOversizeFrameHealthHandlerClosesSession(t *testing.T) {
resetServerLogger(t)
const limit = 32
h := frameSizeTestHealthHandler(limit)
oversizeFrame := bytes.Repeat([]byte("y"), limit+1)
_, err := h.Write(oversizeFrame)
if err != io.ErrClosedPipe {
t.Fatalf("expected io.ErrClosedPipe from health handler Write on oversize frame, got %v", err)
}
select {
case <-h.done.Done():
// expected
default:
t.Fatalf("expected health handler done channel to be closed after oversize frame")
}
}
// TestWriteFrameAtExactLimitIsAccepted verifies that a frame whose length equals
// the limit is still accepted (the guard fires only when the buffer *exceeds* the
// limit). This ensures the boundary condition is correct.
func TestWriteFrameAtExactLimitIsAccepted(t *testing.T) {
resetServerLogger(t)
const limit = 16
h := frameSizeTestServerHandler(limit)
// Frame of exactly `limit` bytes followed by a ';' delimiter — must succeed.
frame := append(bytes.Repeat([]byte("z"), limit), ';')
_, err := h.Write(frame)
if err != nil {
t.Fatalf("expected no error for frame at exact limit, got %v", err)
}
// The session must still be alive.
select {
case <-h.done.Done():
t.Fatalf("expected handler to remain alive for frame at exact limit")
default:
// expected
}
}
// TestWriteNormalFramesBelowLimitAreAccepted confirms that legitimate short
// frames (well below the limit) pass through without triggering the guard.
func TestWriteNormalFramesBelowLimitAreAccepted(t *testing.T) {
resetServerLogger(t)
const limit = 1024 // generously above any test payload
h := frameSizeTestServerHandler(limit)
// Several small frames; none should trigger the limit.
for _, cmd := range []string{"health;", "health;", "health;"} {
if _, err := h.Write([]byte(cmd)); err != nil {
t.Fatalf("unexpected error writing normal frame %q: %v", cmd, err)
}
}
select {
case <-h.done.Done():
t.Fatalf("expected handler to remain alive after small frames")
default:
// expected
}
}
// TestWriteZeroLimitDisablesGuard confirms that setting maxCommandFrameSize to 0
// disables the limit entirely (the guard is not checked), so arbitrarily large
// frames are tolerated. This makes it easy to opt out of the check when needed.
func TestWriteZeroLimitDisablesGuard(t *testing.T) {
resetServerLogger(t)
const limit = 0 // disabled
h := frameSizeTestServerHandler(limit)
// Send 4 KiB without a delimiter — must not trigger the guard.
largeFrame := bytes.Repeat([]byte("a"), 4096)
if _, err := h.Write(largeFrame); err != nil {
t.Fatalf("expected no error when limit is 0 (disabled), got %v", err)
}
select {
case <-h.done.Done():
t.Fatalf("expected handler to remain alive when limit is 0")
default:
// expected
}
}
// TestDefaultMaxCommandFrameSizeMatchesServerConfig validates that the default
// value defined in the config package matches what newDefaultServerConfig
// populates into ServerConfig.MaxCommandFrameSize. This prevents them from
// drifting apart silently.
func TestDefaultMaxCommandFrameSizeMatchesServerConfig(t *testing.T) {
cfg := config.ServerConfig{}
// Retrieve through the exported helper that sets all defaults.
defaultCfg := config.NewDefaultServerConfigForTest()
cfg = *defaultCfg
if cfg.MaxCommandFrameSize != config.DefaultMaxCommandFrameSize {
t.Fatalf("ServerConfig.MaxCommandFrameSize default (%d) != config.DefaultMaxCommandFrameSize (%d)",
cfg.MaxCommandFrameSize, config.DefaultMaxCommandFrameSize)
}
}
|
