starssh/state.go

package starssh

import (
	"errors"

	"golang.org/x/crypto/ssh"
)

var errSSHClientClosing = errors.New("ssh client is closing")

type sshClientRequester interface {
	SendRequest(name string, wantReply bool, payload []byte) (bool, []byte, error)
	Close() error
}

var closeSSHClient = func(client sshClientRequester) error {
	if client == nil {
		return nil
	}
	return client.Close()
}

func (s *StarSSH) snapshotSSHClient() *ssh.Client {
	if s == nil {
		return nil
	}

	s.stateMu.RLock()
	defer s.stateMu.RUnlock()
	return s.Client
}

func (s *StarSSH) requireSSHClient() (*ssh.Client, error) {
	if s == nil {
		return nil, errors.New("ssh client is nil")
	}
	if s.closing.Load() {
		return nil, errSSHClientClosing
	}
	client := s.snapshotSSHClient()
	if client == nil {
		return nil, errors.New("ssh client is nil")
	}
	if s.closing.Load() {
		return nil, errSSHClientClosing
	}
	return client, nil
}

func (s *StarSSH) setTransport(client *ssh.Client, upstream *StarSSH) {
	if s == nil {
		return
	}

	s.stateMu.Lock()
	defer s.stateMu.Unlock()
	s.Client = client
	s.upstream = upstream
	s.online = client != nil
	s.closing.Store(false)
}

func (s *StarSSH) detachTransport() (*ssh.Client, *StarSSH) {
	if s == nil {
		return nil, nil
	}

	s.stateMu.Lock()
	defer s.stateMu.Unlock()

	client := s.Client
	upstream := s.upstream
	s.Client = nil
	s.upstream = nil
	s.online = false
	return client, upstream
}

func (s *StarSSH) takeKeepaliveHandles() (chan struct{}, chan struct{}) {
	if s == nil {
		return nil, nil
	}

	s.keepaliveMu.Lock()
	defer s.keepaliveMu.Unlock()

	stop := s.keepaliveStop
	done := s.keepaliveDone
	s.keepaliveStop = nil
	s.keepaliveDone = nil
	return stop, done
}

func (s *StarSSH) closeTransport(waitKeepalive bool) error {
	if s == nil {
		return nil
	}

	s.closing.Store(true)
	_ = s.closeReusableSFTPClient()
	agentForwarder := s.takeAgentForwarder()

	client, upstream := s.detachTransport()
	stop, done := s.takeKeepaliveHandles()
	if stop != nil {
		close(stop)
	}

	var closeErr error
	if agentForwarder != nil {
		closeErr = normalizeAlreadyClosedError(agentForwarder.Close())
	}
	if client != nil {
		if err := normalizeAlreadyClosedError(closeSSHClient(client)); closeErr == nil {
			closeErr = err
		}
	}
	if waitKeepalive && done != nil {
		<-done
	}
	if upstreamErr := closeUpstream(upstream); closeErr == nil {
		closeErr = upstreamErr
	}
	return closeErr
}

func (s *StarSSH) canAttachAgentForwarder(client *ssh.Client) bool {
	if s == nil || client == nil || s.closing.Load() {
		return false
	}

	s.stateMu.RLock()
	defer s.stateMu.RUnlock()
	return !s.closing.Load() && s.Client == client
}