server_impl.go

package sshserver

import (
	"fmt"
	"net"
	"strings"
	"sync"

	"github.com/containerssh/log"
	"github.com/containerssh/service"
	"golang.org/x/crypto/ssh"
)

type server struct {
	cfg                 Config
	logger              log.Logger
	handler             Handler
	listenSocket        net.Listener
	wg                  *sync.WaitGroup
	lock                *sync.Mutex
	clientSockets       map[*ssh.ServerConn]bool
	nextGlobalRequestID uint64
	nextChannelID       uint64
	hostKeys            []ssh.Signer
}

func (s *server) String() string {
	return "SSH server"
}

func (s *server) RunWithLifecycle(lifecycle service.Lifecycle) error {
	s.lock.Lock()
	alreadyRunning := false
	if s.listenSocket != nil {
		alreadyRunning = true
	} else {
		s.clientSockets = make(map[*ssh.ServerConn]bool)
	}
	s.lock.Unlock()

	if alreadyRunning {
		return fmt.Errorf("SSH server is already running")
	}

	listenConfig := net.ListenConfig{
		Control: s.socketControl,
	}

	netListener, err := listenConfig.Listen(lifecycle.Context(), "tcp", s.cfg.Listen)
	if err != nil {
		return fmt.Errorf("failed to start SSH server on %s (%w)", s.cfg.Listen, err)
	}
	s.listenSocket = netListener
	if err := s.handler.OnReady(); err != nil {
		if err := netListener.Close(); err != nil {
			s.logger.Warningf("failed to close listen socket after failed startup (%v)", err)
		}
		return err
	}
	lifecycle.Running()
	s.logger.Debugf("SSH server running on %s", s.cfg.Listen)

	go func() {
		<-lifecycle.Context().Done()
		if err := s.listenSocket.Close(); err != nil {
			s.logger.Errorf("failed to close listen socket (%v)", err)
		}
	}()
	for {
		tcpConn, err := netListener.Accept()
		if err != nil {
			// Assume listen socket closed
			break
		}
		go s.handleConnection(tcpConn)
	}
	lifecycle.Stopping()
	allClientsExited := make(chan struct{})
	shutdownHandlerExited := make(chan struct{}, 1)
	go s.disconnectClients(lifecycle, allClientsExited)
	go s.shutdownHandler(lifecycle, shutdownHandlerExited)

	s.wg.Wait()
	close(allClientsExited)
	<-shutdownHandlerExited
	return nil
}

func (s *server) disconnectClients(lifecycle service.Lifecycle, allClientsExited chan struct{}) {
	select {
	case <-allClientsExited:
		return
	case <-lifecycle.ShutdownContext().Done():
	}

	s.lock.Lock()
	for serverSocket := range s.clientSockets {
		if err := serverSocket.Close(); err != nil {
			s.logger.Debugf(
				"failed to close client socket for %s (%v)",
				serverSocket.RemoteAddr().String(),
				err,
			)
		}
	}
	s.clientSockets = map[*ssh.ServerConn]bool{}
	s.lock.Unlock()
}

func (s *server) createPasswordAuthenticator(
	handlerNetworkConnection NetworkConnectionHandler,
) func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
	return func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
		authResponse, err := handlerNetworkConnection.OnAuthPassword(conn.User(), password)
		switch authResponse {
		case AuthResponseSuccess:
			return &ssh.Permissions{}, nil
		case AuthResponseFailure:
			return nil, fmt.Errorf("authentication failed")
		case AuthResponseUnavailable:
			s.logger.Warningf("authentication backend currently unavailable (%v)", err)
			return nil, fmt.Errorf("authentication currently unavailable")
		}
		return nil, fmt.Errorf("authentication currently unavailable")
	}
}

