fix: e2e test

Author: rosstimothy

integration/kube_integration_test.go

@@ -34,7 +34,6 @@ import (
 	"os/user"
 	"strconv"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
@@ -2121,7 +2120,9 @@ func testKubeJoin(t *testing.T, suite *KubeSuite) {
 
 	// fooey
 	hostUsername := suite.me.Username
-	participantUsername := suite.me.Username + "-participant"
+	peerUsername := suite.me.Username + "-peer"
+	observer1Username := suite.me.Username + "-observer1"
+	observer2Username := suite.me.Username + "-observer2"
 	kubeGroups := []string{kube.TestImpersonationGroup}
 	kubeUsers := []string{"alice@example.com"}
 	role, err := types.NewRole("kubemaster", types.RoleSpecV6{
@@ -2152,7 +2153,9 @@ func testKubeJoin(t *testing.T, suite *KubeSuite) {
 	})
 	require.NoError(t, err)
 	teleport.AddUserWithRole(hostUsername, role)
-	teleport.AddUserWithRole(participantUsername, joinRole)
+	teleport.AddUserWithRole(peerUsername, joinRole)
+	teleport.AddUserWithRole(observer1Username, joinRole)
+	teleport.AddUserWithRole(observer2Username, joinRole)
 
 	err = teleport.CreateEx(t, nil, tconf)
 	require.NoError(t, err)
@@ -2200,30 +2203,32 @@ func testKubeJoin(t *testing.T, suite *KubeSuite) {
 	// We need to wait for the exec request to be handled here for the session to be
 	// created. Sadly though the k8s API doesn't give us much indication of when that is.
 	var session types.SessionTracker
-	require.Eventually(t, func() bool {
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
 		// We need to wait for the session to be created here. We can't use the
 		// session manager's WaitUntilExists method because it doesn't work for
 		// kubernetes sessions.
 		sessions, err := teleport.Process.GetAuthServer().GetActiveSessionTrackers(context.Background())
-		if err != nil || len(sessions) == 0 {
-			return false
+		assert.NoError(t, err)
+		if assert.Len(t, sessions, 1) {
+			session = sessions[0]
 		}
-
-		session = sessions[0]
-		return true
 	}, 10*time.Second, time.Second)
 
 	participantStdinR, participantStdinW, err := os.Pipe()
 	require.NoError(t, err)
 	participantStdoutR, participantStdoutW, err := os.Pipe()
 	require.NoError(t, err)
-	streamsMu := &sync.Mutex{}
-	streams := make([]*client.KubeSession, 0, 3)
-	observerCaptures := make([]*bytes.Buffer, 0, 2)
+
+	observerCaptures := make([]*bytes.Buffer, 2)
 	albProxy := helpers.MustStartMockALBProxy(t, teleport.Config.Proxy.WebAddr.Addr)
 
 	// join peer by KubeProxyAddr
 	group.Go(func() error {
+		defer func() {
+			// close participant stdout so that we can read it after till EOF
+			participantStdoutW.Close()
+		}()
+
 		tc, err := teleport.NewClient(helpers.ClientConfig{
 			Login:   hostUsername,
 			Cluster: helpers.Site,
@@ -2238,50 +2243,52 @@ func testKubeJoin(t *testing.T, suite *KubeSuite) {
 
 		stream, err := kubeJoin(kube.ProxyConfig{
 			T:          teleport,
-			Username:   participantUsername,
+			Username:   peerUsername,
 			KubeUsers:  kubeUsers,
 			KubeGroups: kubeGroups,
 		}, tc, session, types.SessionPeerMode)
 		if err != nil {
 			return trace.Wrap(err)
 		}
-		streamsMu.Lock()
-		streams = append(streams, stream)
-		streamsMu.Unlock()
+
 		stream.Wait()
-		// close participant stdout so that we can read it after till EOF
-		participantStdoutW.Close()
+
+		t.Cleanup(func() { _ = stream.Close() })
+
 		return nil
 	})
 
 	// join observer by WebProxyAddr
 	group.Go(func() error {
-		stream, capture := kubeJoinByWebAddr(t, teleport, participantUsername, kubeUsers, kubeGroups)
-		streamsMu.Lock()
-		streams = append(streams, stream)
-		observerCaptures = append(observerCaptures, capture)
-		streamsMu.Unlock()
+		stream, capture := kubeJoinByWebAddr(t, teleport, observer1Username, kubeUsers, kubeGroups)
+		observerCaptures[0] = capture
 		stream.Wait()
+
+		t.Cleanup(func() { _ = stream.Close() })
 		return nil
 	})
 
 	// join observer with ALPN conn upgrade
 	group.Go(func() error {
-		stream, capture := kubeJoinByALBAddr(t, teleport, participantUsername, kubeUsers, kubeGroups, albProxy.Addr().String())
-		streamsMu.Lock()
-		streams = append(streams, stream)
-		observerCaptures = append(observerCaptures, capture)
-		streamsMu.Unlock()
+		stream, capture := kubeJoinByALBAddr(t, teleport, observer2Username, kubeUsers, kubeGroups, albProxy.Addr().String())
+		observerCaptures[1] = capture
 		stream.Wait()
+
+		t.Cleanup(func() { _ = stream.Close() })
 		return nil
 	})
 
-	// We wait again for the second user to finish joining the session.
-	// We allow a bit of time to pass here to give the session manager time to recognize the
-	// new IO streams of the second client.
-	time.Sleep(time.Second * 5)
+	// Wait for all users to finish joining the session.
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		session, err := teleport.Process.GetAuthServer().GetSessionTracker(context.Background(), session.GetName())
+		if !assert.NoError(t, err) {
+			return
+		}
+
+		assert.Len(t, session.GetParticipants(), 4)
+	}, 30*time.Second, 500*time.Millisecond)
 
-	// sent a test message from the participant
+	// send a test message from the participant
 	participantStdinW.Write([]byte("\ahi from peer\n\r"))
 
 	// lets type "echo hi" followed by "enter" and then "exit" + "enter":
@@ -2306,8 +2313,8 @@ func testKubeJoin(t *testing.T, suite *KubeSuite) {
 
 	// Verify observers.
 	for _, capture := range observerCaptures {
-		require.Contains(t, capture.String(), "hi from peer")
-		require.Contains(t, capture.String(), "hi from term")
+		assert.Contains(t, capture.String(), "hi from peer")
+		assert.Contains(t, capture.String(), "hi from term")
 	}
 }
 

lib/client/kubesession.go

@@ -49,7 +49,7 @@ type KubeSession struct {
 }
 
 // NewKubeSession joins a live kubernetes session.
-func NewKubeSession(ctx context.Context, tc *TeleportClient, meta types.SessionTracker, tlsServer string, mode types.SessionParticipantMode, tlsConfig *tls.Config) (*KubeSession, error) {
+func NewKubeSession(ctx context.Context, tc *TeleportClient, meta types.SessionTracker, tlsServer string, mode types.SessionParticipantMode, tlsConfig *tls.Config) (_ *KubeSession, err error) {
 	ctx, cancel := context.WithCancel(ctx)
 	joinEndpoint := "wss://" + tc.KubeProxyAddr + "/api/v1/teleport/join/" + meta.GetSessionID()
 
@@ -86,6 +86,15 @@ func NewKubeSession(ctx context.Context, tc *TeleportClient, meta types.SessionT
 
 		return nil, trace.BadParameter("failed to decode remote error: %v", string(body))
 	}
+	defer func() {
+		if err == nil {
+			return
+		}
+
+		if err := ws.Close(); err != nil {
+			log.DebugContext(ctx, "Close stream in response to context termination", "error", err)
+		}
+	}()
 
 	stream, err := streamproto.NewSessionStream(ws, streamproto.ClientHandshake{Mode: mode})
 	if err != nil {
@@ -94,7 +103,9 @@ func NewKubeSession(ctx context.Context, tc *TeleportClient, meta types.SessionT
 	}
 
 	context.AfterFunc(ctx, func() {
-		_ = stream.Close()
+		log.DebugContext(ctx, "Closing stream in response to context termination")
+		err := stream.Close()
+		log.DebugContext(ctx, "Closed stream in response to context termination", "error", err)
 	})
 
 	term, err := terminal.New(tc.Stdin, tc.Stdout, tc.Stderr)