/
dialer.go
117 lines (98 loc) · 2.58 KB
/
dialer.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
package sshd
import (
"context"
"net"
"sync"
"time"
"github.com/dolfly/core/dialer"
md "github.com/dolfly/core/metadata"
ssh_util "github.com/dolfly/x/internal/util/ssh"
"github.com/dolfly/x/registry"
"golang.org/x/crypto/ssh"
)
func init() {
registry.DialerRegistry().Register("sshd", NewDialer)
}
type sshdDialer struct {
sessions map[string]*ssh_util.Session
sessionMutex sync.Mutex
md metadata
options dialer.Options
}
func NewDialer(opts ...dialer.Option) dialer.Dialer {
options := dialer.Options{}
for _, opt := range opts {
opt(&options)
}
return &sshdDialer{
sessions: make(map[string]*ssh_util.Session),
options: options,
}
}
func (d *sshdDialer) Init(md md.Metadata) (err error) {
if err = d.parseMetadata(md); err != nil {
return
}
return nil
}
// Multiplex implements dialer.Multiplexer interface.
func (d *sshdDialer) Multiplex() bool {
return true
}
func (d *sshdDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOption) (conn net.Conn, err error) {
d.sessionMutex.Lock()
defer d.sessionMutex.Unlock()
session, ok := d.sessions[addr]
if session != nil && session.IsClosed() {
delete(d.sessions, addr) // session is dead
ok = false
}
if !ok {
var options dialer.DialOptions
for _, opt := range opts {
opt(&options)
}
conn, err = options.NetDialer.Dial(ctx, "tcp", addr)
if err != nil {
return
}
if d.md.handshakeTimeout > 0 {
conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout))
defer conn.SetDeadline(time.Time{})
}
session, err = d.initSession(ctx, addr, conn)
if err != nil {
conn.Close()
return nil, err
}
if d.md.keepalive {
go session.Keepalive(d.md.keepaliveInterval, d.md.keepaliveTimeout, d.md.keepaliveRetries)
}
go session.Wait()
go session.WaitClose()
d.sessions[addr] = session
}
return ssh_util.NewClientConn(session), nil
}
func (d *sshdDialer) initSession(ctx context.Context, addr string, conn net.Conn) (*ssh_util.Session, error) {
config := ssh.ClientConfig{
Timeout: d.md.handshakeTimeout,
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
if d.options.Auth != nil {
config.User = d.options.Auth.Username()
if password, _ := d.options.Auth.Password(); password != "" {
config.Auth = []ssh.AuthMethod{
ssh.Password(password),
}
}
}
if d.md.signer != nil {
config.Auth = append(config.Auth, ssh.PublicKeys(d.md.signer))
}
sshConn, chans, reqs, err := ssh.NewClientConn(conn, addr, &config)
if err != nil {
return nil, err
}
return ssh_util.NewSession(conn, ssh.NewClient(sshConn, chans, reqs), d.options.Logger), nil
}