metricsintegration/handler.go at main · ContainerSSH/metricsintegration · GitHub

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
package metricsintegration

import (
	"context"
	"net"
	"sync"

	"github.com/containerssh/metrics"
	"github.com/containerssh/sshserver"
)

type metricsHandler struct {
	backend                   sshserver.Handler
	metricsCollector          metrics.Collector
	connectionsMetric         metrics.SimpleGeoCounter
	handshakeSuccessfulMetric metrics.SimpleGeoCounter
	handshakeFailedMetric     metrics.SimpleGeoCounter
	currentConnectionsMetric  metrics.SimpleGeoGauge
}

func (m *metricsHandler) OnReady() error {
	return m.backend.OnReady()
}

func (m *metricsHandler) OnShutdown(shutdownContext context.Context) {
	m.backend.OnShutdown(shutdownContext)
}

func (m *metricsHandler) OnNetworkConnection(
	client net.TCPAddr,
	connectionID string,
) (sshserver.NetworkConnectionHandler, error) {

	networkBackend, err := m.backend.OnNetworkConnection(client, connectionID)
	if err != nil {
		return networkBackend, err
	}
	m.connectionsMetric.Increment(client.IP)
	m.currentConnectionsMetric.Increment(client.IP)
	return &metricsNetworkHandler{
		client:  client,
		backend: networkBackend,
		handler: m,
		lock:    &sync.Mutex{},
	}, nil
}

type metricsNetworkHandler struct {
	backend      sshserver.NetworkConnectionHandler
	client       net.TCPAddr
	handler      *metricsHandler
	lock         *sync.Mutex
	disconnected bool
}

func (m *metricsNetworkHandler) OnShutdown(shutdownContext context.Context) {
	m.backend.OnShutdown(shutdownContext)
}

func (m *metricsNetworkHandler) OnAuthPassword(username string, password []byte) (
	response sshserver.AuthResponse,
	reason error,
) {
	return m.backend.OnAuthPassword(username, password)
}

func (m *metricsNetworkHandler) OnAuthPubKey(username string, pubKey string) (
	response sshserver.AuthResponse,
	reason error,
) {
	return m.backend.OnAuthPubKey(username, pubKey)
}

func (m *metricsNetworkHandler) OnAuthKeyboardInteractive(
	user string,
	challenge func(
		instruction string,
		questions sshserver.KeyboardInteractiveQuestions,
	) (answers sshserver.KeyboardInteractiveAnswers, err error),
) (response sshserver.AuthResponse, reason error) {
	return m.backend.OnAuthKeyboardInteractive(user, challenge)
}

func (m *metricsNetworkHandler) OnHandshakeFailed(reason error) {
	m.handler.handshakeFailedMetric.Increment(m.client.IP)
	m.backend.OnHandshakeFailed(reason)
}

func (m *metricsNetworkHandler) OnHandshakeSuccess(username string) (
	connection sshserver.SSHConnectionHandler,
	failureReason error,
) {
	connectionHandler, failureReason := m.backend.OnHandshakeSuccess(username)
	if failureReason != nil {
		m.handler.handshakeFailedMetric.Increment(m.client.IP)
		return connectionHandler, failureReason
	}
	m.handler.handshakeSuccessfulMetric.Increment(m.client.IP)
	return connectionHandler, failureReason
}

func (m *metricsNetworkHandler) OnDisconnect() {
	m.lock.Lock()
	defer m.lock.Unlock()
	if !m.disconnected {
		m.handler.currentConnectionsMetric.Decrement(m.client.IP)
		m.disconnected = true
	}
	m.backend.OnDisconnect()
}