starssh/keepalive.go

package starssh

import (
	"context"
	"sync"
	"time"
)

var sendKeepAliveRequest = func(ctx context.Context, client sshClientRequester) error {
	_, _, err := client.SendRequest("keepalive@openssh.com", true, nil)
	return err
}

func (s *StarSSH) Ping() error {
	return s.PingContext(context.Background())
}

func (s *StarSSH) PingContext(ctx context.Context) error {
	return s.pingContext(ctx, nil)
}

func (s *StarSSH) PingContextCloseOnCancel(ctx context.Context) error {
	return s.pingContext(ctx, s.Close)
}

func (s *StarSSH) pingContext(ctx context.Context, onCancel func() error) error {
	if ctx == nil {
		ctx = context.Background()
	}

	client, err := s.requireSSHClient()
	if err != nil {
		return err
	}

	runCancel := func() {}
	if onCancel != nil {
		var cancelOnce sync.Once
		runCancel = func() {
			cancelOnce.Do(func() {
				_ = onCancel()
			})
		}

		cancelDone := make(chan struct{})
		defer close(cancelDone)
		go func() {
			select {
			case <-ctx.Done():
				runCancel()
			case <-cancelDone:
			}
		}()
	}

	requestFunc := sendKeepAliveRequest
	pingErr := make(chan error, 1)
	go func() {
		err := requestFunc(ctx, client)
		select {
		case pingErr <- err:
		default:
		}
	}()

	select {
	case err := <-pingErr:
		return err
	case <-ctx.Done():
		runCancel()
		return ctx.Err()
	}
}

func (s *StarSSH) startAutoKeepAlive() {
	if s == nil || s.snapshotSSHClient() == nil {
		return
	}

	interval := s.LoginInfo.KeepAliveInterval
	if interval <= 0 {
		return
	}
	timeout := s.LoginInfo.KeepAliveTimeout
	if timeout <= 0 {
		timeout = defaultKeepAliveTimeout
	}

	stop := make(chan struct{})
	done := make(chan struct{})

	s.keepaliveMu.Lock()
	if s.keepaliveStop != nil {
		s.keepaliveMu.Unlock()
		return
	}
	s.keepaliveStop = stop
	s.keepaliveDone = done
	s.keepaliveMu.Unlock()

	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		defer close(done)

		for {
			select {
			case <-stop:
				return
			case <-ticker.C:
				pingCtx, cancel := context.WithTimeout(context.Background(), timeout)
				err := s.pingContext(pingCtx, nil)
				cancel()
				if err != nil {
					s.closeFromKeepAlive()
					return
				}
			}
		}
	}()
}

func (s *StarSSH) stopAutoKeepAlive() {
	stop, done := s.takeKeepaliveHandles()
	if stop != nil {
		close(stop)
	}
	if done != nil {
		<-done
	}
}

func (s *StarSSH) closeFromKeepAlive() {
	_ = s.closeTransport(false)
}