starssh/state.go

package starssh

import (
	"errors"

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

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) {
	client := s.snapshotSSHClient()
	if client == nil {
		return nil, errors.New("ssh client is nil")
	}
	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
}

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.closeReusableSFTPClient()

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

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