Fix updating digest after its retrieval

Author: otethal

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/hashlib/DigestObject.java

@@ -63,7 +63,6 @@
  */
 public abstract class DigestObject extends PythonBuiltinObject {
     private final String name;
-    private byte[] finalDigest = null;
 
     DigestObject(Object cls, Shape instanceShape, String name) {
         super(cls, instanceShape);
@@ -178,26 +177,82 @@ private PythonBuiltinClassType determineMainDigestType() {
         }
     }
 
-    byte[] digest() {
-        if (finalDigest == null) {
-            finalDigest = finalizeDigest();
-        }
-        return finalDigest;
-    }
+    /**
+     * CPython supports updating after retrieving the digest but the JDK does not. We have to
+     * calculate the digest on a clone, but that does not need to be supported. If it is not, then
+     * we calculate the digest normally, but we must prevent any further updates.
+     *
+     * @see #wasReset(), {@link #update(byte[])}
+     */
+    abstract byte[] digest();
 
-    abstract DigestObject copy(PythonObjectFactory factory) throws CloneNotSupportedException;
-
-    abstract byte[] finalizeDigest();
+    /**
+     * @return true if the digest has already been calculated and the underlying implementation does
+     *         not support cloning, in which case this object can no longer be
+     *         {@linkplain #update(byte[]) updated}
+     */
+    abstract boolean wasReset();
 
+    /**
+     * Must not be called if {@link #wasReset()} returns true.
+     */
     abstract void update(byte[] data);
 
+    abstract DigestObject copy(PythonObjectFactory factory) throws CloneNotSupportedException;
+
     abstract int getDigestLength();
 
     final String getAlgorithm() {
         return name;
     }
 
-    private static final class MessageDigestObject extends DigestObject {
+    /**
+     * Ensures that {@link #update(byte[])} is not called after {@link #digest()} if cloning is not
+     * supported. Also caches the digest and ensures that the cache is cleared on update.
+     */
+    private static abstract class DigestObjectBase extends DigestObject {
+        private byte[] cachedDigest = null;
+        private boolean wasReset;
+
+        DigestObjectBase(Object cls, Shape instanceShape, String name) {
+            super(cls, instanceShape, name);
+        }
+
+        @Override
+        final boolean wasReset() {
+            return wasReset;
+        }
+
+        @Override
+        final byte[] digest() {
+            if (cachedDigest == null) {
+                try {
+                    cachedDigest = calculateDigestOnClone();
+                } catch (CloneNotSupportedException e) {
+                    wasReset = true;
+                    cachedDigest = calculateDigest();
+                }
+            }
+            return cachedDigest;
+        }
+
+        @Override
+        final void update(byte[] data) {
+            if (wasReset) {
+                throw CompilerDirectives.shouldNotReachHere("update() called after digest() on an implementation the does not support clone()");
+            }
+            cachedDigest = null;
+            doUpdate(data);
+        }
+
+        abstract byte[] calculateDigestOnClone() throws CloneNotSupportedException;
+
+        abstract byte[] calculateDigest();
+
+        abstract void doUpdate(byte[] data);
+    }
+
+    private static final class MessageDigestObject extends DigestObjectBase {
         private final MessageDigest digest;
 
         MessageDigestObject(PythonBuiltinClassType digestType, Shape instanceShape, String name, MessageDigest digest) {
@@ -213,19 +268,19 @@ DigestObject copy(PythonObjectFactory factory) throws CloneNotSupportedException
 
         @Override
         @TruffleBoundary
-        byte[] finalizeDigest() {
-            try {
-                // Try to finalize a clone, because CPython supports updating
-                // after retrieving the digest and the JDK does not.
-                return ((MessageDigest) digest.clone()).digest();
-            } catch (CloneNotSupportedException e) {
-                return digest.digest();
-            }
+        byte[] calculateDigestOnClone() throws CloneNotSupportedException {
+            return ((MessageDigest) digest.clone()).digest();
         }
 
         @Override
         @TruffleBoundary
-        void update(byte[] data) {
+        byte[] calculateDigest() {
+            return digest.digest();
+        }
+
+        @Override
+        @TruffleBoundary
+        void doUpdate(byte[] data) {
             digest.update(data);
         }
 
@@ -236,7 +291,7 @@ int getDigestLength() {
         }
     }
 
-    private static final class MacDigestObject extends DigestObject {
+    private static final class MacDigestObject extends DigestObjectBase {
         private final Mac mac;
 
         MacDigestObject(PythonBuiltinClassType digestType, Shape instanceShape, String name, Mac mac) {
@@ -252,13 +307,19 @@ DigestObject copy(PythonObjectFactory factory) throws CloneNotSupportedException
 
         @Override
         @TruffleBoundary
-        byte[] finalizeDigest() {
+        byte[] calculateDigestOnClone() throws CloneNotSupportedException {
+            return ((Mac) mac.clone()).doFinal();
+        }
+
+        @Override
+        @TruffleBoundary
+        byte[] calculateDigest() {
             return mac.doFinal();
         }
 
         @Override
         @TruffleBoundary
-        void update(byte[] data) {
+        void doUpdate(byte[] data) {
             mac.update(data);
         }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/hashlib/DigestObjectBuiltins.java

@@ -52,6 +52,7 @@
 import com.oracle.graal.python.builtins.objects.buffer.PythonBufferAccessLibrary;
 import com.oracle.graal.python.builtins.objects.bytes.BytesNodes;
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
+import com.oracle.graal.python.nodes.ErrorMessages;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryClinicBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
@@ -122,6 +123,9 @@ protected ArgumentClinicProvider getArgumentClinic() {
         @Specialization(limit = "3")
         PNone update(VirtualFrame frame, DigestObject self, Object buffer,
                         @CachedLibrary("buffer") PythonBufferAccessLibrary bufferLib) {
+            if (self.wasReset()) {
+                raise(PythonBuiltinClassType.ValueError, ErrorMessages.UPDATING_FINALIZED_DIGEST_IS_NOT_SUPPORTED);
+            }
             try {
                 self.update(bufferLib.getInternalOrCopiedByteArray(buffer));
             } finally {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/ErrorMessages.java

@@ -784,6 +784,7 @@ public abstract class ErrorMessages {
     public static final TruffleString UNSUPPORTED_STR_TYPE = tsLiteral("unsupported string type: %s");
     public static final TruffleString UNSUPPORTED_TARGET_SIZE = tsLiteral("Unsupported target size: %d");
     public static final TruffleString UNSUPPORTED_USE_OF_SYS_EXECUTABLE = tsLiteral("internal error: unsupported use of sys.executable");
+    public static final TruffleString UPDATING_FINALIZED_DIGEST_IS_NOT_SUPPORTED = tsLiteral("internal error: updating a finalized digest is not supported");
     public static final TruffleString UTIME_CANNOT_USE_DIR_FD_AND_FOLLOW_SYMLINKS = tsLiteral("utime: cannot use dir_fd and follow_symlinks together on this platform");
     public static final TruffleString VALUE_TOO_LARGE_TO_FIT_INTO_INDEX = tsLiteral("value too large to fit into index-sized integer");
     public static final TruffleString VARS_ARGUMENT_MUST_HAVE_DICT = tsLiteral("vars() argument must have __dict__ attribute");