func (s *server) createPubKeyAuthenticator(
	handlerNetworkConnection NetworkConnectionHandler,
) func(conn ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
	return func(conn ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
		authorizedKey := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pubKey)))
		authResponse, err := handlerNetworkConnection.OnAuthPubKey(conn.User(), authorizedKey)
		switch authResponse {
		case AuthResponseSuccess:
			return &ssh.Permissions{}, nil
		case AuthResponseFailure:
			return nil, fmt.Errorf("authentication failed")
		case AuthResponseUnavailable:
			s.logger.Warningf("authentication backend currently unavailable (%v)", err)
			return nil, fmt.Errorf("authentication currently unavailable")
		}
		return nil, fmt.Errorf("authentication currently unavailable")
	}
}

func (s *server) createConfiguration(handlerNetworkConnection NetworkConnectionHandler) *ssh.ServerConfig {
	serverConfig := &ssh.ServerConfig{
		Config: ssh.Config{
			KeyExchanges: s.cfg.getKex(),
			Ciphers:      s.cfg.getCiphers(),
			MACs:         s.cfg.getMACs(),
		},
		NoClientAuth:      false,
		MaxAuthTries:      6,
		PasswordCallback:  s.createPasswordAuthenticator(handlerNetworkConnection),
		PublicKeyCallback: s.createPubKeyAuthenticator(handlerNetworkConnection),
		ServerVersion:     s.cfg.ServerVersion,
		BannerCallback:    func(conn ssh.ConnMetadata) string { return s.cfg.Banner },
	}
	for _, key := range s.hostKeys {
		serverConfig.AddHostKey(key)
	}
	return serverConfig
}

func (s *server) handleConnection(conn net.Conn) {
	addr := conn.RemoteAddr().(*net.TCPAddr)
	handlerNetworkConnection, err := s.handler.OnNetworkConnection(*addr, GenerateConnectionID())
	if err != nil {
		s.logger.Infoe(err)
		_ = conn.Close()
		return
	}
	s.logger.Debugf("client connected: %s", addr.IP.String())

	sshConn, channels, globalRequests, err := ssh.NewServerConn(conn, s.createConfiguration(handlerNetworkConnection))
	if err != nil {
		s.logger.Debugf("SSH handshake failed for %s (%v)", addr.IP.String(), err)
		handlerNetworkConnection.OnHandshakeFailed(err)
		handlerNetworkConnection.OnDisconnect()
		_ = conn.Close()
		return
	}
	s.lock.Lock()
	s.clientSockets[sshConn] = true
	s.lock.Unlock()
	s.wg.Add(1)
	go func() {
		_ = sshConn.Wait()
		s.logger.Debugf("client disconnected: %s", addr.IP.String())
		handlerNetworkConnection.OnDisconnect()
		s.wg.Done()
	}()
	handlerSSHConnection, failureReason := handlerNetworkConnection.OnHandshakeSuccess(sshConn.User())
	if failureReason != nil {
		s.logger.Debugf("handshake success handler returned with error (%v)", err)
		// No need to close connection, already closed.
		return
	}
	go s.handleChannels(channels, handlerSSHConnection)
	go s.handleGlobalRequests(globalRequests, handlerSSHConnection)
}

func (s *server) handleGlobalRequests(requests <-chan *ssh.Request, connection SSHConnectionHandler) {
	for {
		request, ok := <-requests
		if !ok {
			break
		}
		requestID := s.nextGlobalRequestID
		s.nextGlobalRequestID++
		connection.OnUnsupportedGlobalRequest(requestID, request.Type, request.Payload)
		if request.WantReply {
			if err := request.Reply(false, []byte("request type not supported")); err != nil {
				s.logger.Debugf("failed to send reply to global request type %s (%v)", request.Type, err)
			}
		}
	}
}

func (s *server) handleChannels(channels <-chan ssh.NewChannel, connection SSHConnectionHandler) {
	for {
		newChannel, ok := <-channels
		if !ok {
			break
		}
		channelID := s.nextChannelID
		s.nextChannelID++
		if newChannel.ChannelType() != "session" {
			connection.OnUnsupportedChannel(channelID, newChannel.ChannelType(), newChannel.ExtraData())
			if err := newChannel.Reject(ssh.UnknownChannelType, "unsupported channel type"); err != nil {
				s.logger.Debugf("failed to send channel rejection for channel type %s", newChannel.ChannelType())
			}
			continue
		}
		go s.handleSessionChannel(channelID, newChannel, connection)
	}
}

