From 67250bbd8d1ff05f5b391c3624a7afc0bc447917 Mon Sep 17 00:00:00 2001
From: Jaap Droogers <jaap@boerendroogers.nl>
Date: Thu, 24 Aug 2023 21:09:55 +0200
Subject: [PATCH] 1048: Updated tests for cookie.

 Make tests to work with an unvalid Base64 string and a Base64 string without encrypted content.

1048: Solved base64 problem with remember me cookie.
---
 .../web/mgt/CookieRememberMeManager.java      |  4 +-
 .../web/mgt/CookieRememberMeManagerTest.java  | 51 ++++++++++++++++---
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/web/src/main/java/org/apache/shiro/web/mgt/CookieRememberMeManager.java b/web/src/main/java/org/apache/shiro/web/mgt/CookieRememberMeManager.java
index 291751ef77..61c734daf0 100644
--- a/web/src/main/java/org/apache/shiro/web/mgt/CookieRememberMeManager.java
+++ b/web/src/main/java/org/apache/shiro/web/mgt/CookieRememberMeManager.java
@@ -251,11 +251,11 @@ protected byte[] getRememberedSerializedIdentity(SubjectContext subjectContext)
      * @param base64 the base64 encoded String that may need to be padded
      * @return the base64 String padded if necessary.
      */
-    private String ensurePadding(String base64) {
+    protected String ensurePadding(String base64) {
         int length = base64.length();
         if (length % 4 != 0) {
             StringBuilder sb = new StringBuilder(base64);
-            for (int i = 0; i < length % 4; ++i) {
+            while (sb.length() % 4 != 0) {
                 sb.append('=');
             }
             base64 = sb.toString();
diff --git a/web/src/test/java/org/apache/shiro/web/mgt/CookieRememberMeManagerTest.java b/web/src/test/java/org/apache/shiro/web/mgt/CookieRememberMeManagerTest.java
index 7ccc47be0e..dad1fd559d 100644
--- a/web/src/test/java/org/apache/shiro/web/mgt/CookieRememberMeManagerTest.java
+++ b/web/src/test/java/org/apache/shiro/web/mgt/CookieRememberMeManagerTest.java
@@ -45,7 +45,7 @@
  *
  * @since 1.0
  */
-public class CookieRememberMeManagerTest {
+class CookieRememberMeManagerTest {
 
     @Test
     void onSuccessfulLogin() {
@@ -142,7 +142,7 @@ void getRememberedPrincipals() {
         replay(mockRequest);
 
         CookieRememberMeManager mgr = new CookieRememberMeManager();
-        mgr.setCipherKey( Base64.decode("kPH+bIxk5D2deZiIxcaaaA=="));
+        mgr.setCipherKey(Base64.decode("kPH+bIxk5D2deZiIxcaaaA=="));
         PrincipalCollection collection = mgr.getRememberedPrincipals(context);
 
         verify(mockRequest);
@@ -193,7 +193,7 @@ void getRememberedPrincipalsNoMoreDefaultCipher() {
     // SHIRO-69
 
     @Test
-    void getRememberedPrincipalsDecryptionError() {
+    void getRememberedPrincipalsDecryptionError_whenWrongCookieValue() {
         HttpServletRequest mockRequest = createNiceMock(HttpServletRequest.class);
         HttpServletResponse mockResponse = createNiceMock(HttpServletResponse.class);
 
@@ -203,8 +203,8 @@ void getRememberedPrincipalsDecryptionError() {
 
         expect(mockRequest.getAttribute(ShiroHttpServletRequest.IDENTITY_REMOVED_KEY)).andReturn(null);
 
-        // Simulate a bad return value here (for example if this was encrypted with a different key
-        final String userPCAesBase64 = "garbage";
+        // Simulate a bad return value here (valid Base64 does not contain key)
+        String userPCAesBase64 = "garbage";
         Cookie[] cookies = new Cookie[]{
                 new Cookie(CookieRememberMeManager.DEFAULT_REMEMBER_ME_COOKIE_NAME, userPCAesBase64)
         };
@@ -213,7 +213,31 @@ void getRememberedPrincipalsDecryptionError() {
         replay(mockRequest);
 
         CookieRememberMeManager mgr = new CookieRememberMeManager();
-        final PrincipalCollection rememberedPrincipals = mgr.getRememberedPrincipals(context);
+        assertThrows(CryptoException.class, () -> mgr.getRememberedPrincipals(context), "CryptoException should be thrown on invalid cookies");
+    }
+
+    @Test
+    void getRememberedPrincipalsDecryptionError_whenInvalidBase64() {
+        HttpServletRequest mockRequest = createNiceMock(HttpServletRequest.class);
+        HttpServletResponse mockResponse = createNiceMock(HttpServletResponse.class);
+
+        WebSubjectContext context = new DefaultWebSubjectContext();
+        context.setServletRequest(mockRequest);
+        context.setServletResponse(mockResponse);
+
+        expect(mockRequest.getAttribute(ShiroHttpServletRequest.IDENTITY_REMOVED_KEY)).andReturn(null);
+
+        // Simulate a bad return value here (not valid Base64)
+        String userPCAesBase64 = "InvalidBase64";
+        Cookie[] cookies = new Cookie[]{
+                new Cookie(CookieRememberMeManager.DEFAULT_REMEMBER_ME_COOKIE_NAME, userPCAesBase64)
+        };
+
+        expect(mockRequest.getCookies()).andReturn(cookies).anyTimes();
+        replay(mockRequest);
+
+        CookieRememberMeManager mgr = new CookieRememberMeManager();
+        PrincipalCollection rememberedPrincipals = mgr.getRememberedPrincipals(context);
         assertNull(rememberedPrincipals, "rememberedPrincipals should be null on invalid cookies.");
     }
 
@@ -270,4 +294,19 @@ void shouldIgnoreInvalidCookieValues() {
         // then
         assertNull(rememberedSerializedIdentity, "should ignore invalid cookie values");
     }
+
+    @Test
+    void ensurePaddingShouldAddEnoughEquals() {
+        CookieRememberMeManager mgr = new CookieRememberMeManager();
+        StringBuilder stringToTest = new StringBuilder("A string to test padding");
+        for (int i = 0; i < 10; i++) {
+            stringToTest.append("x");
+            String encoded = Base64.encodeToString(stringToTest.toString().getBytes());
+            while (encoded.endsWith("=")) {
+                encoded = encoded.substring(0, encoded.length() - 1);
+            }
+            String base64 = mgr.ensurePadding(encoded);
+            assertDoesNotThrow(() -> Base64.decode(base64), "Error decoding " + stringToTest);
+        }
+    }
 }