From bb39c9d49d8ea59b8fecfdd673bae142c171937b Mon Sep 17 00:00:00 2001
From: Noah Stride <noah.stride@goteleport.com>
Date: Thu, 20 Jun 2024 18:04:41 +0100
Subject: [PATCH] [v14] Fix gRPC connections being disconnected regardless of
 DisconnectCertExpiry (#43292)

* Fix gRPC connections being disconnected regardless of DisconnectCertExpiry

* Remove log lines

* Add test

* back out the changes that added Authorizer
---
 lib/auth/transport_credentials.go |  7 +++++--
 lib/authz/permissions.go          |  7 ++++++-
 lib/authz/permissions_test.go     | 22 +++++++++++++++++++++-
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/lib/auth/transport_credentials.go b/lib/auth/transport_credentials.go
index 4c357130dab77..c1bc0dc9d929d 100644
--- a/lib/auth/transport_credentials.go
+++ b/lib/auth/transport_credentials.go
@@ -179,8 +179,11 @@ func newTimeoutConn(conn net.Conn, clock clockwork.Clock, expires time.Time) (ne
 	}
 
 	return &timeoutConn{
-		Conn:  conn,
-		timer: clock.AfterFunc(expires.Sub(clock.Now()), func() { conn.Close() }),
+		Conn: conn,
+		timer: clock.AfterFunc(expires.Sub(clock.Now()), func() {
+			log.Debug("Closing gRPC connection due to certificate expiry")
+			conn.Close()
+		}),
 	}, nil
 }
 
diff --git a/lib/authz/permissions.go b/lib/authz/permissions.go
index 6368f68b651ee..ee620b0a242c4 100644
--- a/lib/authz/permissions.go
+++ b/lib/authz/permissions.go
@@ -265,7 +265,12 @@ func (c *Context) GetDisconnectCertExpiry(authPref types.AuthPreference) time.Ti
 	// See https://github.com/gravitational/teleport/issues/18544
 
 	// If the session doesn't need to be disconnected on cert expiry just return the default value.
-	if c.Checker != nil && !c.Checker.AdjustDisconnectExpiredCert(authPref.GetDisconnectExpiredCert()) {
+	disconnectExpiredCert := authPref.GetDisconnectExpiredCert()
+	if c.Checker != nil {
+		disconnectExpiredCert = c.Checker.AdjustDisconnectExpiredCert(disconnectExpiredCert)
+	}
+
+	if !disconnectExpiredCert {
 		return time.Time{}
 	}
 
diff --git a/lib/authz/permissions_test.go b/lib/authz/permissions_test.go
index 5fc5971382c55..416195ae5ed9d 100644
--- a/lib/authz/permissions_test.go
+++ b/lib/authz/permissions_test.go
@@ -51,12 +51,14 @@ func TestGetDisconnectExpiredCertFromIdentity(t *testing.T) {
 		name                    string
 		expires                 time.Time
 		previousIdentityExpires time.Time
+		checker                 services.AccessChecker
 		mfaVerified             bool
 		disconnectExpiredCert   bool
 		expected                time.Time
 	}{
 		{
 			name:                    "mfa overrides expires when set",
+			checker:                 &fakeCtxChecker{},
 			expires:                 now,
 			previousIdentityExpires: inAnHour,
 			mfaVerified:             true,
@@ -65,6 +67,7 @@ func TestGetDisconnectExpiredCertFromIdentity(t *testing.T) {
 		},
 		{
 			name:                  "expires returned when mfa unset",
+			checker:               &fakeCtxChecker{},
 			expires:               now,
 			mfaVerified:           false,
 			disconnectExpiredCert: true,
@@ -72,11 +75,28 @@ func TestGetDisconnectExpiredCertFromIdentity(t *testing.T) {
 		},
 		{
 			name:                    "unset when disconnectExpiredCert is false",
+			checker:                 &fakeCtxChecker{},
 			expires:                 now,
 			previousIdentityExpires: inAnHour,
 			mfaVerified:             true,
 			disconnectExpiredCert:   false,
 		},
+		{
+			name:                  "no expiry returned when checker nil and disconnectExpiredCert false",
+			checker:               nil,
+			expires:               now,
+			mfaVerified:           false,
+			disconnectExpiredCert: false,
+			expected:              time.Time{},
+		},
+		{
+			name:                  "expiry returned when checker nil and disconnectExpiredCert true",
+			checker:               nil,
+			expires:               now,
+			mfaVerified:           false,
+			disconnectExpiredCert: true,
+			expected:              now,
+		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			var mfaVerified string
@@ -92,7 +112,7 @@ func TestGetDisconnectExpiredCertFromIdentity(t *testing.T) {
 			authPref := types.DefaultAuthPreference()
 			authPref.SetDisconnectExpiredCert(test.disconnectExpiredCert)
 
-			ctx := Context{Checker: &fakeCtxChecker{}, Identity: WrapIdentity(identity)}
+			ctx := Context{Checker: test.checker, Identity: WrapIdentity(identity)}
 
 			got := ctx.GetDisconnectCertExpiry(authPref)
 			require.Equal(t, test.expected, got)