type envRequestPayload struct {
	Name  string
	Value string
}

type execRequestPayload struct {
	Exec string
}

type ptyRequestPayload struct {
	Term     string
	Columns  uint32
	Rows     uint32
	Width    uint32
	Height   uint32
	ModeList []byte
}

type shellRequestPayload struct {
}

type signalRequestPayload struct {
	Signal string
}

type subsystemRequestPayload struct {
	Subsystem string
}

type windowRequestPayload struct {
	Columns uint32
	Rows    uint32
	Width   uint32
	Height  uint32
}

type exitStatusPayload struct {
	exitStatus ExitStatus
}

type requestType string

const (
	requestTypeEnv       requestType = "env"
	requestTypePty       requestType = "pty"
	requestTypeShell     requestType = "shell"
	requestTypeExec      requestType = "Exec"
	requestTypeSubsystem requestType = "Subsystem"
	requestTypeWindow    requestType = "window-change"
	requestTypeSignal    requestType = "Signal"
)

func (s *server) handleSessionChannel(channelID uint64, newChannel ssh.NewChannel, connection SSHConnectionHandler) {
	handlerChannel, rejection := connection.OnSessionChannel(channelID, newChannel.ExtraData())
	if rejection != nil {
		if err := newChannel.Reject(rejection.Reason(), rejection.Message()); err != nil {
			s.logger.Debugf("failed to send session channel rejection", err)
		}
		return
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		s.logger.Debugf("failed to accept session channel (%v)", err)
		return
	}
	nextRequestID := uint64(0)
	for {
		request, ok := <-requests
		if !ok {
			break
		}
		requestID := nextRequestID
		nextRequestID++
		go s.handleChannelRequest(requestID, request, channel, handlerChannel)
	}
}

