package org.graalvm.compiler.hotspot.replacements;
import static org.graalvm.compiler.hotspot.GraalHotSpotVMConfig.INJECTED_INTRINSIC_CONTEXT;
import static org.graalvm.compiler.hotspot.GraalHotSpotVMConfig.INJECTED_METAACCESS;
import static org.graalvm.compiler.hotspot.HotSpotBackend.DECRYPT_BLOCK;
import static org.graalvm.compiler.hotspot.HotSpotBackend.DECRYPT_BLOCK_WITH_ORIGINAL_KEY;
import static org.graalvm.compiler.hotspot.HotSpotBackend.ENCRYPT_BLOCK;
import static org.graalvm.compiler.hotspot.replacements.CipherBlockChainingSubstitutions.aesCryptType;
import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.VERY_SLOW_PATH_PROBABILITY;
import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.probability;
import static org.graalvm.compiler.replacements.ReplacementsUtil.getArrayBaseOffset;
import org.graalvm.compiler.api.replacements.ClassSubstitution;
import org.graalvm.compiler.api.replacements.Fold;
import org.graalvm.compiler.api.replacements.MethodSubstitution;
import org.graalvm.compiler.core.common.spi.ForeignCallDescriptor;
import org.graalvm.compiler.graph.Node.ConstantNodeParameter;
import org.graalvm.compiler.graph.Node.NodeIntrinsic;
import org.graalvm.compiler.nodes.ComputeObjectAddressNode;
import org.graalvm.compiler.nodes.DeoptimizeNode;
import org.graalvm.compiler.nodes.PiNode;
import org.graalvm.compiler.nodes.extended.ForeignCallNode;
import org.graalvm.compiler.nodes.extended.RawLoadNode;
import org.graalvm.compiler.nodes.graphbuilderconf.IntrinsicContext;
import org.graalvm.compiler.word.Word;
import jdk.internal.vm.compiler.word.LocationIdentity;
import jdk.internal.vm.compiler.word.Pointer;
import jdk.internal.vm.compiler.word.WordFactory;
import jdk.vm.ci.meta.DeoptimizationAction;
import jdk.vm.ci.meta.DeoptimizationReason;
import jdk.vm.ci.meta.JavaKind;
@ClassSubstitution(className = "com.sun.crypto.provider.AESCrypt", optional = true)
public class AESCryptSubstitutions {
static final int AES_BLOCK_SIZE_IN_BYTES = 16;
@Fold
static long kOffset(@Fold.InjectedParameter IntrinsicContext context) {
return HotSpotReplacementsUtil.getFieldOffset(aesCryptType(context), "K");
}
@Fold
static long lastKeyOffset(@Fold.InjectedParameter IntrinsicContext context) {
return HotSpotReplacementsUtil.getFieldOffset(aesCryptType(context), "lastKey");
}
@MethodSubstitution(isStatic = false)
static void encryptBlock(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, true, false);
}
@MethodSubstitution(isStatic = false)
static void implEncryptBlock(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, true, false);
}
@MethodSubstitution(isStatic = false)
static void decryptBlock(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, false, false);
}
@MethodSubstitution(isStatic = false)
static void implDecryptBlock(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, false, false);
}
@MethodSubstitution(value = "decryptBlock", isStatic = false)
static void decryptBlockWithOriginalKey(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, false, true);
}
@MethodSubstitution(value = "implDecryptBlock", isStatic = false)
static void implDecryptBlockWithOriginalKey(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset) {
crypt(rcvr, in, inOffset, out, outOffset, false, true);
}
private static void crypt(Object rcvr, byte[] in, int inOffset, byte[] out, int outOffset, boolean encrypt, boolean withOriginalKey) {
checkArgs(in, inOffset, out, outOffset);
Object realReceiver = PiNode.piCastNonNull(rcvr, aesCryptType(INJECTED_INTRINSIC_CONTEXT));
Object kObject = RawLoadNode.load(realReceiver, kOffset(INJECTED_INTRINSIC_CONTEXT), JavaKind.Object, LocationIdentity.any());
Pointer kAddr = Word.objectToTrackedPointer(kObject).add(getArrayBaseOffset(INJECTED_METAACCESS, JavaKind.Int));
Word inAddr = WordFactory.unsigned(ComputeObjectAddressNode.get(in, getArrayBaseOffset(INJECTED_METAACCESS, JavaKind.Byte) + inOffset));
Word outAddr = WordFactory.unsigned(ComputeObjectAddressNode.get(out, getArrayBaseOffset(INJECTED_METAACCESS, JavaKind.Byte) + outOffset));
if (encrypt) {
encryptBlockStub(ENCRYPT_BLOCK, inAddr, outAddr, kAddr);
} else {
if (withOriginalKey) {
Object lastKeyObject = RawLoadNode.load(realReceiver, lastKeyOffset(INJECTED_INTRINSIC_CONTEXT), JavaKind.Object, LocationIdentity.any());
Pointer lastKeyAddr = Word.objectToTrackedPointer(lastKeyObject).add(getArrayBaseOffset(INJECTED_METAACCESS, JavaKind.Byte));
decryptBlockWithOriginalKeyStub(DECRYPT_BLOCK_WITH_ORIGINAL_KEY, inAddr, outAddr, kAddr, lastKeyAddr);
} else {
decryptBlockStub(DECRYPT_BLOCK, inAddr, outAddr, kAddr);
}
}
}
static void checkArgs(byte[] in, int inOffset, byte[] out, int outOffset) {
if (probability(VERY_SLOW_PATH_PROBABILITY, inOffset < 0 || in.length - AES_BLOCK_SIZE_IN_BYTES < inOffset || outOffset < 0 || out.length - AES_BLOCK_SIZE_IN_BYTES < outOffset)) {
DeoptimizeNode.deopt(DeoptimizationAction.None, DeoptimizationReason.RuntimeConstraint);
}
}
@NodeIntrinsic(ForeignCallNode.class)
public static native void encryptBlockStub(@ConstantNodeParameter ForeignCallDescriptor descriptor, Word in, Word out, Pointer key);
@NodeIntrinsic(ForeignCallNode.class)
public static native void decryptBlockStub(@ConstantNodeParameter ForeignCallDescriptor descriptor, Word in, Word out, Pointer key);
@NodeIntrinsic(ForeignCallNode.class)
public static native void decryptBlockWithOriginalKeyStub(@ConstantNodeParameter ForeignCallDescriptor descriptor, Word in, Word out, Pointer key, Pointer originalKey);
}