go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add keepalive for ssh

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

31
dialer/ssh/conn.go
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()
}

82
dialer/ssh/dialer.go
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
}

21
dialer/ssh/metadata.go
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
}

31
dialer/sshd/conn.go
View File

@ -1,31 +0,0 @@
package sshd
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()
}

85
dialer/sshd/dialer.go
View File

@ -2,7 +2,6 @@ package sshd
import (
"context"
"errors"
"net"
"sync"
"time"
@ -19,7 +18,7 @@ func init() {
}
type sshdDialer struct {
sessions map[string]*sshSession
sessions map[string]*ssh_util.Session
sessionMutex sync.Mutex
md metadata
options dialer.Options
@ -32,7 +31,7 @@ func NewDialer(opts ...dialer.Option) dialer.Dialer {
}
return &sshdDialer{
sessions: make(map[string]*sshSession),
sessions: make(map[string]*ssh_util.Session),
options: options,
}
}
@ -70,68 +69,31 @@ func (d *sshdDialer) Dial(ctx context.Context, addr string, opts ...dialer.DialO
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{})
}
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 session.conn, err
return ssh_util.NewClientConn(session), nil
}
// Handshake implements dialer.Handshaker
func (d *sshdDialer) 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{})
}
log := d.options.Logger
session, ok := d.sessions[opts.Addr]
if session != nil && session.conn != conn {
err := errors.New("ssh: unrecognized connection")
log.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)
if err != nil {
log.Error(err)
conn.Close()
delete(d.sessions, opts.Addr)
return nil, err
}
session = s
go func() {
s.wait()
log.Debug("session closed")
}()
d.sessions[opts.Addr] = session
}
if session.IsClosed() {
delete(d.sessions, opts.Addr)
return nil, ssh_util.ErrSessionDead
}
return ssh_util.NewClientConn(session.conn, session.client), nil
}
func (d *sshdDialer) initSession(ctx context.Context, addr string, conn net.Conn) (*sshSession, error) {
func (d *sshdDialer) initSession(ctx context.Context, addr string, conn net.Conn) (*ssh_util.Session, error) {
config := ssh.ClientConfig{
// Timeout: timeout,
Timeout: d.md.handshakeTimeout,
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
if d.options.Auth != nil {
@ -151,10 +113,5 @@ func (d *sshdDialer) 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
}

13
dialer/sshd/metadata.go
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 *sshdDialer) parseMetadata(md mdata.Metadata) (err error) {
@ -40,5 +44,10 @@ func (d *sshdDialer) 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
}

14
dialer/ws/dialer.go
View File

@ -105,21 +105,21 @@ func (d *wsDialer) Handshake(ctx context.Context, conn net.Conn, options ...dial
cc := ws_util.Conn(c)
if d.md.keepAlive > 0 {
c.SetReadDeadline(time.Now().Add(d.md.keepAlive * 2))
if d.md.keepaliveInterval > 0 {
c.SetReadDeadline(time.Now().Add(d.md.keepaliveInterval * 2))
c.SetPongHandler(func(string) error {
c.SetReadDeadline(time.Now().Add(d.md.keepAlive * 2))
d.options.Logger.Debugf("pong: set read deadline: %v", d.md.keepAlive*2)
c.SetReadDeadline(time.Now().Add(d.md.keepaliveInterval * 2))
d.options.Logger.Debugf("pong: set read deadline: %v", d.md.keepaliveInterval*2)
return nil
})
go d.keepAlive(cc)
go d.keepalive(cc)
}
return cc, nil
}
func (d *wsDialer) keepAlive(conn ws_util.WebsocketConn) {
ticker := time.NewTicker(d.md.keepAlive)
func (d *wsDialer) keepalive(conn ws_util.WebsocketConn) {
ticker := time.NewTicker(d.md.keepaliveInterval)
defer ticker.Stop()
for range ticker.C {

16
dialer/ws/metadata.go
View File

@ -23,8 +23,8 @@ type metadata struct {
writeBufferSize int
enableCompression bool
header http.Header
keepAlive time.Duration
header http.Header
keepaliveInterval time.Duration
}
func (d *wsDialer) parseMetadata(md mdata.Metadata) (err error) {
@ -38,9 +38,7 @@ func (d *wsDialer) parseMetadata(md mdata.Metadata) (err error) {
writeBufferSize = "writeBufferSize"
enableCompression = "enableCompression"
header = "header"
keepAlive = "keepAlive"
keepAlivePeriod = "ttl"
header = "header"
)
d.md.host = mdutil.GetString(md, host)
@ -64,10 +62,10 @@ func (d *wsDialer) parseMetadata(md mdata.Metadata) (err error) {
d.md.header = h
}
if mdutil.GetBool(md, keepAlive) {
d.md.keepAlive = mdutil.GetDuration(md, keepAlivePeriod)
if d.md.keepAlive <= 0 {
d.md.keepAlive = defaultKeepAlivePeriod
if mdutil.GetBool(md, "keepalive") {
d.md.keepaliveInterval = mdutil.GetDuration(md, "ttl", "keepalive.interval")
if d.md.keepaliveInterval <= 0 {
d.md.keepaliveInterval = defaultKeepAlivePeriod
}
}