starssh/session.go

package starssh

import (
	"errors"
	"io"
	"net"
	"strings"

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

var newSSHSession = func(client *ssh.Client) (*ssh.Session, error) {
	return client.NewSession()
}

var requestSessionPTY = func(session *ssh.Session, config TerminalConfig) error {
	return session.RequestPty(config.Term, config.Rows, config.Columns, config.Modes)
}

func (s *StarSSH) Close() error {
	return s.closeTransport(true)
}

func (s *StarSSH) NewSession() (*ssh.Session, error) {
	return s.NewPTYSession(nil)
}

func (s *StarSSH) NewExecSession() (*ssh.Session, error) {
	client, err := s.requireSSHClient()
	if err != nil {
		return nil, err
	}
	session, err := NewExecSession(client)
	if err != nil {
		return nil, err
	}
	if err := s.maybeRequestAgentForwarding(session); err != nil {
		_ = session.Close()
		return nil, err
	}
	return session, nil
}

func (s *StarSSH) NewPTYSession(config *TerminalConfig) (*ssh.Session, error) {
	client, err := s.requireSSHClient()
	if err != nil {
		return nil, err
	}
	session, err := NewPTYSession(client, config)
	if err != nil {
		return nil, err
	}
	if err := s.maybeRequestAgentForwarding(session); err != nil {
		_ = session.Close()
		return nil, err
	}
	return session, nil
}

func NewTransferSession(client *ssh.Client) (*ssh.Session, error) {
	return NewExecSession(client)
}

func NewExecSession(client *ssh.Client) (*ssh.Session, error) {
	if client == nil {
		return nil, errors.New("ssh client is nil")
	}
	return newSSHSession(client)
}

func NewSession(client *ssh.Client) (*ssh.Session, error) {
	return NewPTYSession(client, nil)
}

func NewPTYSession(client *ssh.Client, config *TerminalConfig) (*ssh.Session, error) {
	if client == nil {
		return nil, errors.New("ssh client is nil")
	}

	session, err := newSSHSession(client)
	if err != nil {
		return nil, err
	}

	cfg := normalizeTerminalConfig(config)
	if err := requestSessionPTY(session, cfg); err != nil {
		_ = session.Close()
		return nil, err
	}

	return session, nil
}

func normalizeTerminalConfig(config *TerminalConfig) TerminalConfig {
	cfg := TerminalConfig{
		Term:    defaultPTYTerm,
		Rows:    defaultPTYRows,
		Columns: defaultPTYColumns,
		Modes: ssh.TerminalModes{
			ssh.ECHO:          1,
			ssh.TTY_OP_ISPEED: 14400,
			ssh.TTY_OP_OSPEED: 14400,
		},
	}
	if config == nil {
		return cfg
	}

	if strings.TrimSpace(config.Term) != "" {
		cfg.Term = config.Term
	}
	if config.Rows > 0 {
		cfg.Rows = config.Rows
	}
	if config.Columns > 0 {
		cfg.Columns = config.Columns
	}
	if len(config.Modes) > 0 {
		cfg.Modes = config.Modes
	}
	return cfg
}

func normalizeAlreadyClosedError(err error) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, net.ErrClosed) {
		return nil
	}
	if errors.Is(err, io.EOF) {
		return nil
	}
	if strings.Contains(err.Error(), "use of closed network connection") {
		return nil
	}
	return err
}

func closeUpstream(upstream *StarSSH) error {
	if upstream == nil {
		return nil
	}
	return upstream.Close()
}