Merge pull request #34 from kkoz/pickle-error-logging

Log session when session unpickling fails
glencoesoftware · Feb 13, 2025 · abdd5d3 · abdd5d3
2 parents 13bb969 + 29755a4
commit abdd5d3
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/src/main/java/com/glencoesoftware/omero/ms/core/PickledSessionConnector.java b/src/main/java/com/glencoesoftware/omero/ms/core/PickledSessionConnector.java
@@ -21,6 +21,7 @@
 import java.math.BigInteger;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -69,14 +70,23 @@ public class PickledSessionConnector implements IConnector {
 
     private Long userId;
 
+    private String b64Session;
+
     protected PickledSessionConnector() {
     }
 
     public PickledSessionConnector(byte[] serialized) {
-        init(serialized);
+        try {
+            init(serialized);
+        } catch (Exception e) {
+            log.error("Pickled Session: {}",
+                    Base64.getEncoder().encodeToString(serialized));
+            throw e;
+        }
     }
 
     protected void init(byte[] sessionData) {
+        b64Session = Base64.getEncoder().encodeToString(sessionData);
         ByteBufferKaitaiStream bbks = new ByteBufferKaitaiStream(sessionData);
         PythonPickle pickleData = new PythonPickle(bbks);
         List<Op> ops = pickleData.ops();
@@ -154,6 +164,9 @@ private void deserializeConnector(Iterator<Op> opIterator) {
                         case "is_public":
                             isPublic = deserializeBooleanField(opIterator);
                             break;
+                        default:
+                            log.warn("Unexpected field name in connector: {}",
+                                     fieldName);
                     }
                 } catch (Exception e) {
                     log.error("Exception while deserializing: {}", fieldName);
@@ -259,6 +272,11 @@ public static String handleStringValue(
         } else if (value.code() == PythonPickle.Opcode.BINGET
                 || value.code() == PythonPickle.Opcode.LONG_BINGET) {
             v = memo.get(value.arg());
+            if (v == null) {
+                throw new IllegalArgumentException(
+                        "Failed to find key in memo: " +
+                        value.arg().toString());
+            }
         } else if (throwOnUnexpected){
             throw new IllegalArgumentException(
                     "Unexpected opcode for string field: " + value.code());

diff --git a/src/test/java/com/glencoesoftware/omero/ms/core/PickledSessionConnectorTest.java b/src/test/java/com/glencoesoftware/omero/ms/core/PickledSessionConnectorTest.java
@@ -421,6 +421,16 @@ public class PickledSessionConnectorTest {
     private static String BINUNICODE8_TEST = "gASVEgAAAAAAAAB9lI0EAAAAAAAAAHRlc3SUjARkaWN0lHMu";
     //python3>>> base64.b64encode(b'\x80\x04\x95\x12\x00\x00\x00\x00\x00\x00\x00}\x94\x8d\x04\x00\x00\x00\x00\x00\x00\x00test\x94\x8c\x04dict\x94s.')
 
+    private static String STRING_USER_ID = "gASVeAAAAAAAAAB9lCiMBmJhbmFuYZRLFo"
+            + "wJY29ubmVjdG9ylH2UKIwJc2VydmVyX2lklIwCLTGUjAlpc19zZWN1cmWUiIwHd"
+            + "XNlcl9pZJSMAzEyM5SMEW9tZXJvX3Nlc3Npb25fa2V5lIwGYWJjMTIzlIwJaXNf"
+            + "cHVibGljlIl1dS4=";
+
+    private static String MEMOIZED_SESSION_ID = "gASVfwAAAAAAAAB9lCiMBmJhbmFuY"
+            + "ZRLFowJY29ubmVjdG9ylH2UKIwGa2V5ZHVwlIwGYWJjMTIzlIwJc2VydmVyX2lk"
+            + "lIwCLTGUjAlpc19zZWN1cmWUiIwHdXNlcl9pZJRLe4wRb21lcm9fc2Vzc2lvbl9"
+            + "rZXmUaAWMCWlzX3B1YmxpY5SJdXUu";
+
     @Test
     public void testUnpicklingJDBC() {
         IConnector v = new JDBCPickledSessionConnector(DB_SESSION_DATA);
@@ -502,6 +512,18 @@ public void testUnpicklingRedisWithBinget() {
         Assert.assertEquals(v.getUserId(), Long.valueOf(2L));
     }
 
+    @Test(expectedExceptions={IllegalArgumentException.class})
+    public void testUnpicklingRedisWithStringUserId() {
+        IConnector v = new PickledSessionConnector(
+                Base64.getDecoder().decode(STRING_USER_ID));
+    }
+
+    @Test(expectedExceptions={IllegalArgumentException.class})
+    public void testUnpicklingRedisWithSessionIdMemoizedInConnector() {
+        IConnector v = new PickledSessionConnector(
+                Base64.getDecoder().decode(MEMOIZED_SESSION_ID));
+    }
+
     @Test
     public void nonzeroLongTest() {
         byte[] data = Base64.getDecoder().decode(LONG_NONZERO);