|
| 1 | +/* |
| 2 | + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one |
| 3 | + * or more contributor license agreements. Licensed under the "Elastic License |
| 4 | + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side |
| 5 | + * Public License v 1"; you may not use this file except in compliance with, at |
| 6 | + * your election, the "Elastic License 2.0", the "GNU Affero General Public |
| 7 | + * License v3.0 only", or the "Server Side Public License, v 1". |
| 8 | + */ |
| 9 | + |
| 10 | +package org.elasticsearch.entitlement.instrumentation.impl; |
| 11 | + |
| 12 | +import org.elasticsearch.entitlement.api.EntitlementChecks; |
| 13 | +import org.elasticsearch.entitlement.api.EntitlementProvider; |
| 14 | +import org.elasticsearch.entitlement.instrumentation.InstrumentationService; |
| 15 | +import org.elasticsearch.entitlement.instrumentation.MethodKey; |
| 16 | +import org.elasticsearch.logging.LogManager; |
| 17 | +import org.elasticsearch.logging.Logger; |
| 18 | +import org.elasticsearch.test.ESTestCase; |
| 19 | +import org.junit.Before; |
| 20 | + |
| 21 | +import java.lang.reflect.InvocationTargetException; |
| 22 | +import java.lang.reflect.Method; |
| 23 | +import java.util.Map; |
| 24 | + |
| 25 | +import static org.elasticsearch.entitlement.instrumentation.impl.ASMUtils.bytecode2text; |
| 26 | + |
| 27 | +/** |
| 28 | + * This tests {@link InstrumenterImpl} in isolation, without a java agent. |
| 29 | + * It causes the methods to be instrumented, and verifies that the instrumentation is called as expected. |
| 30 | + * Problems with bytecode generation are easier to debug this way than in the context of an agent. |
| 31 | + */ |
| 32 | +@ESTestCase.WithoutSecurityManager |
| 33 | +public class InstrumenterTests extends ESTestCase { |
| 34 | + final InstrumentationService instrumentationService = new InstrumentationServiceImpl(); |
| 35 | + |
| 36 | + private static TestEntitlementManager getTestChecks() { |
| 37 | + return (TestEntitlementManager) EntitlementProvider.checks(); |
| 38 | + } |
| 39 | + |
| 40 | + @Before |
| 41 | + public void initialize() { |
| 42 | + getTestChecks().isActive = false; |
| 43 | + } |
| 44 | + |
| 45 | + /** |
| 46 | + * Contains all the virtual methods from {@link ClassToInstrument}, |
| 47 | + * allowing this test to call them on the dynamically loaded instrumented class. |
| 48 | + */ |
| 49 | + public interface Testable {} |
| 50 | + |
| 51 | + /** |
| 52 | + * This is a placeholder for real class library methods. |
| 53 | + * Without the java agent, we can't instrument the real methods, so we instrument this instead. |
| 54 | + * <p> |
| 55 | + * Methods of this class must have the same signature and the same static/virtual condition as the corresponding real method. |
| 56 | + * They should assert that the arguments came through correctly. |
| 57 | + * They must not throw {@link TestException}. |
| 58 | + */ |
| 59 | + public static class ClassToInstrument implements Testable { |
| 60 | + public static void systemExit(int status) { |
| 61 | + assertEquals(123, status); |
| 62 | + } |
| 63 | + } |
| 64 | + |
| 65 | + static final class TestException extends RuntimeException {} |
| 66 | + |
| 67 | + /** |
| 68 | + * We're not testing the permission checking logic here. |
| 69 | + * This is a trivial implementation of {@link EntitlementChecks} that just always throws, |
| 70 | + * just to demonstrate that the injected bytecodes succeed in calling these methods. |
| 71 | + */ |
| 72 | + public static class TestEntitlementManager implements EntitlementChecks { |
| 73 | + /** |
| 74 | + * This allows us to test that the instrumentation is correct in both cases: |
| 75 | + * if the check throws, and if it doesn't. |
| 76 | + */ |
| 77 | + volatile boolean isActive; |
| 78 | + |
| 79 | + @Override |
| 80 | + public void checkSystemExit(Class<?> callerClass, int status) { |
| 81 | + assertSame(InstrumenterTests.class, callerClass); |
| 82 | + assertEquals(123, status); |
| 83 | + throwIfActive(); |
| 84 | + } |
| 85 | + |
| 86 | + private void throwIfActive() { |
| 87 | + if (isActive) { |
| 88 | + throw new TestException(); |
| 89 | + } |
| 90 | + } |
| 91 | + } |
| 92 | + |
| 93 | + public void test() throws Exception { |
| 94 | + // This test doesn't replace ClassToInstrument in-place but instead loads a separate |
| 95 | + // class ClassToInstrument_NEW that contains the instrumentation. Because of this, |
| 96 | + // we need to configure the Transformer to use a MethodKey and instrumentationMethod |
| 97 | + // with slightly different signatures (using the common interface Testable) which |
| 98 | + // is not what would happen when it's run by the agent. |
| 99 | + |
| 100 | + MethodKey k1 = instrumentationService.methodKeyForTarget(ClassToInstrument.class.getMethod("systemExit", int.class)); |
| 101 | + Method v1 = EntitlementChecks.class.getMethod("checkSystemExit", Class.class, int.class); |
| 102 | + var instrumenter = new InstrumenterImpl("_NEW", Map.of(k1, v1)); |
| 103 | + |
| 104 | + byte[] newBytecode = instrumenter.instrumentClassFile(ClassToInstrument.class).bytecodes(); |
| 105 | + |
| 106 | + if (logger.isTraceEnabled()) { |
| 107 | + logger.trace("Bytecode after instrumentation:\n{}", bytecode2text(newBytecode)); |
| 108 | + } |
| 109 | + |
| 110 | + Class<?> newClass = new TestLoader(Testable.class.getClassLoader()).defineClassFromBytes( |
| 111 | + ClassToInstrument.class.getName() + "_NEW", |
| 112 | + newBytecode |
| 113 | + ); |
| 114 | + |
| 115 | + // Before checking is active, nothing should throw |
| 116 | + callStaticSystemExit(newClass, 123); |
| 117 | + |
| 118 | + getTestChecks().isActive = true; |
| 119 | + |
| 120 | + // After checking is activated, everything should throw |
| 121 | + assertThrows(TestException.class, () -> callStaticSystemExit(newClass, 123)); |
| 122 | + } |
| 123 | + |
| 124 | + /** |
| 125 | + * Calling a static method of a dynamically loaded class is significantly more cumbersome |
| 126 | + * than calling a virtual method. |
| 127 | + */ |
| 128 | + private static void callStaticSystemExit(Class<?> c, int status) throws NoSuchMethodException, IllegalAccessException { |
| 129 | + try { |
| 130 | + c.getMethod("systemExit", int.class).invoke(null, status); |
| 131 | + } catch (InvocationTargetException e) { |
| 132 | + Throwable cause = e.getCause(); |
| 133 | + if (cause instanceof TestException n) { |
| 134 | + // Sometimes we're expecting this one! |
| 135 | + throw n; |
| 136 | + } else { |
| 137 | + throw new AssertionError(cause); |
| 138 | + } |
| 139 | + } |
| 140 | + } |
| 141 | + |
| 142 | + static class TestLoader extends ClassLoader { |
| 143 | + TestLoader(ClassLoader parent) { |
| 144 | + super(parent); |
| 145 | + } |
| 146 | + |
| 147 | + public Class<?> defineClassFromBytes(String name, byte[] bytes) { |
| 148 | + return defineClass(name, bytes, 0, bytes.length); |
| 149 | + } |
| 150 | + } |
| 151 | + |
| 152 | + private static final Logger logger = LogManager.getLogger(InstrumenterTests.class); |
| 153 | +} |
0 commit comments