forked from jumpserver/koko
/
sshclients.go
113 lines (99 loc) · 2.36 KB
/
sshclients.go
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
package srvconn
import (
"strings"
"sync"
"time"
"github.com/eeeeethan2333/koko/pkg/logger"
)
type UserSSHClient struct {
ID string // userID_assetID_systemUserID_systemUsername
clients map[*sshClient]int64
mu sync.Mutex
}
func (u *UserSSHClient) AddClient(client *sshClient) {
u.mu.Lock()
defer u.mu.Unlock()
u.clients[client] = time.Now().UnixNano()
logger.Infof("Store new client(%s) remain %d", client, len(u.clients))
}
func (u *UserSSHClient) DeleteClient(client *sshClient) {
u.mu.Lock()
defer u.mu.Unlock()
delete(u.clients, client)
logger.Infof("Remove client(%s) remain %d", client, len(u.clients))
}
func (u *UserSSHClient) GetClient() *sshClient {
u.mu.Lock()
defer u.mu.Unlock()
if len(u.clients) == 0 {
return nil
}
var client *sshClient
var latest int64
for item, timestamp := range u.clients {
if timestamp > latest {
latest = timestamp
client = item
}
}
return client
}
func (u *UserSSHClient) count() int {
u.mu.Lock()
defer u.mu.Unlock()
return len(u.clients)
}
type SSHManager struct {
data map[string]*UserSSHClient
mu sync.Mutex
}
func (s *SSHManager) getClientFromCache(key string) (*sshClient, bool) {
s.mu.Lock()
defer s.mu.Unlock()
if userClient, ok := s.data[key]; ok {
client := userClient.GetClient()
if client != nil {
return client, true
}
}
return nil, false
}
func (s *SSHManager) AddClientCache(key string, client *sshClient) {
s.mu.Lock()
defer s.mu.Unlock()
if userClient, ok := s.data[key]; ok {
userClient.AddClient(client)
} else {
userClient = &UserSSHClient{
ID: key,
clients: make(map[*sshClient]int64),
}
userClient.AddClient(client)
s.data[key] = userClient
logger.Infof("Add new user cache current count: %d", len(s.data))
}
}
func (s *SSHManager) deleteClientFromCache(key string, client *sshClient) {
s.mu.Lock()
defer s.mu.Unlock()
if userClient, ok := s.data[key]; ok {
userClient.DeleteClient(client)
if userClient.count() == 0 {
delete(s.data, key)
logger.Infof("Delete user cache current count: %d", len(s.data))
}
}
}
func (s *SSHManager) searchSSHClientFromCache(prefixKey string) (client *sshClient, ok bool) {
s.mu.Lock()
defer s.mu.Unlock()
for key, userClient := range s.data {
if strings.HasPrefix(key, prefixKey) {
client := userClient.GetClient()
if client != nil {
return client, true
}
}
}
return nil, false
}