VoiletSshTools/smoketest/sshd/main.go

// Command sshd is a miniature SSH server used only to smoke-test the sshtools
// clients locally. It is not a real sshd: one hard-coded account, commands run
// through /bin/sh, no PTY and no port forwarding.
//
// It prints the address it listens on and the fingerprint of its host key on
// stdout, so a test script can drive it without hard-coded ports.
package main

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"net"
	"os"
	"os/exec"
	"syscall"

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

type exitStatusMsg struct{ Status uint32 }

type exitSignalMsg struct {
	Signal     string
	CoreDumped bool
	Errmsg     string
	Lang       string
}

var signalNames = map[syscall.Signal]string{
	syscall.SIGHUP: "HUP", syscall.SIGINT: "INT", syscall.SIGQUIT: "QUIT",
	syscall.SIGILL: "ILL", syscall.SIGABRT: "ABRT", syscall.SIGFPE: "FPE",
	syscall.SIGKILL: "KILL", syscall.SIGSEGV: "SEGV", syscall.SIGPIPE: "PIPE",
	syscall.SIGALRM: "ALRM", syscall.SIGTERM: "TERM", syscall.SIGUSR1: "USR1",
	syscall.SIGUSR2: "USR2",
}

func main() {
	addr := flag.String("addr", "127.0.0.1:2222", "listen address")
	user := flag.String("user", "tester", "accepted username")
	pass := flag.String("pass", "s3cr3t", "accepted password")
	auth := flag.String("auth", "both", "offered auth methods: both|password|keyboard-interactive|publickey")
	authorizedKey := flag.String("authorized-key", "", "public key file to accept; enables public key auth")
	hostKeyFile := flag.String("hostkey", "", "PEM file holding the host key (generated when missing)")
	knownHostsOut := flag.String("knownhosts-out", "", "write a known_hosts line for this server here")
	flag.Parse()

	config := &ssh.ServerConfig{}
	switch *auth {
	case "both":
		config.PasswordCallback = passwordCallback(*user, *pass)
		config.KeyboardInteractiveCallback = kbdCallback(*user, *pass)
	case "password":
		config.PasswordCallback = passwordCallback(*user, *pass)
	case "keyboard-interactive":
		config.KeyboardInteractiveCallback = kbdCallback(*user, *pass)
	case "publickey":
		// installed below
	default:
		fatalf("unknown -auth %q", *auth)
	}
	if *authorizedKey != "" {
		callback, err := publicKeyCallback(*user, *authorizedKey)
		if err != nil {
			fatalf("authorized key: %v", err)
		}
		config.PublicKeyCallback = callback
	} else if *auth == "publickey" {
		fatalf("-auth publickey needs -authorized-key")
	}

	signer, err := hostSigner(*hostKeyFile)
	if err != nil {
		fatalf("host key: %v", err)
	}
	config.AddHostKey(signer)

	listener, err := net.Listen("tcp", *addr)
	if err != nil {
		fatalf("listen: %v", err)
	}
	if *knownHostsOut != "" {
		line := knownhosts.Line([]string{knownhosts.Normalize(listener.Addr().String())}, signer.PublicKey())
		if err := os.WriteFile(*knownHostsOut, []byte(line+"\n"), 0o600); err != nil {
			fatalf("write known_hosts: %v", err)
		}
	}
	// Printed last: their presence tells the test script that everything is ready.
	fmt.Printf("ADDR=%s\n", listener.Addr().String())
	fmt.Printf("HOSTKEY_SHA256=%s\n", ssh.FingerprintSHA256(signer.PublicKey()))

	for {
		conn, err := listener.Accept()
		if err != nil {
			return
		}
		go serve(conn, config)
	}
}

func passwordCallback(user, pass string) func(ssh.ConnMetadata, []byte) (*ssh.Permissions, error) {
	return func(c ssh.ConnMetadata, given []byte) (*ssh.Permissions, error) {
		if c.User() == user && string(given) == pass {
			return nil, nil
		}
		return nil, fmt.Errorf("password rejected for %q", c.User())
	}
}

