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
|
package handlers
import (
"encoding/base64"
"fmt"
"io"
"strconv"
"strings"
"time"
"github.com/mimecast/dtail/internal"
"github.com/mimecast/dtail/internal/io/logger"
"github.com/mimecast/dtail/internal/version"
)
type baseHandler struct {
done *internal.Done
server string
shellStarted bool
commands chan string
receiveBuf []byte
status int
}
func (h *baseHandler) Server() string {
return h.server
}
func (h *baseHandler) Status() int {
return h.status
}
func (h *baseHandler) Done() <-chan struct{} {
return h.done.Done()
}
func (h *baseHandler) Shutdown() {
h.done.Shutdown()
}
// SendMessage to the server.
func (h *baseHandler) SendMessage(command string) error {
encoded := base64.StdEncoding.EncodeToString([]byte(command))
logger.Debug("Sending command", h.server, command, encoded)
select {
case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded):
case <-time.After(time.Second * 5):
return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded)
case <-h.Done():
return nil
}
return nil
}
// Read data from the dtail server via Writer interface.
func (h *baseHandler) Write(p []byte) (n int, err error) {
for _, b := range p {
h.receiveBuf = append(h.receiveBuf, b)
if b == '\n' {
if len(h.receiveBuf) == 0 {
continue
}
message := string(h.receiveBuf)
h.handleMessageType(message)
}
}
return len(p), nil
}
// Send data to the dtail server via Reader interface.
func (h *baseHandler) Read(p []byte) (n int, err error) {
select {
case command := <-h.commands:
n = copy(p, []byte(command))
case <-h.Done():
return 0, io.EOF
}
return
}
// Handle various message types.
func (h *baseHandler) handleMessageType(message string) {
if len(h.receiveBuf) == 0 {
return
}
// Hidden server commands starti with a dot "."
if h.receiveBuf[0] == '.' {
h.handleHiddenMessage(message)
h.receiveBuf = h.receiveBuf[:0]
return
}
logger.Raw(message)
h.receiveBuf = h.receiveBuf[:0]
}
// Handle messages received from server which are not meant to be displayed
// to the end user.
func (h *baseHandler) handleHiddenMessage(message string) {
switch {
case strings.HasPrefix(message, ".syn close connection"):
h.SendMessage(".ack close connection")
select {
case <-time.After(time.Second * 5):
logger.Debug("Shutting down client after timeout and sending ack to server")
h.Shutdown()
case <-h.Done():
return
}
case strings.HasPrefix(message, ".run exitstatus"):
splitted := strings.Split(strings.TrimSuffix(message, "\n"), " ")
if len(splitted) != 3 {
logger.Error("Unable to retrieve exitstatus", message)
return
}
i, err := strconv.Atoi(splitted[2])
if err != nil {
logger.Error("Unable to retrieve exitstatus", message, err)
return
}
logger.Debug("Retrieved exitstatus", h.status)
if i > h.status {
h.status = i
}
}
}
|
