Add WebSocket handler for WebUI database sessions

lib/client/alpn.go

@@ -85,14 +85,33 @@ type ALPNAuthTunnelConfig struct {
 	// RouteToDatabase contains the destination server that must receive the connection.
 	// Specific for database proxying.
 	RouteToDatabase proto.RouteToDatabase
+
+	// TLSCert specifies the TLS certificate used on the proxy connection.
+	TLSCert *tls.Certificate
+}
+
+func (c *ALPNAuthTunnelConfig) CheckAndSetDefaults(ctx context.Context) error {
+	if c.AuthClient == nil {
+		return trace.BadParameter("missing auth client")
+	}
+
+	if c.TLSCert == nil {
+		tlsCert, err := getUserCerts(ctx, c.AuthClient, c.MFAResponse, c.Expires, c.RouteToDatabase, c.ConnectionDiagnosticID)
+		if err != nil {
+			return trace.BadParameter("failed to parse private key: %v", err)
+		}
+
+		c.TLSCert = &tlsCert
+	}
+
+	return nil
 }
 
 // RunALPNAuthTunnel runs a local authenticated ALPN proxy to another service.
 // At least one Route (which defines the service) must be defined
 func RunALPNAuthTunnel(ctx context.Context, cfg ALPNAuthTunnelConfig) error {
-	tlsCert, err := getUserCerts(ctx, cfg.AuthClient, cfg.MFAResponse, cfg.Expires, cfg.RouteToDatabase, cfg.ConnectionDiagnosticID)
-	if err != nil {
-		return trace.BadParameter("failed to parse private key: %v", err)
+	if err := cfg.CheckAndSetDefaults(ctx); err != nil {
+		return trace.Wrap(err)
 	}
 
 	lp, err := alpnproxy.NewLocalProxy(alpnproxy.LocalProxyConfig{
@@ -101,7 +120,7 @@ func RunALPNAuthTunnel(ctx context.Context, cfg ALPNAuthTunnelConfig) error {
 		Protocols:          []alpn.Protocol{cfg.Protocol},
 		Listener:           cfg.Listener,
 		ParentContext:      ctx,
-		Cert:               tlsCert,
+		Cert:               *cfg.TLSCert,
 	}, alpnproxy.WithALPNConnUpgradeTest(ctx, getClusterCACertPool(cfg.AuthClient)))
 	if err != nil {
 		return trace.Wrap(err)

lib/client/db/repl/repl.go

@@ -0,0 +1,91 @@
+// Teleport
+// Copyright (C) 2024 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package repl
+
+import (
+	"context"
+	"io"
+	"net"
+	"sync"
+
+	"github.com/gravitational/trace"
+
+	"github.com/gravitational/teleport/lib/tlsca"
+)
+
+// NewREPLConfig represents the database REPL constructor config.
+type NewREPLConfig struct {
+	// Client is the user terminal client.
+	Client io.ReadWriter
+	// ServerConn is the database server connection.
+	ServerConn net.Conn
+	// Route is the session routing information.
+	Route tlsca.RouteToDatabase
+}
+
+// REPLNewFunc defines the constructor function for database REPL
+// sessions.
+type REPLNewFunc func(context.Context, *NewREPLConfig) (REPLInstance, error)
+
+// REPLInstance represents a REPL instance.
+type REPLInstance interface {
+	// Run executes the REPL. This is a blocking operation.
+	Run(context.Context) error
+}
+
+// REPLGetter is an interface for retrieving REPL constructor functions given
+// the database protocol.
+type REPLGetter interface {
+	// GetREPL returns a start function for the specified protocol.
+	GetREPL(ctx context.Context, dbProtocol string) (REPLNewFunc, error)
+}
+
+// registry implements a default package level registry.
+var registry = &REPLRegistry{repl: make(map[string]REPLNewFunc)}
+
+// REPLRegistry implements a database REPL registry.
+type REPLRegistry struct {
+	mu   sync.Mutex
+	repl map[string]REPLNewFunc
+}
+
+// GetREPL implements REPLGetter.
+func (r *REPLRegistry) GetREPL(_ context.Context, dbProtocol string) (REPLNewFunc, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if startFunc, ok := r.repl[dbProtocol]; ok {
+		return startFunc, nil
+	}
+
+	return nil, trace.NotImplemented("REPL not registered for protocol %q", dbProtocol)
+}
+
+func (r *REPLRegistry) register(dbProtocol string, f REPLNewFunc) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.repl[dbProtocol] = f
+}
+
+// DefaultGetter implements a default instance of REPLGetter.
+var DefaultGetter REPLGetter = registry
+
+// RegisterREPL registers a new REPL for the database protocol.
+func RegisterREPL(dbProtocol string, f REPLNewFunc) {
+	registry.register(dbProtocol, f)
+}

lib/defaults/defaults.go

@@ -721,6 +721,10 @@ const (
 
 	// WebsocketKubeExec provides latency information for a session.
 	WebsocketKubeExec = "k"
+
+	// WebsocketDatabaseSessionRequest is received when a new database session
+	// is requested.
+	WebsocketDatabaseSessionRequest = "d"
 )
 
 // The following are cryptographic primitives Teleport does not support in

lib/service/service.go

@@ -108,6 +108,7 @@ import (
 	_ "github.com/gravitational/teleport/lib/backend/pgbk"
 	"github.com/gravitational/teleport/lib/bpf"
 	"github.com/gravitational/teleport/lib/cache"
+	dbrepl "github.com/gravitational/teleport/lib/client/db/repl"
 	"github.com/gravitational/teleport/lib/cloud"
 	"github.com/gravitational/teleport/lib/cloud/gcp"
 	"github.com/gravitational/teleport/lib/cloud/imds"
@@ -1067,6 +1068,10 @@ func NewTeleport(cfg *servicecfg.Config) (*TeleportProcess, error) {
 		cfg.PluginRegistry = plugin.NewRegistry()
 	}
 
+	if cfg.DatabaseREPLGetter == nil {
+		cfg.DatabaseREPLGetter = dbrepl.DefaultGetter
+	}
+
 	var cloudLabels labels.Importer
 
 	// Check if we're on a cloud instance, and if we should override the node's hostname.
@@ -4642,6 +4647,7 @@ func (process *TeleportProcess) initProxyEndpoint(conn *Connector) error {
 			AutomaticUpgradesChannels: cfg.Proxy.AutomaticUpgradesChannels,
 			IntegrationAppHandler:     connectionsHandler,
 			FeatureWatchInterval:      retryutils.HalfJitter(web.DefaultFeatureWatchInterval * 2),
+			DatabaseREPLGetter:        cfg.DatabaseREPLGetter,
 		}
 		webHandler, err := web.NewHandler(webConfig)
 		if err != nil {

lib/service/servicecfg/config.go

@@ -43,6 +43,7 @@ import (
 	"github.com/gravitational/teleport/lib/auth/state"
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/lite"
+	dbrepl "github.com/gravitational/teleport/lib/client/db/repl"
 	"github.com/gravitational/teleport/lib/cloud/imds"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/events"
@@ -265,6 +266,10 @@ type Config struct {
 	// AccessGraph represents AccessGraph server config
 	AccessGraph AccessGraphConfig
 
+	// DatabaseREPLGetter is used to retrieve datatabase REPL given the
+	// protocol.
+	DatabaseREPLGetter dbrepl.REPLGetter
+
 	// token is either the token needed to join the auth server, or a path pointing to a file
 	// that contains the token
 	//

lib/web/apiserver.go

@@ -80,6 +80,7 @@ import (
 	"github.com/gravitational/teleport/lib/authz"
 	"github.com/gravitational/teleport/lib/automaticupgrades"
 	"github.com/gravitational/teleport/lib/client"
+	dbrepl "github.com/gravitational/teleport/lib/client/db/repl"
 	"github.com/gravitational/teleport/lib/client/sso"
 	"github.com/gravitational/teleport/lib/defaults"
 	dtconfig "github.com/gravitational/teleport/lib/devicetrust/config"
@@ -312,6 +313,9 @@ type Config struct {
 	// FeatureWatchInterval is the interval between pings to the auth server
 	// to fetch new cluster features
 	FeatureWatchInterval time.Duration
+
+	// DatabaseREPLGetter is used for retrieving database REPL.
+	DatabaseREPLGetter dbrepl.REPLGetter
 }
 
 // SetDefaults ensures proper default values are set if
@@ -819,6 +823,7 @@ func (h *Handler) bindDefaultEndpoints() {
 	h.GET("/webapi/sites/:site/sessions", h.WithClusterAuth(h.clusterActiveAndPendingSessionsGet)) // get list of active and pending sessions
 
 	h.GET("/webapi/sites/:site/kube/exec/ws", h.WithClusterAuthWebSocket(h.podConnect)) // connect to a pod with exec (via websocket, with auth over websocket)
+	h.GET("/webapi/sites/:site/db/exec/ws", h.WithClusterAuthWebSocket(h.dbConnect))
 
 	// Audit events handlers.
 	h.GET("/webapi/sites/:site/events/search", h.WithClusterAuth(h.clusterSearchEvents))                 // search site events
@@ -3561,6 +3566,7 @@ func (h *Handler) siteNodeConnect(
 	}
 
 	term, err := NewTerminal(ctx, TerminalHandlerConfig{
+		Logger:             h.logger,
 		Term:               req.Term,
 		SessionCtx:         sessionCtx,
 		UserAuthClient:     clt,
@@ -3713,6 +3719,7 @@ func (h *Handler) podConnect(
 		ws:                  ws,
 		keepAliveInterval:   keepAliveInterval,
 		log:                 h.log.WithField(teleport.ComponentKey, "pod"),
+		logger:              h.logger.With(teleport.ComponentKey, "pod"),
 		userClient:          clt,
 		localCA:             hostCA,
 		configServerAddr:    serverAddr,

lib/web/apiserver_test.go

@@ -111,6 +111,7 @@ import (
 	"github.com/gravitational/teleport/lib/bpf"
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/client/conntest"
+	dbrepl "github.com/gravitational/teleport/lib/client/db/repl"
 	"github.com/gravitational/teleport/lib/cryptosuites"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/events"
@@ -210,6 +211,9 @@ type webSuiteConfig struct {
 
 	// clock to use for all server components
 	clock clockwork.FakeClock
+
+	// databaseREPLGetter allows setting custom database REPLs.
+	databaseREPLGetter dbrepl.REPLGetter
 }
 
 func newWebSuiteWithConfig(t *testing.T, cfg webSuiteConfig) *WebSuite {
@@ -508,6 +512,7 @@ func newWebSuiteWithConfig(t *testing.T, cfg webSuiteConfig) *WebSuite {
 			return &proxyClientCert, nil
 		},
 		IntegrationAppHandler: &mockIntegrationAppHandler{},
+		DatabaseREPLGetter:    cfg.databaseREPLGetter,
 	}
 
 	if handlerConfig.HealthCheckAppServer == nil {