add keepalive for ssh

This commit is contained in:
ginuerzh
2023-03-03 13:06:20 +08:00
parent 54046d2003
commit 7a21c7eb6f
13 changed files with 213 additions and 219 deletions

View File

@ -1,31 +0,0 @@
package ssh
import (
"net"
"golang.org/x/crypto/ssh"
)
type sshSession struct {
addr string
conn net.Conn
client *ssh.Client
closed chan struct{}
dead chan struct{}
}
func (s *sshSession) IsClosed() bool {
select {
case <-s.dead:
return true
case <-s.closed:
return true
default:
}
return false
}
func (s *sshSession) wait() error {
defer close(s.closed)
return s.client.Wait()
}

View File

@ -2,13 +2,11 @@ package ssh
import (
"context"
"errors"
"net"
"sync"
"time"
"github.com/go-gost/core/dialer"
"github.com/go-gost/core/logger"
md "github.com/go-gost/core/metadata"
ssh_util "github.com/go-gost/x/internal/util/ssh"
"github.com/go-gost/x/registry"
@ -20,9 +18,8 @@ func init() {
}
type sshDialer struct {
sessions map[string]*sshSession
sessions map[string]*ssh_util.Session
sessionMutex sync.Mutex
logger logger.Logger
md metadata
options dialer.Options
}
@ -34,8 +31,7 @@ func NewDialer(opts ...dialer.Option) dialer.Dialer {
}
return &sshDialer{
sessions: make(map[string]*sshSession),
logger: options.Logger,
sessions: make(map[string]*ssh_util.Session),
options: options,
}
}
@ -72,73 +68,36 @@ func (d *sshDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialOp
if err != nil {
return
}
session = &sshSession{
addr: addr,
conn: conn,
if d.md.handshakeTimeout > 0 {
conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout))
defer conn.SetDeadline(time.Time{})
}
d.sessions[addr] = session
}
return session.conn, err
}
// Handshake implements dialer.Handshaker
func (d *sshDialer) Handshake(ctx context.Context, conn net.Conn, options ...dialer.HandshakeOption) (net.Conn, error) {
opts := &dialer.HandshakeOptions{}
for _, option := range options {
option(opts)
}
d.sessionMutex.Lock()
defer d.sessionMutex.Unlock()
if d.md.handshakeTimeout > 0 {
conn.SetDeadline(time.Now().Add(d.md.handshakeTimeout))
defer conn.SetDeadline(time.Time{})
}
session, ok := d.sessions[opts.Addr]
if session != nil && session.conn != conn {
err := errors.New("ssh: unrecognized connection")
d.logger.Error(err)
conn.Close()
delete(d.sessions, opts.Addr)
return nil, err
}
if !ok || session.client == nil {
s, err := d.initSession(ctx, opts.Addr, conn)
session, err = d.initSession(ctx, addr, conn)
if err != nil {
d.logger.Error(err)
conn.Close()
delete(d.sessions, opts.Addr)
return nil, err
}
session = s
go func() {
s.wait()
d.logger.Debug("session closed")
}()
d.sessions[opts.Addr] = session
}
if session.IsClosed() {
delete(d.sessions, opts.Addr)
return nil, ssh_util.ErrSessionDead
}
if d.md.keepalive {
go session.Keepalive(d.md.keepaliveInterval, d.md.keepaliveTimeout, d.md.keepaliveRetries)
}
go session.Wait()
go session.WaitClose()
channel, reqs, err := session.client.OpenChannel(ssh_util.GostSSHTunnelRequest, nil)
d.sessions[addr] = session
}
channel, reqs, err := session.OpenChannel(ssh_util.GostSSHTunnelRequest)
if err != nil {
return nil, err
}
go ssh.DiscardRequests(reqs)
return ssh_util.NewConn(conn, channel), nil
return ssh_util.NewConn(session, channel), nil
}
func (d *sshDialer) initSession(ctx context.Context, addr string, conn net.Conn) (*sshSession, error) {
func (d *sshDialer) initSession(ctx context.Context, addr string, conn net.Conn) (*ssh_util.Session, error) {
config := ssh.ClientConfig{
Timeout: 30 * time.Second,
Timeout: d.md.handshakeTimeout,
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
if d.options.Auth != nil {
@ -158,10 +117,5 @@ func (d *sshDialer) initSession(ctx context.Context, addr string, conn net.Conn)
return nil, err
}
return &sshSession{
conn: conn,
client: ssh.NewClient(sshConn, chans, reqs),
closed: make(chan struct{}),
dead: make(chan struct{}),
}, nil
return ssh_util.NewSession(conn, ssh.NewClient(sshConn, chans, reqs), d.options.Logger), nil
}

View File

@ -10,8 +10,12 @@ import (
)
type metadata struct {
handshakeTimeout time.Duration
signer ssh.Signer
handshakeTimeout time.Duration
signer ssh.Signer
keepalive bool
keepaliveInterval time.Duration
keepaliveTimeout time.Duration
keepaliveRetries int
}
func (d *sshDialer) parseMetadata(md mdata.Metadata) (err error) {
@ -27,11 +31,10 @@ func (d *sshDialer) parseMetadata(md mdata.Metadata) (err error) {
return err
}
pp := mdutil.GetString(md, passphrase)
if pp == "" {
d.md.signer, err = ssh.ParsePrivateKey(data)
} else {
if pp := mdutil.GetString(md, passphrase); pp != "" {
d.md.signer, err = ssh.ParsePrivateKeyWithPassphrase(data, []byte(pp))
} else {
d.md.signer, err = ssh.ParsePrivateKey(data)
}
if err != nil {
return err
@ -40,5 +43,11 @@ func (d *sshDialer) parseMetadata(md mdata.Metadata) (err error) {
d.md.handshakeTimeout = mdutil.GetDuration(md, handshakeTimeout)
if d.md.keepalive = mdutil.GetBool(md, "keepalive"); d.md.keepalive {
d.md.keepaliveInterval = mdutil.GetDuration(md, "ttl", "keepalive.interval")
d.md.keepaliveTimeout = mdutil.GetDuration(md, "keepalive.timeout")
d.md.keepaliveRetries = mdutil.GetInt(md, "keepalive.retries")
}
return
}