From 3d845a4ad1a4ac0174b71da9921e29b8e1efc526 Mon Sep 17 00:00:00 2001
From: Koichi Kato <koichi.kato@synacor.com>
Date: Wed, 11 Sep 2024 09:54:18 +0900
Subject: [PATCH] ZCS-15807: fixed 2FA trusted device handling

---
 .../src/java/com/zimbra/cs/service/account/Auth.java  | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/store/src/java/com/zimbra/cs/service/account/Auth.java b/store/src/java/com/zimbra/cs/service/account/Auth.java
index db0c99c001..f40bae0125 100644
--- a/store/src/java/com/zimbra/cs/service/account/Auth.java
+++ b/store/src/java/com/zimbra/cs/service/account/Auth.java
@@ -328,7 +328,7 @@ public Element handle(Element request, Map<String, Object> context) throws Servi
                         if (responseElement != null) {
                             return responseElement;
                         }
-                        return needTwoFactorAuth(context, request, acct, twoFactorManager, zsc, tokenType, recoveryCode);
+                        return needTwoFactorAuth(context, request, acct, twoFactorManager, zsc, tokenType, recoveryCode, trustedToken);
                     }
                 } else {
                     if (password != null || recoveryCode != null) {
@@ -448,7 +448,7 @@ private void verifyTrustedDevice(Account account, TrustedDeviceToken td, Map<Str
     }
 
     private Element needTwoFactorAuth(Map<String, Object> context, Element requestElement, Account account, TwoFactorAuth auth,
-            ZimbraSoapContext zsc, TokenType tokenType, String recoveryCode) throws ServiceException {
+            ZimbraSoapContext zsc, TokenType tokenType, String recoveryCode, TrustedDeviceToken td) throws ServiceException {
         /* two cases here:
          * 1) the user needs to provide a two-factor code.
          *    in this case, the server returns a two-factor auth token in the response header that the client
@@ -459,20 +459,23 @@ private Element needTwoFactorAuth(Map<String, Object> context, Element requestEl
         if (!auth.twoFactorAuthEnabled()) {
             throw AccountServiceException.TWO_FACTOR_SETUP_REQUIRED(getTwoFactorAuthRequiredSetupErrorMessage(account));
         } else {
+            HttpServletRequest httpReq = (HttpServletRequest) context.get(SoapServlet.SERVLET_REQUEST);
+            HttpServletResponse httpResp = (HttpServletResponse) context.get(SoapServlet.SERVLET_RESPONSE);
             Element response = zsc.createElement(AccountConstants.AUTH_RESPONSE);
             AuthToken authToken = AuthProvider.getAuthToken(account, recoveryCode != null ? Usage.RESET_PASSWORD : Usage.TWO_FACTOR_AUTH, tokenType);
             response.addUniqueElement(AccountConstants.E_TWO_FACTOR_AUTH_REQUIRED).setText("true");
             response.addAttribute(AccountConstants.E_LIFETIME, authToken.getExpires() - System.currentTimeMillis(), Element.Disposition.CONTENT);
             authToken.encodeAuthResp(response, false);
             if (recoveryCode != null) {
-                HttpServletRequest httpReq = (HttpServletRequest) context.get(SoapServlet.SERVLET_REQUEST);
-                HttpServletResponse httpResp = (HttpServletResponse) context.get(SoapServlet.SERVLET_RESPONSE);
                 boolean rememberMe = requestElement.getAttributeBool(AccountConstants.A_PERSIST_AUTH_TOKEN_COOKIE, false);
                 authToken.setIgnoreSameSite(requestElement.getAttributeBool(AccountConstants.A_IGNORE_SAME_SITE_COOKIE, false));
                 authToken.encode(httpReq, httpResp, false, ZimbraCookie.secureCookie(httpReq), rememberMe);
             }
             response.addUniqueElement(AccountConstants.E_TRUSTED_DEVICES_ENABLED).setText(account.isFeatureTrustedDevicesEnabled() ? "true" : "false");
             AccountUtil.addTwoFactorAttributes(response, account);
+            if (td != null) {
+                td.encode(httpResp, response, ZimbraCookie.secureCookie(httpReq));
+            }
             return response;
         }
     }