refactor(YouTube - Litho): Use a simpler hook that does not require using a thread local (#5281)

LisoUseInAIKyrios · web-flow · commit 8cc8a91bec3c · 2025-06-28T11:51:29.000+04:00
diff --git a/extensions/youtube/src/main/java/app/revanced/extension/youtube/patches/components/LithoFilterPatch.java b/extensions/youtube/src/main/java/app/revanced/extension/youtube/patches/components/LithoFilterPatch.java
@@ -48,7 +48,7 @@ public String toString() {
         /**
          * Search through a byte array for all ASCII strings.
          */
-        private static void findAsciiStrings(StringBuilder builder, byte[] buffer) {
+        static void findAsciiStrings(StringBuilder builder, byte[] buffer) {
             // Valid ASCII values (ignore control characters).
             final int minimumAscii = 32;  // 32 = space character
             final int maximumAscii = 126; // 127 = delete character
@@ -96,7 +96,7 @@ private static void findAsciiStrings(StringBuilder builder, byte[] buffer) {
     private static final class DummyFilter extends Filter { }
 
     private static final Filter[] filters = new Filter[] {
-            new DummyFilter() // Replaced by patch.
+            new DummyFilter() // Replaced patching, do not touch.
     };
 
     private static final StringTrieSearch pathSearchTree = new StringTrieSearch();
@@ -108,11 +108,7 @@ private static final class DummyFilter extends Filter { }
      * Because litho filtering is multi-threaded and the buffer is passed in from a different injection point,
      * the buffer is saved to a ThreadLocal so each calling thread does not interfere with other threads.
      */
-    private static final ThreadLocal<ByteBuffer> bufferThreadLocal = new ThreadLocal<>();
-    /**
-     * Results of calling {@link #filter(String, StringBuilder)}.
-     */
-    private static final ThreadLocal<Boolean> filterResult = new ThreadLocal<>();
+    private static final ThreadLocal<byte[]> bufferThreadLocal = new ThreadLocal<>();
 
     static {
         for (Filter filter : filters) {
@@ -168,57 +164,50 @@ private static void filterUsingCallbacks(StringTrieSearch pathSearchTree,
     /**
      * Injection point.  Called off the main thread.
      */
-    @SuppressWarnings("unused")
-    public static void setProtoBuffer(@Nullable ByteBuffer protobufBuffer) {
+    public static void setProtoBuffer(byte[] buffer) {
         // Set the buffer to a thread local.  The buffer will remain in memory, even after the call to #filter completes.
         // This is intentional, as it appears the buffer can be set once and then filtered multiple times.
         // The buffer will be cleared from memory after a new buffer is set by the same thread,
         // or when the calling thread eventually dies.
-        if (protobufBuffer == null) {
+        bufferThreadLocal.set(buffer);
+    }
+
+    /**
+     * Injection point.  Called off the main thread.
+     * Targets 20.21 and lower.
+     */
+    public static void setProtoBuffer(@Nullable ByteBuffer buffer) {
+        // Set the buffer to a thread local.  The buffer will remain in memory, even after the call to #filter completes.
+        // This is intentional, as it appears the buffer can be set once and then filtered multiple times.
+        // The buffer will be cleared from memory after a new buffer is set by the same thread,
+        // or when the calling thread eventually dies.
+        if (buffer == null || !buffer.hasArray()) {
             // It appears the buffer can be cleared out just before the call to #filter()
             // Ignore this null value and retain the last buffer that was set.
-            Logger.printDebug(() -> "Ignoring null protobuffer");
+            Logger.printDebug(() -> "Ignoring null or empty buffer: " + buffer);
         } else {
-            bufferThreadLocal.set(protobufBuffer);
+            setProtoBuffer(buffer.array());
         }
     }
 
     /**
      * Injection point.
      */
-    public static boolean shouldFilter() {
-        Boolean shouldFilter = filterResult.get();
-        return shouldFilter != null && shouldFilter;
-    }
-
-    /**
-     * Injection point.  Called off the main thread, and commonly called by multiple threads at the same time.
-     */
-    public static void filter(@Nullable String lithoIdentifier, StringBuilder pathBuilder) {
-        filterResult.set(handleFiltering(lithoIdentifier, pathBuilder));
-    }
-
-    private static boolean handleFiltering(@Nullable String lithoIdentifier, StringBuilder pathBuilder) {
+    public static boolean shouldFilter(@Nullable String lithoIdentifier, StringBuilder pathBuilder) {
         try {
             if (pathBuilder.length() == 0) {
                 return false;
             }
 
-            ByteBuffer protobufBuffer = bufferThreadLocal.get();
-            final byte[] bufferArray;
+            byte[] buffer = bufferThreadLocal.get();
             // Potentially the buffer may have been null or never set up until now.
             // Use an empty buffer so the litho id/path filters still work correctly.
-            if (protobufBuffer == null) {
-                bufferArray = EMPTY_BYTE_ARRAY;
-            } else if (!protobufBuffer.hasArray()) {
-                Logger.printDebug(() -> "Proto buffer does not have an array, using an empty buffer array");
-                bufferArray = EMPTY_BYTE_ARRAY;
-            } else {
-                bufferArray = protobufBuffer.array();
+            if (buffer == null) {
+                buffer = EMPTY_BYTE_ARRAY;
             }
 
-            LithoFilterParameters parameter = new LithoFilterParameters(lithoIdentifier,
-                    pathBuilder.toString(), bufferArray);
+            LithoFilterParameters parameter = new LithoFilterParameters(
+                    lithoIdentifier, pathBuilder.toString(), buffer);
             Logger.printDebug(() -> "Searching " + parameter);
 
             if (parameter.identifier != null && identifierSearchTree.matches(parameter.identifier, parameter)) {
diff --git a/patches/src/main/kotlin/app/revanced/patches/youtube/misc/litho/filter/Fingerprints.kt b/patches/src/main/kotlin/app/revanced/patches/youtube/misc/litho/filter/Fingerprints.kt
@@ -7,26 +7,18 @@ import com.android.tools.smali.dexlib2.AccessFlags
 import com.android.tools.smali.dexlib2.Opcode
 
 internal val componentContextParserFingerprint = fingerprint {
-    strings(
-        "TreeNode result must be set.",
-        // String is a partial match and changed slightly in 20.03+
-        "it was removed due to duplicate converter bindings."
-    )
+    strings("Number of bits must be positive")
 }
 
-/**
- * Resolves to the class found in [componentContextParserFingerprint].
- * When patching 19.16 this fingerprint matches the same method as [componentContextParserFingerprint].
- */
-internal val componentContextSubParserFingerprint = fingerprint {
+internal val componentCreateFingerprint = fingerprint {
     strings(
-        "Number of bits must be positive"
+        "Element missing correct type extension",
+        "Element missing type"
     )
 }
 
 internal val lithoFilterFingerprint = fingerprint {
     accessFlags(AccessFlags.STATIC, AccessFlags.CONSTRUCTOR)
-    returns("V")
     custom { _, classDef ->
         classDef.endsWith("/LithoFilterPatch;")
     }
@@ -58,7 +50,7 @@ internal val lithoThreadExecutorFingerprint = fingerprint {
     parameters("I", "I", "I")
     custom { method, classDef ->
         classDef.superclass == "Ljava/util/concurrent/ThreadPoolExecutor;" &&
-            method.containsLiteralInstruction(1L) // 1L = default thread timeout.
+                method.containsLiteralInstruction(1L) // 1L = default thread timeout.
     }
 }
 
diff --git a/patches/src/main/kotlin/app/revanced/patches/youtube/misc/litho/filter/LithoFilterPatch.kt b/patches/src/main/kotlin/app/revanced/patches/youtube/misc/litho/filter/LithoFilterPatch.kt
@@ -9,13 +9,13 @@ import app.revanced.patcher.extensions.InstructionExtensions.removeInstructions
 import app.revanced.patcher.extensions.InstructionExtensions.replaceInstruction
 import app.revanced.patcher.patch.bytecodePatch
 import app.revanced.patches.youtube.misc.extension.sharedExtensionPatch
+import app.revanced.patches.youtube.misc.playservice.is_19_17_or_greater
 import app.revanced.patches.youtube.misc.playservice.is_19_25_or_greater
 import app.revanced.patches.youtube.misc.playservice.is_20_05_or_greater
 import app.revanced.patches.youtube.misc.playservice.versionCheckPatch
 import app.revanced.patches.youtube.shared.conversionContextFingerprintToString
 import app.revanced.util.addInstructionsAtControlFlowLabel
 import app.revanced.util.findFreeRegister
-import app.revanced.util.findInstructionIndicesReversedOrThrow
 import app.revanced.util.getReference
 import app.revanced.util.indexOfFirstInstructionOrThrow
 import app.revanced.util.indexOfFirstInstructionReversedOrThrow
@@ -24,7 +24,6 @@ import com.android.tools.smali.dexlib2.Opcode
 import com.android.tools.smali.dexlib2.iface.instruction.OneRegisterInstruction
 import com.android.tools.smali.dexlib2.iface.instruction.ReferenceInstruction
 import com.android.tools.smali.dexlib2.iface.reference.FieldReference
-import com.android.tools.smali.dexlib2.iface.reference.MethodReference
 
 lateinit var addLithoFilter: (String) -> Unit
     private set
@@ -66,17 +65,11 @@ val lithoFilterPatch = bytecodePatch(
      *   }
      * }
      *
-     * class ComponentContextParser {
-     *    public Component parseComponent() {
+     * class CreateComponentClass {
+     *    public Component createComponent() {
      *        ...
      *
-     *        // Checks if the component should be filtered.
-     *        // Sets a thread local with the filtering result.
-     *        extensionClass.filter(identifier, pathBuilder);  // Inserted by this patch.
-     *
-     *        ...
-     *
-     *        if (extensionClass.shouldFilter()) {  // Inserted by this patch.
+     *        if (extensionClass.shouldFilter(identifier, path)) {  // Inserted by this patch.
      *            return emptyComponent;
      *        }
      *        return originalUnpatchedComponent; // Original code.
@@ -116,95 +109,68 @@ val lithoFilterPatch = bytecodePatch(
         // Allow the method to run to completion, and override the
         // return value with an empty component if it should be filtered.
         // It is important to allow the original code to always run to completion,
-        // otherwise memory leaks and poor app performance can occur.
-        //
-        // The extension filtering result needs to be saved off somewhere, but cannot
-        // save to a class field since the target class is called by multiple threads.
-        // It would be great if there was a way to change the register count of the
-        // method implementation and save the result to a high register to later use
-        // in the method, but there is no simple way to do that.
-        // Instead save the extension filter result to a thread local and check the
-        // filtering result at each method return index.
-        // String field for the litho identifier.
-        componentContextParserFingerprint.method.apply {
-            val conversionContextClass = conversionContextFingerprintToString.originalClassDef
-
-            val conversionContextIdentifierField = componentContextSubParserFingerprint.match(
-                componentContextParserFingerprint.originalClassDef
-            ).let {
-                // Identifier field is loaded just before the string declaration.
-                val index = it.method.indexOfFirstInstructionReversedOrThrow(
-                    it.stringMatches!!.first().index
-                ) {
-                    val reference = getReference<FieldReference>()
-                    reference?.definingClass == conversionContextClass.type
-                            && reference.type == "Ljava/lang/String;"
-                }
-                it.method.getInstruction<ReferenceInstruction>(index).getReference<FieldReference>()
+        // otherwise high memory usage and poor app performance can occur.
+
+        // Find the identifier/path fields of the conversion context.
+        val conversionContextIdentifierField = componentContextParserFingerprint.let {
+            // Identifier field is loaded just before the string declaration.
+            val index = it.method.indexOfFirstInstructionReversedOrThrow(
+                it.stringMatches!!.first().index
+            ) {
+                val reference = getReference<FieldReference>()
+                reference?.definingClass == conversionContextFingerprintToString.originalClassDef.type
+                        && reference.type == "Ljava/lang/String;"
             }
 
-            // StringBuilder field for the litho path.
-            val conversionContextPathBuilderField = conversionContextClass.fields
-                .single { field -> field.type == "Ljava/lang/StringBuilder;" }
+            it.method.getInstruction<ReferenceInstruction>(index).getReference<FieldReference>()!!
+        }
 
-            val conversionContextResultIndex = indexOfFirstInstructionOrThrow {
-                val reference = getReference<MethodReference>()
-                reference?.returnType == conversionContextClass.type
-            } + 1
+        val conversionContextPathBuilderField = conversionContextFingerprintToString.originalClassDef
+            .fields.single { field -> field.type == "Ljava/lang/StringBuilder;" }
 
-            val conversionContextResultRegister = getInstruction<OneRegisterInstruction>(
-                conversionContextResultIndex
-            ).registerA
+        // Find class and methods to create an empty component.
+        val builderMethodDescriptor = emptyComponentFingerprint.classDef.methods.single {
+            // The only static method in the class.
+                method -> AccessFlags.STATIC.isSet(method.accessFlags)
+        }
+        val emptyComponentField = classBy {
+            // Only one field that matches.
+            it.type == builderMethodDescriptor.returnType
+        }!!.immutableClass.fields.single()
+
+        componentCreateFingerprint.method.apply {
+            val insertIndex = if (is_19_17_or_greater) {
+                indexOfFirstInstructionOrThrow(Opcode.RETURN_OBJECT)
+            } else {
+                // 19.16 clobbers p2 so must check at start of the method and not at the return index.
+                0
+            }
 
-            val identifierRegister = findFreeRegister(
-                conversionContextResultIndex, conversionContextResultRegister
-            )
-            val stringBuilderRegister = findFreeRegister(
-                conversionContextResultIndex, conversionContextResultRegister, identifierRegister
-            )
+            val freeRegister = findFreeRegister(insertIndex)
+            val identifierRegister = findFreeRegister(insertIndex, freeRegister)
+            val pathRegister = findFreeRegister(insertIndex, freeRegister, identifierRegister)
 
-            // Check if the component should be filtered, and save the result to a thread local.
             addInstructionsAtControlFlowLabel(
-                conversionContextResultIndex + 1,
+                insertIndex,
                 """
-                    iget-object v$identifierRegister, v$conversionContextResultRegister, $conversionContextIdentifierField
-                    iget-object v$stringBuilderRegister, v$conversionContextResultRegister, $conversionContextPathBuilderField
-                    invoke-static { v$identifierRegister, v$stringBuilderRegister }, $EXTENSION_CLASS_DESCRIPTOR->filter(Ljava/lang/String;Ljava/lang/StringBuilder;)V
+                    move-object/from16 v$freeRegister, p2
+                    iget-object v$identifierRegister, v$freeRegister, $conversionContextIdentifierField
+                    iget-object v$pathRegister, v$freeRegister, $conversionContextPathBuilderField
+                    invoke-static { v$identifierRegister, v$pathRegister }, $EXTENSION_CLASS_DESCRIPTOR->shouldFilter(Ljava/lang/String;Ljava/lang/StringBuilder;)Z
+                    move-result v$freeRegister
+                    if-eqz v$freeRegister, :unfiltered
+                    
+                    # Return an empty component
+                    move-object/from16 v$freeRegister, p1
+                    invoke-static { v$freeRegister }, $builderMethodDescriptor
+                    move-result-object v$freeRegister
+                    iget-object v$freeRegister, v$freeRegister, $emptyComponentField
+                    return-object v$freeRegister
+        
+                    :unfiltered
+                    nop
                 """
             )
-
-            // Get the only static method in the class.
-            val builderMethodDescriptor = emptyComponentFingerprint.classDef.methods.single {
-                method -> AccessFlags.STATIC.isSet(method.accessFlags)
-            }
-            // Only one field.
-            val emptyComponentField = classBy { classDef ->
-                classDef.type == builderMethodDescriptor.returnType
-            }!!.immutableClass.fields.single()
-
-            // Check at each return value if the component is filtered,
-            // and return an empty component if filtering is needed.
-            findInstructionIndicesReversedOrThrow(Opcode.RETURN_OBJECT).forEach { returnIndex ->
-                val freeRegister = findFreeRegister(returnIndex)
-
-                addInstructionsAtControlFlowLabel(
-                    returnIndex,
-                    """
-                        invoke-static { }, $EXTENSION_CLASS_DESCRIPTOR->shouldFilter()Z
-                        move-result v$freeRegister
-                        if-eqz v$freeRegister, :unfiltered
-        
-                        move-object/from16 v$freeRegister, p1
-                        invoke-static { v$freeRegister }, $builderMethodDescriptor
-                        move-result-object v$freeRegister
-                        iget-object v$freeRegister, v$freeRegister, $emptyComponentField
-                        return-object v$freeRegister
-        
-                        :unfiltered
-                        nop    
-                    """
-                )
-            }
         }
 
         // endregion
