summaryrefslogtreecommitdiff
path: root/internal/server/handlers/controlhandler.go
blob: daa98351b7538be645d4d1effa941efdf17b4db9 (plain)
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
package handlers

import (
	"context"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/mimecast/dtail/internal/io/logger"
	user "github.com/mimecast/dtail/internal/user/server"
)

// ControlHandler is used for control functions and health monitoring.
type ControlHandler struct {
	ctx            context.Context
	done           chan struct{}
	hostname       string
	payload        []byte
	serverMessages chan string
	user           *user.User
}

// NewControlHandler returns a new control handler.
func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, <-chan struct{}) {
	logger.Debug(user, "Creating control handler")

	h := ControlHandler{
		ctx:            ctx,
		done:           make(chan struct{}),
		serverMessages: make(chan string, 10),
		user:           user,
	}

	fqdn, err := os.Hostname()
	if err != nil {
		logger.FatalExit(err)
	}

	s := strings.Split(fqdn, ".")
	h.hostname = s[0]

	return &h, h.done
}

// Read is to send data to the client via the Reader interface.
func (h *ControlHandler) Read(p []byte) (n int, err error) {
	for {
		select {
		case message := <-h.serverMessages:
			wholePayload := []byte(fmt.Sprintf("SERVER|%s|%s\n", h.hostname, message))
			n = copy(p, wholePayload)
			return
		case <-h.ctx.Done():
			return 0, io.EOF
		}
	}
}

// Write is to read data to the client via the Writer interface.
func (h *ControlHandler) Write(p []byte) (n int, err error) {
	for _, c := range p {
		switch c {
		case ';':
			wholePayload := strings.TrimSpace(string(h.payload))
			h.handleCommand(h.ctx, wholePayload)
			h.payload = nil

		default:
			h.payload = append(h.payload, c)
		}
	}

	n = len(p)
	return
}

func (h *ControlHandler) handleCommand(ctx context.Context, command string) {
	logger.Info(h.user, command)
	s := strings.Split(command, " ")
	logger.Debug(h.user, "Receiving command", command, s)

	switch s[0] {
	case "health":
		h.serverMessages <- "OK: DTail SSH Server seems fine"
		h.serverMessages <- "done;"
	case "debug":
		h.serverMessages <- logger.Debug(h.user, "Receiving debug command", command, s)
	default:
		h.serverMessages <- logger.Warn(h.user, "Received unknown control command", command, s)
	}
}