feat: Improve impersonation logic

frontend/src/layout/navigation/ProjectNotice.tsx

@@ -1,4 +1,5 @@
 import { IconGear, IconPlus } from '@posthog/icons'
+import { Spinner } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { dayjs } from 'lib/dayjs'
 import { LemonBanner } from 'lib/lemon-ui/LemonBanner'
@@ -22,26 +23,38 @@ interface ProjectNoticeBlueprint {
     closeable?: boolean
 }
 
-function CountDown({ datetime }: { datetime: dayjs.Dayjs }): JSX.Element {
+function CountDown({ datetime, callback }: { datetime: dayjs.Dayjs; callback?: () => void }): JSX.Element {
     const [now, setNow] = useState(dayjs())
 
+    // Format the time difference as 00:00:00
+    const duration = dayjs.duration(datetime.diff(now))
+    const pastCountdown = duration.seconds() < 0
+
+    const countdown = pastCountdown
+        ? 'Expired'
+        : duration.hours() > 0
+        ? duration.format('HH:mm:ss')
+        : duration.format('mm:ss')
+
     useEffect(() => {
         const interval = setInterval(() => setNow(dayjs()), 1000)
         return () => clearInterval(interval)
     }, [])
 
-    // Format the time difference as 00:00:00
-    const duration = dayjs.duration(datetime.diff(now))
-    const countdown = duration.hours() > 0 ? duration.format('HH:mm:ss') : duration.format('mm:ss')
+    useEffect(() => {
+        if (pastCountdown) {
+            callback?.()
+        }
+    }, [pastCountdown])
 
     return <>{countdown}</>
 }
 
 export function ProjectNotice(): JSX.Element | null {
     const { projectNoticeVariant } = useValues(navigationLogic)
     const { currentOrganization } = useValues(organizationLogic)
-    const { logout } = useActions(userLogic)
-    const { user } = useValues(userLogic)
+    const { logout, loadUser } = useActions(userLogic)
+    const { user, userLoading } = useValues(userLogic)
     const { closeProjectNotice } = useActions(navigationLogic)
     const { showInviteModal } = useActions(inviteLogic)
     const { requestVerificationLink } = useActions(verifyEmailLogic)
@@ -124,7 +137,14 @@ export function ProjectNotice(): JSX.Element | null {
                     You are currently logged in as a customer.{' '}
                     {user?.is_impersonated_until && (
                         <>
-                            Expires in <CountDown datetime={dayjs(user.is_impersonated_until)} />
+                            Expires in <CountDown datetime={dayjs(user.is_impersonated_until)} callback={loadUser} />
+                            {userLoading ? (
+                                <Spinner />
+                            ) : (
+                                <Link className="ml-2" onClick={() => loadUser()}>
+                                    Refresh
+                                </Link>
+                            )}
                         </>
                     )}
                 </>

posthog/middleware.py

@@ -677,7 +677,19 @@ def get_impersonated_session_expires_at(request: HttpRequest) -> Optional[dateti
 
     init_time = get_or_set_session_cookie_created_at(request=request)
 
-    return datetime.fromtimestamp(init_time) + timedelta(seconds=settings.IMPERSONATION_TIMEOUT_SECONDS)
+    last_activity_time = request.session.get(settings.IMPERSONATION_COOKIE_LAST_ACTIVITY_KEY, init_time)
+
+    # If the last activity time is less than the idle timeout, we extend the session
+    if time.time() - last_activity_time < settings.IMPERSONATION_IDLE_TIMEOUT_SECONDS:
+        last_activity_time = request.session[settings.IMPERSONATION_COOKIE_LAST_ACTIVITY_KEY] = time.time()
+        request.session.modified = True
+
+    idle_expiry_time = datetime.fromtimestamp(last_activity_time) + timedelta(
+        seconds=settings.IMPERSONATION_IDLE_TIMEOUT_SECONDS
+    )
+    total_expiry_time = datetime.fromtimestamp(init_time) + timedelta(seconds=settings.IMPERSONATION_TIMEOUT_SECONDS)
+
+    return min(idle_expiry_time, total_expiry_time)
 
 
 class AutoLogoutImpersonateMiddleware:

posthog/settings/web.py

@@ -376,7 +376,15 @@
 # Used only to display in the UI to inform users of allowlist options
 PUBLIC_EGRESS_IP_ADDRESSES = get_list(os.getenv("PUBLIC_EGRESS_IP_ADDRESSES", ""))
 
-IMPERSONATION_TIMEOUT_SECONDS = get_from_env("IMPERSONATION_TIMEOUT_SECONDS", 15 * 60, type_cast=int)
+# The total time allowed for an impersonated session
+IMPERSONATION_TIMEOUT_SECONDS = get_from_env("IMPERSONATION_TIMEOUT_SECONDS", 60 * 60 * 2, type_cast=int)
+# The time allowed for an impersonated session to be idle before it expires
+IMPERSONATION_IDLE_TIMEOUT_SECONDS = get_from_env("IMPERSONATION_IDLE_TIMEOUT_SECONDS", 30 * 60, type_cast=int)
+# Impersonation cookie last activity key
+IMPERSONATION_COOKIE_LAST_ACTIVITY_KEY = get_from_env(
+    "IMPERSONATION_COOKIE_LAST_ACTIVITY_KEY", "impersonation_last_activity"
+)
+
 SESSION_COOKIE_CREATED_AT_KEY = get_from_env("SESSION_COOKIE_CREATED_AT_KEY", "session_created_at")
 
 PROJECT_SWITCHING_TOKEN_ALLOWLIST = get_list(os.getenv("PROJECT_SWITCHING_TOKEN_ALLOWLIST", "sTMFPsFhdP1Ssg"))

posthog/test/test_middleware.py

@@ -501,7 +501,8 @@ def test_logout(self):
         self.assertNotIn("ph_current_project_name", response.cookies)
 
 
-@override_settings(IMPERSONATION_TIMEOUT_SECONDS=30)
+@override_settings(IMPERSONATION_TIMEOUT_SECONDS=100)
+@override_settings(IMPERSONATION_IDLE_TIMEOUT_SECONDS=20)
 class TestAutoLogoutImpersonateMiddleware(APIBaseTest):
     other_user: User
 
@@ -538,21 +539,65 @@ def test_not_staff_user_cannot_login(self):
         assert response.status_code == 200
         assert self.client.get("/api/users/@me").json()["email"] == self.user.email
 
-    def test_after_timeout_api_requests_401(self):
-        now = datetime.now()
+    def test_after_idle_timeout_api_requests_401(self):
+        now = datetime(2024, 1, 1, 12, 0, 0)
         with freeze_time(now):
             self.login_as_other_user()
             res = self.client.get("/api/users/@me")
             assert res.status_code == 200
             assert res.json()["email"] == "[email protected]"
+            assert res.json()["is_impersonated_until"] == "2024-01-01T12:00:20+00:00"
             assert self.client.session.get("session_created_at") == now.timestamp()
 
-        with freeze_time(now + timedelta(seconds=10)):
+        # Move forward by 19
+        now = now + timedelta(seconds=19)
+        with freeze_time(now):
             res = self.client.get("/api/users/@me")
             assert res.status_code == 200
             assert res.json()["email"] == "[email protected]"
+            assert res.json()["is_impersonated_until"] == "2024-01-01T12:00:39+00:00"
 
-        with freeze_time(now + timedelta(seconds=35)):
+        # Past idle timeout
+        now = now + timedelta(seconds=21)
+
+        with freeze_time(now):
+            res = self.client.get("/api/users/@me")
+            assert res.status_code == 401
+
+    def test_after_total_timeout_api_requests_401(self):
+        now = datetime(2024, 1, 1, 12, 0, 0)
+        with freeze_time(now):
+            self.login_as_other_user()
+            res = self.client.get("/api/users/@me")
+            assert res.status_code == 200
+            assert res.json()["email"] == "[email protected]"
+            assert res.json()["is_impersonated_until"] == "2024-01-01T12:00:20+00:00"
+            assert self.client.session.get("session_created_at") == now.timestamp()
+
+        for _ in range(4):
+            # Move forward by 19 seconds 4 times for a total of 76 seconds
+            now = now + timedelta(seconds=19)
+            with freeze_time(now):
+                res = self.client.get("/api/users/@me")
+                assert res.status_code == 200
+                assert res.json()["email"] == "[email protected]"
+                # Format exactly like the date above
+                assert res.json()["is_impersonated_until"] == (now + timedelta(seconds=20)).strftime(
+                    "%Y-%m-%dT%H:%M:%S+00:00"
+                )
+
+        now = now + timedelta(seconds=19)
+        with freeze_time(now):
+            res = self.client.get("/api/users/@me")
+            assert res.status_code == 200
+            assert res.json()["email"] == "[email protected]"
+            # Even though below the idle timeout, we now see the total timeout as that is earlier
+            assert res.json()["is_impersonated_until"] == "2024-01-01T12:01:40+00:00"
+
+        # Now even less than the idle time will take us past the total timeout
+        now = now + timedelta(seconds=10)
+
+        with freeze_time(now):
             res = self.client.get("/api/users/@me")
             assert res.status_code == 401