From 6ce4580299c048fc5e6627f937fffb5ab46b4ccf Mon Sep 17 00:00:00 2001
From: Michael Wilson <mike@mdwn.dev>
Date: Wed, 29 Mar 2023 14:13:48 -0400
Subject: [PATCH] Support app servers on different types of tunnels. (#23749)

* Support app servers on different types of tunnels.

`AppServer`s can now be accessed through different types of tunnels, not just
on `AppTunnel`s. The first additional tunnel type is the `OktaTunnel` type.
This tunnel type has been added to the reverse tunnel. The proxy app transport
has been adjusted to use the `GetTunnelType` from the `AppServer` instead of
`AppTunnel` directly.

* Update api/types/appserver.go

Co-authored-by: Tiago Silva <tiago.silva@goteleport.com>

---------

Co-authored-by: Tiago Silva <tiago.silva@goteleport.com>
---
 api/types/appserver.go      | 12 ++++++++
 api/types/appserver_test.go | 56 +++++++++++++++++++++++++++++++++++++
 api/types/constants.go      |  3 ++
 lib/reversetunnel/srv.go    |  2 ++
 lib/web/app/transport.go    |  2 +-
 5 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 api/types/appserver_test.go

diff --git a/api/types/appserver.go b/api/types/appserver.go
index fb09c6a90bf..f4b6d802de7 100644
--- a/api/types/appserver.go
+++ b/api/types/appserver.go
@@ -51,6 +51,8 @@ type AppServer interface {
 	GetApp() Application
 	// SetApp sets the app this app server proxies.
 	SetApp(Application) error
+	// GetTunnelType returns the tunnel type associated with the app server.
+	GetTunnelType() TunnelType
 	// ProxiedService provides common methods for a proxied service.
 	ProxiedService
 }
@@ -178,6 +180,16 @@ func (s *AppServerV3) SetApp(app Application) error {
 	return nil
 }
 
+// GetTunnelType returns the tunnel type associated with the app server.
+func (s *AppServerV3) GetTunnelType() TunnelType {
+	switch {
+	case s.Origin() == OriginOkta:
+		return OktaTunnel
+	default:
+		return AppTunnel
+	}
+}
+
 // String returns the server string representation.
 func (s *AppServerV3) String() string {
 	return fmt.Sprintf("AppServer(Name=%v, Version=%v, Hostname=%v, HostID=%v, App=%v)",
diff --git a/api/types/appserver_test.go b/api/types/appserver_test.go
new file mode 100644
index 00000000000..d725bd0b0ef
--- /dev/null
+++ b/api/types/appserver_test.go
@@ -0,0 +1,56 @@
+/*
+Copyright 2023 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetTunnelType(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		appServer AppServer
+		expected  TunnelType
+	}{
+		{
+			name:      "default",
+			appServer: &AppServerV3{},
+			expected:  AppTunnel,
+		},
+		{
+			name: "okta",
+			appServer: &AppServerV3{
+				Metadata: Metadata{
+					Labels: map[string]string{
+						OriginLabel: OriginOkta,
+					},
+				},
+			},
+			expected: OktaTunnel,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			require.Equal(t, test.expected, test.appServer.GetTunnelType())
+		})
+	}
+}
diff --git a/api/types/constants.go b/api/types/constants.go
index 293d9eaf8f1..bebf263c332 100644
--- a/api/types/constants.go
+++ b/api/types/constants.go
@@ -526,6 +526,9 @@ const (
 
 	// WindowsDesktopTunnel is a tunnel where the Windows desktop service dials back to the proxy.
 	WindowsDesktopTunnel TunnelType = "windows_desktop"
+
+	// OktaTunnel is a tunnel where the Okta service dials back to the proxy.
+	OktaTunnel TunnelType = "okta"
 )
 
 type TunnelStrategyType string
diff --git a/lib/reversetunnel/srv.go b/lib/reversetunnel/srv.go
index 282f409c732..2c0471a310e 100644
--- a/lib/reversetunnel/srv.go
+++ b/lib/reversetunnel/srv.go
@@ -729,6 +729,8 @@ func (s *server) handleHeartbeat(conn net.Conn, sconn *ssh.ServerConn, nch ssh.N
 		s.handleNewCluster(conn, sconn, nch)
 	case types.RoleWindowsDesktop:
 		s.handleNewService(role, conn, sconn, nch, types.WindowsDesktopTunnel)
+	case types.RoleOkta:
+		s.handleNewService(role, conn, sconn, nch, types.OktaTunnel)
 	// Unknown role.
 	default:
 		s.log.Errorf("Unsupported role attempting to connect: %v", val)
diff --git a/lib/web/app/transport.go b/lib/web/app/transport.go
index 2ce9cfb1a2e..cf0c5fbc0dc 100644
--- a/lib/web/app/transport.go
+++ b/lib/web/app/transport.go
@@ -296,7 +296,7 @@ func dialAppServer(ctx context.Context, proxyClient reversetunnel.Tunnel, cluste
 		To:                    &utils.NetAddr{AddrNetwork: "tcp", Addr: reversetunnel.LocalNode},
 		OriginalClientDstAddr: originalDst,
 		ServerID:              fmt.Sprintf("%v.%v", server.GetHostID(), clusterName),
-		ConnType:              types.AppTunnel,
+		ConnType:              server.GetTunnelType(),
 		ProxyIDs:              server.GetProxyIDs(),
 	})
 	return conn, trace.Wrap(err)