starting ssh proxy test harness

internal/testenv/upstreams/ssh.go

@@ -0,0 +1,177 @@
+package upstreams
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net"
+
+	"github.com/pomerium/pomerium/internal/testenv"
+	"github.com/pomerium/pomerium/internal/testenv/values"
+	"golang.org/x/crypto/ssh"
+)
+
+type SSHUpstreamOptions struct {
+	displayName  string
+	serverConfig ssh.ServerConfig
+}
+
+type SSHUpstreamOption interface {
+	applySSH(*SSHUpstreamOptions)
+}
+
+type sshUpstreamOption func(o *SSHUpstreamOptions)
+
+func (s sshUpstreamOption) applySSH(o *SSHUpstreamOptions) { s(o) }
+
+func WithPublicKeyAuthAlgorithms(algs []string) SSHUpstreamOption {
+	return sshUpstreamOption(func(o *SSHUpstreamOptions) {
+		o.serverConfig.PublicKeyAuthAlgorithms = algs
+	})
+}
+
+func WithHostKeys(keys ...ssh.Signer) SSHUpstreamOption {
+	return sshUpstreamOption(func(o *SSHUpstreamOptions) {
+		for _, key := range keys {
+			o.serverConfig.AddHostKey(key)
+		}
+	})
+}
+
+func WithBannerCallback(c func(ssh.ConnMetadata) string) SSHUpstreamOption {
+	return sshUpstreamOption(func(o *SSHUpstreamOptions) {
+		o.serverConfig.BannerCallback = c
+	})
+}
+
+func WithPublicKeyCallback(c func(ssh.ConnMetadata, ssh.PublicKey) (*ssh.Permissions, error)) SSHUpstreamOption {
+	return sshUpstreamOption(func(o *SSHUpstreamOptions) {
+		o.serverConfig.PublicKeyCallback = c
+	})
+}
+
+type ServerConnCallback func(*ssh.ServerConn, <-chan ssh.NewChannel, <-chan *ssh.Request)
+
+var closeConnCallback ServerConnCallback = func(conn *ssh.ServerConn, ncc <-chan ssh.NewChannel, rq <-chan *ssh.Request) {
+	conn.Close()
+}
+
+// SSHUpstream represents an ssh server which can be used as the target for
+// one or more Pomerium routes in a test environment.
+//
+// Use SetServerConnCallback() to define the behavior of this server once
+// a connection is established.
+//
+// Dial() can be used to make a client-side connection through the Pomerium
+// route, while DirectDial() can be used to connect bypassing Pomerium.
+type SSHUpstream interface {
+	testenv.Upstream
+
+	SetServerConnCallback(callback ServerConnCallback)
+
+	Dial(r testenv.Route, config *ssh.ClientConfig) (*ssh.Client, error)
+	DirectDial(r testenv.Route, config *ssh.ClientConfig) (*ssh.Client, error)
+}
+
+type sshUpstream struct {
+	SSHUpstreamOptions
+	testenv.Aggregate
+	serverPort values.MutableValue[int]
+
+	// XXX: does it make sense to cache clients?
+	//clientCache sync.Map // map[testenv.Route]*ssh.Client
+
+	serverConnCallback ServerConnCallback
+}
+
+var (
+	_ testenv.Upstream = (*sshUpstream)(nil)
+	_ SSHUpstream      = (*sshUpstream)(nil)
+)
+
+// SSH creates a new ssh upstream server.
+func SSH(opts ...SSHUpstreamOption) SSHUpstream {
+	options := SSHUpstreamOptions{
+		displayName: "SSH Upstream",
+	}
+	for _, op := range opts {
+		op.applySSH(&options)
+	}
+	up := &sshUpstream{
+		SSHUpstreamOptions: options,
+		serverPort:         values.Deferred[int](),
+		serverConnCallback: closeConnCallback, // default handler, to avoid hanging connections
+	}
+	up.RecordCaller()
+	return up
+}
+
+// Port implements SSHUpstream.
+func (h *sshUpstream) Port() values.Value[int] {
+	return h.serverPort
+}
+
+// Router implements SSHUpstream.
+func (h *sshUpstream) SetServerConnCallback(callback ServerConnCallback) {
+	h.serverConnCallback = callback
+}
+
+// Route implements SSHUpstream.
+func (h *sshUpstream) Route() testenv.RouteStub {
+	r := &testenv.PolicyRoute{}
+	protocol := "ssh"
+	r.To(values.Bind(h.serverPort, func(port int) string {
+		return fmt.Sprintf("%s://127.0.0.1:%d", protocol, port)
+	}))
+	h.Add(r)
+	return r
+}
+
+// Run implements SSHUpstream.
+func (h *sshUpstream) Run(ctx context.Context) error {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		return err
+	}
+	h.serverPort.Resolve(listener.Addr().(*net.TCPAddr).Port)
+
+	go func() {
+		<-ctx.Done()
+		listener.Close()
+	}()
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			return err
+		}
+		go h.handleConnection(ctx, conn)
+	}
+}
+
+func (h *sshUpstream) handleConnection(ctx context.Context, conn net.Conn) {
+	serverConn, ncc, rc, err := ssh.NewServerConn(conn, &h.serverConfig)
+	if err != nil {
+		// XXX: figure out the right way to log this
+		log.Println("ssh connection handshake unsuccessful:", err)
+		return
+	}
+	go func() {
+		<-ctx.Done()
+		conn.Close()
+	}()
+	h.serverConnCallback(serverConn, ncc, rc)
+}
+
+// Dial implements SSHUpstream.
+func (h *sshUpstream) Dial(r testenv.Route, config *ssh.ClientConfig) (*ssh.Client, error) {
+	// XXX: need to add ssh listener configuration to Env
+	//ssh.Dial("tcp", h.Env().)
+	return nil, fmt.Errorf("not implemented")
+}
+
+// DirectDial implements SSHUpstream.
+func (h *sshUpstream) DirectDial(r testenv.Route, config *ssh.ClientConfig) (*ssh.Client, error) {
+	addr := fmt.Sprintf("127.0.0.1:%d", h.serverPort.Value())
+	return ssh.Dial("tcp", addr, config)
+}