/
session.go
126 lines (109 loc) · 2.24 KB
/
session.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
114
115
116
117
118
119
120
121
122
123
124
125
126
package ssh
import (
"context"
"net"
"time"
"github.com/jxo-me/netx/core/logger"
"golang.org/x/crypto/ssh"
)
const (
defaultKeepaliveInterval = 30 * time.Second
defaultKeepaliveTimeout = 15 * time.Second
defaultkeepaliveRetries = 1
)
type Session struct {
net.Conn
client *ssh.Client
closed chan struct{}
dead chan struct{}
log logger.ILogger
}
func NewSession(c net.Conn, client *ssh.Client, log logger.ILogger) *Session {
return &Session{
Conn: c,
client: client,
closed: make(chan struct{}),
dead: make(chan struct{}),
log: log,
}
}
func (s *Session) OpenChannel(name string) (ssh.Channel, <-chan *ssh.Request, error) {
return s.client.OpenChannel(name, nil)
}
func (s *Session) IsClosed() bool {
select {
case <-s.dead:
return true
case <-s.closed:
return true
default:
}
return false
}
func (s *Session) Wait() error {
defer close(s.closed)
return s.client.Wait()
}
func (s *Session) WaitClose() {
defer s.client.Close()
select {
case <-s.dead:
s.log.Debugf("session is dead")
case <-s.closed:
s.log.Debugf("session is closed")
}
}
func (s *Session) Keepalive(interval, timeout time.Duration, retries int) {
if interval <= 0 {
interval = defaultKeepaliveInterval
}
if timeout <= 0 {
timeout = defaultKeepaliveTimeout
}
if retries <= 0 {
retries = defaultkeepaliveRetries
}
s.log.Debugf("keepalive is enabled, interval: %v, timeout: %v, retries: %d", interval, timeout, retries)
defer close(s.dead)
t := time.NewTicker(interval)
defer t.Stop()
count := retries
for {
select {
case <-t.C:
start := time.Now()
err := func() error {
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
select {
case err := <-s.ping():
return err
case <-ctx.Done():
return ctx.Err()
}
}()
if err != nil {
s.log.Debugf("ssh ping: %v", err)
count--
if count == 0 {
return
}
continue
}
s.log.Debugf("ssh ping OK, RTT: %v", time.Since(start))
count = retries
case <-s.closed:
return
}
}
}
func (s *Session) ping() <-chan error {
ch := make(chan error, 1)
go func() {
defer close(ch)
if _, _, err := s.client.SendRequest("ping", true, nil); err != nil {
ch <- err
}
}()
return ch
}