func kbdCallback(user, pass string) func(ssh.ConnMetadata, ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	return func(c ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
		answers, err := challenge("", "", []string{"Password: "}, []bool{false})
		if err != nil {
			return nil, err
		}
		if c.User() == user && len(answers) == 1 && answers[0] == pass {
			return nil, nil
		}
		return nil, fmt.Errorf("keyboard-interactive rejected for %q", c.User())
	}
}

func publicKeyCallback(user, path string) (func(ssh.ConnMetadata, ssh.PublicKey) (*ssh.Permissions, error), error) {
	authorized, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	want, _, _, _, err := ssh.ParseAuthorizedKey(authorized)
	if err != nil {
		return nil, err
	}
	return func(c ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
		if c.User() == user && bytes.Equal(key.Marshal(), want.Marshal()) {
			return nil, nil
		}
		return nil, fmt.Errorf("public key rejected for %q", c.User())
	}, nil
}

func hostSigner(path string) (ssh.Signer, error) {
	if path != "" {
		pemBytes, err := os.ReadFile(path)
		switch {
		case err == nil:
			return ssh.ParsePrivateKey(pemBytes)
		case !errors.Is(err, os.ErrNotExist):
			return nil, err
		}
	}
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}
	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		return nil, err
	}
	if path != "" {
		block, err := ssh.MarshalPrivateKey(priv, "sshtools smoketest host key")
		if err != nil {
			return nil, err
		}
		if err := os.WriteFile(path, pem.EncodeToMemory(block), 0o600); err != nil {
			return nil, err
		}
	}
	return signer, nil
}

func serve(nConn net.Conn, config *ssh.ServerConfig) {
	defer nConn.Close()
	conn, chans, reqs, err := ssh.NewServerConn(nConn, config)
	if err != nil {
		fmt.Fprintf(os.Stderr, "sshd: handshake: %v\n", err)
		return
	}
	defer conn.Close()
	go ssh.DiscardRequests(reqs)
	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "only session channels are supported")
			continue
		}
		ch, requests, err := newChan.Accept()
		if err != nil {
			continue
		}
		go runSession(ch, requests)
	}
}

func runSession(ch ssh.Channel, requests <-chan *ssh.Request) {
	defer ch.Close()
	for req := range requests {
		if req.Type != "exec" {
			req.Reply(false, nil)
			continue
		}
		var payload struct{ Command string }
		if err := ssh.Unmarshal(req.Payload, &payload); err != nil {
			req.Reply(false, nil)
			return
		}
		req.Reply(true, nil)
		execCommand(ch, payload.Command)
		return
	}
}

func execCommand(ch ssh.Channel, command string) {
	cmd := exec.Command("/bin/sh", "-c", command)
	cmd.Stdin = ch
	cmd.Stdout = ch
	cmd.Stderr = ch.Stderr()
	err := cmd.Run()
	if err == nil {
		sendExitStatus(ch, 0)
		return
	}
	var exitErr *exec.ExitError
	if errors.As(err, &exitErr) {
		if ws, ok := exitErr.Sys().(syscall.WaitStatus); ok && ws.Signaled() {
			name, ok := signalNames[ws.Signal()]
			if !ok {
				name = ws.Signal().String()
			}
			_, _ = ch.SendRequest("exit-signal", false, ssh.Marshal(exitSignalMsg{Signal: name}))
			return
		}
		code := exitErr.ExitCode()
		if code < 0 {
			code = 1
		}
		sendExitStatus(ch, code)
		return
	}
	fmt.Fprintf(ch.Stderr(), "sshd: %v\n", err)
	sendExitStatus(ch, 127)
}

func sendExitStatus(ch ssh.Channel, status int) {
	_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(exitStatusMsg{Status: uint32(status)}))
}

func fatalf(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "sshd: "+format+"\n", args...)
	os.Exit(1)
}