func (s *server) unmarshalEnv(request *ssh.Request) (payload envRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalPty(request *ssh.Request) (payload ptyRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalShell(request *ssh.Request) (payload shellRequestPayload, err error) {
	if len(request.Payload) != 0 {
		err = ssh.Unmarshal(request.Payload, &payload)
	}
	return payload, err
}

func (s *server) unmarshalExec(request *ssh.Request) (payload execRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalSubsystem(request *ssh.Request) (payload subsystemRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalWindow(request *ssh.Request) (payload windowRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalSignal(request *ssh.Request) (payload signalRequestPayload, err error) {
	return payload, ssh.Unmarshal(request.Payload, &payload)
}

func (s *server) unmarshalPayload(request *ssh.Request) (payload interface{}, err error) {
	switch requestType(request.Type) {
	case requestTypeEnv:
		return s.unmarshalEnv(request)
	case requestTypePty:
		return s.unmarshalPty(request)
	case requestTypeShell:
		return s.unmarshalShell(request)
	case requestTypeExec:
		return s.unmarshalExec(request)
	case requestTypeSubsystem:
		return s.unmarshalSubsystem(request)
	case requestTypeWindow:
		return s.unmarshalWindow(request)
	case requestTypeSignal:
		return s.unmarshalSignal(request)
	default:
		return nil, nil
	}
}

func (s *server) handleChannelRequest(requestID uint64, request *ssh.Request, sshChannel ssh.Channel, sessionChannel SessionChannelHandler) {
	reply := func(success bool, message string, reason error) {
		if request.WantReply {
			if err := request.Reply(
				success,
				[]byte(message),
			); err != nil {
				s.logger.Debugf("failed to send reply (%v)", err)
			}
		}
	}
	onExit := func(exitCode ExitStatus) {
		if _, err := sshChannel.SendRequest(
			"exit-status",
			false,
			ssh.Marshal(exitStatusPayload{
				exitStatus: exitCode,
			})); err != nil {
			s.logger.Debugf("failed to send exit status to client (%v)", err)
		}
		if err := sshChannel.Close(); err != nil {
			s.logger.Debugf("failed to close SSH channel (%v)", err)
		}
	}
	payload, err := s.unmarshalPayload(request)
	if payload == nil {
		sessionChannel.OnUnsupportedChannelRequest(requestID, request.Type, request.Payload)
		reply(false, fmt.Sprintf("unsupported request type: %s", request.Type), nil)
		return
	}
	if err != nil {
		s.logger.Debugf("failed to unmarshal %s request payload %v (%v)", request.Type, request.Payload, err)
		sessionChannel.OnFailedDecodeChannelRequest(requestID, request.Type, request.Payload, err)
		reply(false, "failed to unmarshal payload", nil)
		return
	}
	if err := s.handleDecodedChannelRequest(
		requestID,
		requestType(request.Type),
		payload,
		sshChannel,
		sessionChannel,
		onExit,
	); err != nil {
		reply(false, err.Error(), err)
		return
	}
	reply(true, "", nil)
}

func (s *server) handleDecodedChannelRequest(
	requestID uint64,
	requestType requestType,
	payload interface{},
	channel ssh.Channel,
	sessionChannel SessionChannelHandler,
	onExit func(exitCode ExitStatus),
) error {
	switch requestType {
	case requestTypeEnv:
		return s.onEnvRequest(requestID, sessionChannel, payload)
	case requestTypePty:
		return s.onPtyRequest(requestID, sessionChannel, payload)
	case requestTypeShell:
		return s.onShell(requestID, sessionChannel, channel, onExit)
	case requestTypeExec:
		return s.onExec(requestID, sessionChannel, payload, channel, onExit)
	case requestTypeSubsystem:
		return s.onSubsystem(requestID, sessionChannel, payload, channel, onExit)
	case requestTypeWindow:
		return s.onChannel(requestID, sessionChannel, payload)
	case requestTypeSignal:
		return s.onSignal(requestID, sessionChannel, payload)
	}
	return nil
}

func (s *server) onEnvRequest(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}) error {
	return sessionChannel.OnEnvRequest(
		requestID,
		payload.(envRequestPayload).Name,
		payload.(envRequestPayload).Value,
	)
}

func (s *server) onPtyRequest(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}) error {
	return sessionChannel.OnPtyRequest(
		requestID,
		payload.(ptyRequestPayload).Term,
		payload.(ptyRequestPayload).Columns,
		payload.(ptyRequestPayload).Rows,
		payload.(ptyRequestPayload).Width,
		payload.(ptyRequestPayload).Height,
		payload.(ptyRequestPayload).ModeList,
	)
}

func (s *server) onShell(requestID uint64, sessionChannel SessionChannelHandler, channel ssh.Channel, onExit func(exitCode ExitStatus)) error {
	return sessionChannel.OnShell(
		requestID,
		channel,
		channel,
		channel.Stderr(),
		onExit,
	)
}

func (s *server) onExec(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}, channel ssh.Channel, onExit func(exitCode ExitStatus)) error {
	return sessionChannel.OnExecRequest(
		requestID,
		payload.(execRequestPayload).Exec,
		channel,
		channel,
		channel.Stderr(),
		onExit,
	)
}

func (s *server) onSignal(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}) error {
	return sessionChannel.OnSignal(
		requestID,
		payload.(signalRequestPayload).Signal,
	)
}

func (s *server) onSubsystem(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}, channel ssh.Channel, onExit func(exitCode ExitStatus)) error {
	return sessionChannel.OnSubsystem(
		requestID,
		payload.(subsystemRequestPayload).Subsystem,
		channel,
		channel,
		channel.Stderr(),
		onExit,
	)
}

func (s *server) onChannel(requestID uint64, sessionChannel SessionChannelHandler, payload interface{}) error {
	return sessionChannel.OnWindow(
		requestID,
		payload.(windowRequestPayload).Columns,
		payload.(windowRequestPayload).Rows,
		payload.(windowRequestPayload).Width,
		payload.(windowRequestPayload).Height,
	)
}

func (s *server) shutdownHandler(lifecycle service.Lifecycle, exited chan struct{}) {
	s.handler.OnShutdown(lifecycle.ShutdownContext())
	exited <- struct{}{}
}