Skip to content

Commit 9f35b01

Browse files
committed
add tests for sshd package
1 parent 5d0db51 commit 9f35b01

File tree

11 files changed

+524
-101
lines changed

11 files changed

+524
-101
lines changed

Makefile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ vet:
1717
go vet $(PACKAGES)
1818

1919
test: fmt vet
20-
go test -p 2 -v $(PACKAGES)
20+
go test -timeout 1m -v $(PACKAGES) && echo "\nall ok."
2121

2222
run: all $(ID_RSA)
2323
bin/daemon -k $(ID_RSA) -d

src/daemon/agent/docker_handler.go

Lines changed: 45 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ type DockerHandler struct {
1717
exec *docker.Exec
1818
closer docker.CloseWaiter
1919
filter *payload.Request
20+
closed bool
2021
}
2122

2223
func NewDockerClient() (*docker.Client, error) {
@@ -55,8 +56,6 @@ func (h *DockerHandler) Handle(req *HandleRequest) error {
5556
return errors.New(fmt.Sprintf("Could not found container for %v", h.filter))
5657
}
5758

58-
log.Debugf("Found container %s", matched.ID[:10])
59-
6059
return h.execCommand(matched, req)
6160
}
6261

@@ -100,22 +99,11 @@ func (h *DockerHandler) execCommand(container *docker.Container, req *HandleRequ
10099
log.Debugf("Container session successfuly created %s", session.ID[:10])
101100

102101
success := make(chan struct{}, 1)
103-
started := make(chan error, 1)
104-
105-
go func() {
106-
select {
107-
case <-success:
108-
success <- struct{}{}
109-
started <- nil
110-
case <-time.After(15 * time.Second):
111-
started <- errors.New("Could not wait session within 15 seconds")
112-
}
113-
}()
114102

115103
startExecOptions := docker.StartExecOptions{
116-
InputStream: req.Reader,
117-
OutputStream: req.Writer,
118-
ErrorStream: req.Writer,
104+
InputStream: req.Stdin,
105+
OutputStream: req.Stdout,
106+
ErrorStream: req.Stderr,
119107
Detach: false,
120108
Tty: false,
121109
RawTerminal: true,
@@ -132,6 +120,18 @@ func (h *DockerHandler) execCommand(container *docker.Container, req *HandleRequ
132120
}
133121
h.closer = closer
134122

123+
started := make(chan error, 1)
124+
125+
go func() {
126+
select {
127+
case <-success:
128+
success <- struct{}{}
129+
started <- nil
130+
case <-time.After(15 * time.Second):
131+
started <- errors.New("Could not wait session within 15 seconds")
132+
}
133+
}()
134+
135135
select {
136136
case err := <-started:
137137
if err != nil {
@@ -150,26 +150,42 @@ func (h *DockerHandler) execCommand(container *docker.Container, req *HandleRequ
150150
return nil
151151
}
152152

153-
func (h *DockerHandler) Wait() error {
153+
func (h *DockerHandler) Wait() (int, error) {
154154
if h.closer != nil {
155155
log.Debug("Starting wait for container session response")
156156
if err := h.closer.Wait(); err != nil {
157-
return errors.New(fmt.Sprintf("Could wait container session (%s)", err))
157+
return 1, errors.New(fmt.Sprintf("Could wait container session (%s)", err))
158158
}
159159
}
160-
return nil
160+
161+
return h.exitCode()
162+
}
163+
164+
func (h *DockerHandler) exitCode() (int, error) {
165+
if h.exec == nil {
166+
return 1, errors.New("Exec instance is undefined")
167+
}
168+
169+
inspect, err := h.cli.InspectExec(h.exec.ID)
170+
if err != nil {
171+
return 1, errors.New(fmt.Sprintf("Could not inspect exec=%s (%s)", h.exec.ID, err))
172+
}
173+
174+
log.Debugf("Process exited with code %d", inspect.ExitCode)
175+
176+
return inspect.ExitCode, nil
161177
}
162178

163179
func (h *DockerHandler) isMatched(container *docker.Container) bool {
164-
if container.ID == h.filter.ContainerId {
165-
log.Debugf("Match container by ID=%s", container.ID)
180+
if len(h.filter.ContainerId) > 8 && strings.HasPrefix(container.ID, h.filter.ContainerId) {
181+
log.Debugf("Match container by id=%s", container.ID)
166182
return true
167183
}
168184

169185
if h.filter.ContainerEnv != "" {
170186
for _, env := range container.Config.Env {
171187
if env == h.filter.ContainerEnv {
172-
log.Debugf("Match container by env %s", env)
188+
log.Debugf("Match container by env %s id=%s", env, container.ID)
173189
return true
174190
}
175191
}
@@ -186,7 +202,7 @@ func (h *DockerHandler) isMatched(container *docker.Container) bool {
186202

187203
for name, value := range container.Config.Labels {
188204
if name == fieldName && value == fieldValue {
189-
log.Debugf("Match container by label %s=%s", name, value)
205+
log.Debugf("Match container by label %s=%s id=%s", name, value, container.ID)
190206
return true
191207
}
192208
}
@@ -209,13 +225,17 @@ func (h *DockerHandler) Resize(req *ResizeRequest) error {
209225

210226
func (h *DockerHandler) Close() error {
211227

228+
if h.closed {
229+
log.Warnf("Close session called multiple times")
230+
return nil
231+
}
232+
h.closed = true
233+
212234
if h.exec != nil && h.closer != nil {
213235
err := h.closer.Close()
214236
if err != nil {
215237
return errors.New(fmt.Sprintf("Could not close container session (%s)", err))
216238
}
217-
218-
h.exec = nil
219239
log.Info("Container session successfuly closed")
220240
}
221241

src/daemon/agent/docker_handler_test.go

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ func Test_DockerHandler_shouldSuccessfullyRunInteractiveSession(t *testing.T) {
3737

3838
handleReq := &HandleRequest{
3939
Tty: tty,
40-
Reader: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
41-
Writer: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
40+
Stdin: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
41+
Stdout: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
42+
Stderr: iotest.NewWriteLogger("[e]: ", pipe.IoWriter()),
4243
}
4344

4445
require.NoError(t, handler.Handle(handleReq))
@@ -50,7 +51,9 @@ func Test_DockerHandler_shouldSuccessfullyRunInteractiveSession(t *testing.T) {
5051

5152
require.NoError(t, pipe.WaitStringReceived("complete."))
5253
require.NoError(t, handler.Close())
53-
require.NoError(t, handler.Wait())
54+
55+
_, err := handler.Wait()
56+
require.NoError(t, err)
5457

5558
require.Contains(t, pipe.String(), "echo term is $TERM\r\n")
5659
require.Contains(t, pipe.String(), "term is xterm\r\n")
@@ -75,15 +78,18 @@ func Test_DockerHandler_shouldSuccessfullyRunNonInteractiveSession(t *testing.T)
7578
pipe := testutils.NewTestingPipe()
7679

7780
handleReq := &HandleRequest{
78-
Reader: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
79-
Writer: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
81+
Stdin: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
82+
Stdout: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
83+
Stderr: iotest.NewWriteLogger("[e]: ", pipe.IoWriter()),
8084
Exec: "ls -la ; echo complete.",
8185
}
8286

8387
require.NoError(t, handler.Handle(handleReq))
8488
require.NoError(t, pipe.WaitStringReceived("complete."))
8589
require.NoError(t, handler.Close())
86-
require.NoError(t, handler.Wait())
90+
91+
_, err := handler.Wait()
92+
require.NoError(t, err)
8793

8894
require.Contains(t, pipe.String(), ".dockerenv\n")
8995
}
@@ -102,15 +108,18 @@ func Test_DockerHandler_shouldSuccessfullyFindContainers(t *testing.T) {
102108
pipe := testutils.NewTestingPipe()
103109

104110
handleReq := &HandleRequest{
105-
Reader: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
106-
Writer: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
111+
Stdin: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
112+
Stdout: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
113+
Stderr: iotest.NewWriteLogger("[e]: ", pipe.IoWriter()),
107114
Exec: "echo complete.",
108115
}
109116

110117
require.NoError(t, handler.Handle(handleReq))
111118
require.NoError(t, pipe.WaitStringReceived("complete."))
112119
require.NoError(t, handler.Close())
113-
require.NoError(t, handler.Wait())
120+
121+
_, err := handler.Wait()
122+
require.NoError(t, err)
114123
}
115124

116125
t.Run("container.ID", func(t *testing.T) {
@@ -132,7 +141,7 @@ func Test_DockerHandler_shouldSuccessfullyFindContainers(t *testing.T) {
132141
})
133142
}
134143

135-
func Test_DockerHandler_shouldFailToHandleRequest(t *testing.T) {
144+
func Test_DockerHandler_shouldFailToHandleRequests(t *testing.T) {
136145
cli := NewTestDockerClient(t)
137146
container := NewTestDockerContainer(t, cli, "FOO=BAR", map[string]string{})
138147
defer RemoveTestDockerContainer(t, cli, container)
@@ -144,16 +153,20 @@ func Test_DockerHandler_shouldFailToHandleRequest(t *testing.T) {
144153
pipe := testutils.NewTestingPipe()
145154

146155
handleReq := &HandleRequest{
147-
Reader: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
148-
Writer: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
156+
Stdin: iotest.NewReadLogger("[r]: ", pipe.IoReader()),
157+
Stdout: iotest.NewWriteLogger("[w]: ", pipe.IoWriter()),
158+
Stderr: iotest.NewWriteLogger("[e]: ", pipe.IoWriter()),
149159
Exec: "true",
150160
}
151161

152162
err := handler.Handle(handleReq)
153163
require.Error(t, err)
154164
require.Contains(t, err.Error(), expect)
155165
require.NoError(t, handler.Close())
156-
require.NoError(t, handler.Wait())
166+
167+
code, err := handler.Wait()
168+
require.Error(t, err)
169+
require.Equal(t, 1, code)
157170
}
158171

159172
t.Run("container not found", func(t *testing.T) {

src/daemon/agent/echo_handler.go

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package agent
2+
3+
import (
4+
log "github.com/Sirupsen/logrus"
5+
"io"
6+
)
7+
8+
type EchoHandlerErrors struct {
9+
Handle error
10+
Wait error
11+
Close error
12+
}
13+
14+
type EchoHandler struct {
15+
completed chan error
16+
errors EchoHandlerErrors
17+
}
18+
19+
func NewEchoHandler(errors EchoHandlerErrors) *EchoHandler {
20+
return &EchoHandler{
21+
completed: make(chan error),
22+
errors: errors,
23+
}
24+
}
25+
26+
func (h *EchoHandler) Handle(req *HandleRequest) error {
27+
go func() {
28+
_, err := io.Copy(req.Stdout, req.Stdin)
29+
if err != nil {
30+
log.Warnf("Could not copy io streams (%s)", err)
31+
}
32+
h.completed <- err
33+
}()
34+
35+
if h.errors.Handle != nil {
36+
return h.errors.Handle
37+
}
38+
39+
if req.Exec != "" {
40+
if _, err := req.Stdout.Write([]byte(req.Exec)); err != nil {
41+
return err
42+
}
43+
}
44+
45+
return nil
46+
}
47+
48+
func (h *EchoHandler) Resize(tty *ResizeRequest) error {
49+
return nil
50+
}
51+
52+
func (h *EchoHandler) Wait() (int, error) {
53+
if h.errors.Wait != nil {
54+
return 1, h.errors.Wait
55+
}
56+
57+
select {
58+
case err := <-h.completed:
59+
if err != nil {
60+
return 1, err
61+
}
62+
}
63+
64+
return 0, nil
65+
}
66+
67+
func (h *EchoHandler) Close() error {
68+
if h.errors.Close != nil {
69+
return h.errors.Close
70+
}
71+
return nil
72+
}

src/daemon/agent/handler.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,16 @@ type TtyRequest struct {
1717

1818
type HandleRequest struct {
1919
Tty *TtyRequest
20-
Reader io.Reader
21-
Writer io.Writer
20+
Stdin io.Reader
21+
Stdout io.Writer
22+
Stderr io.Writer
2223
Exec string
2324
}
2425

2526
type Handler interface {
2627
Handle(req *HandleRequest) error
2728
Resize(tty *ResizeRequest) error
28-
Wait() error
29+
Wait() (int, error)
2930
Close() error
3031
}
3132

0 commit comments

Comments
 (0)