Support dynamic registration for native methods using --dynamic-register (amimo#63)

* Add methods dynamic register supported

Author: Bzi-Han

.gitignore

@@ -0,0 +1,8 @@
+# VSCode
+.vscode
+
+# Test data
+data
+
+# Cache
+__pycache__

dcc.py

@@ -339,17 +339,17 @@ def archive_compiled_code(project_dir):
     outfile = shutil.make_archive(outfile, 'zip', project_dir)
     return outfile
 
-
-def compile_dex(apkfile, filtercfg):
+def compile_dex(apkfile, filtercfg, dynamic_register):
     show_logging(level=logging.INFO)
 
     d = auto_vm(apkfile)
     dx = analysis.Analysis(d)
 
     method_filter = MethodFilter(filtercfg, d)
 
-    compiler = Dex2C(d, dx)
+    compiler = Dex2C(d, dx, dynamic_register)
 
+    native_method_prototype = {}
     compiled_method_code = {}
     errors = []
 
@@ -372,20 +372,85 @@ def compile_dex(apkfile, filtercfg):
                 errors.append('%s:%s' % (full_name, str(e)))
                 continue
 
-            if code:
-                compiled_method_code[method_triple] = code
+            if code[0]:
+                compiled_method_code[method_triple] = code[0]
+                native_method_prototype[jni_longname] = code[1]
 
-    return compiled_method_code, errors
+    return compiled_method_code, native_method_prototype, errors
 
 def is_apk(name):
     return name.endswith('.apk')
 
-def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, source_archive='project-source.zip'):
+def write_dummy_dynamic_register(project_dir):
+    source_dir = os.path.join(project_dir, 'jni', 'nc')
+    if not os.path.exists(source_dir):
+        os.makedirs(source_dir)
+
+    filepath = os.path.join(source_dir, 'DynamicRegister.cpp')
+    with open(filepath, 'w') as fp:
+        fp.write('#include "DynamicRegister.h"\n\nconst char *dynamic_register_compile_methods(JNIEnv *env) { return nullptr; }')
+
+def write_dynamic_register(project_dir, compiled_methods, method_prototypes):
+    source_dir = os.path.join(project_dir, 'jni', 'nc')
+    if not os.path.exists(source_dir):
+        os.makedirs(source_dir)
+
+    export_list = {}
+
+    # Make export list
+    for method_triple in sorted(compiled_methods.keys()):
+        full_name = JniLongName(*method_triple)
+        if not full_name in method_prototypes:
+            raise Exception('Method %s prototype info could not be found' % full_name)
+
+        class_path = method_triple[0][1:-1].replace('.', '/')
+        method_name = method_triple[1]
+        method_signature = method_triple[2]
+        method_native_name = full_name
+        method_native_prototype = method_prototypes[full_name]
+
+        if not class_path in export_list:
+            export_list[class_path] = [] # methods
+        
+        export_list[class_path].append((method_name, method_signature, method_native_name, method_native_prototype))
+
+    if len(export_list) == 0:
+        logger.info('No export methods')
+        return
+    
+    # Generate extern block and export block
+    extern_block = []
+    export_block = ['\njclass clazz;\n']
+    export_block_template = 'clazz = env->FindClass("%s");\nif (clazz == nullptr)\n    return "Class not found: %s";\n'
+    export_block_template += 'const JNINativeMethod export_method_%d[] = {\n%s\n};\n'
+    export_block_template += 'env->RegisterNatives(clazz, export_method_%d, %d);\n'
+    export_block_template += 'env->DeleteLocalRef(clazz);\n'
+
+    for index, class_path in enumerate(sorted(export_list.keys())):
+        methods = export_list[class_path]
+
+        extern_block.append('\n'.join(['extern %s;' % method[3] for method in methods]))
+        
+        export_methods = ',\n'.join(['{"%s", "%s", (void *)%s}' % (method[0], method[1], method[2]) for method in methods])
+        export_block.append(export_block_template % (class_path, class_path, index, export_methods, index, len(methods)))
+
+    export_block.append('return nullptr;\n')
+
+    # Write DynamicRegister.cpp
+    filepath = os.path.join(source_dir, 'DynamicRegister.cpp')
+    with open(filepath, 'w') as fp:
+        fp.write('#include "DynamicRegister.h"\n\n')
+        fp.write('\n'.join(extern_block))
+        fp.write('\n\nconst char *dynamic_register_compile_methods(JNIEnv *env) {')
+        fp.write('\n'.join(export_block))
+        fp.write('}')
+
+def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, source_archive='project-source.zip', dynamic_register=False):
     if not os.path.exists(apkfile):
         logger.error("file %s is not exists", apkfile)
         return
 
-    compiled_methods, errors = compile_dex(apkfile, filtercfg)
+    compiled_methods, method_prototypes, errors = compile_dex(apkfile, filtercfg, dynamic_register)
 
     if errors:
         logger.warning('================================')
@@ -400,11 +465,22 @@ def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, sour
         if not os.path.exists(project_dir):
             shutil.copytree('project', project_dir)
         write_compiled_methods(project_dir, compiled_methods)
+
+        if dynamic_register:
+            write_dynamic_register(project_dir, compiled_methods, method_prototypes)
+        else:
+            write_dummy_dynamic_register(project_dir)
     else:
         project_dir = make_temp_dir('dcc-project-')
         shutil.rmtree(project_dir)
         shutil.copytree('project', project_dir)
         write_compiled_methods(project_dir, compiled_methods)
+
+        if dynamic_register:
+            write_dynamic_register(project_dir, compiled_methods, method_prototypes)
+        else:
+            write_dummy_dynamic_register(project_dir)
+
         src_zip = archive_compiled_code(project_dir)
         shutil.move(src_zip, source_archive)
 
@@ -428,6 +504,7 @@ def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, sour
     parser.add_argument('-o', '--out', nargs='?', help='Output APK file name')
     parser.add_argument('--sign', action='store_true', default=False, help='Sign apk')
     parser.add_argument('--filter', default='filter.txt', help='Method filter configure file')
+    parser.add_argument('--dynamic-register', action='store_true', default=False, help='Export native methods using RegisterNatives')
     parser.add_argument('--no-build', action='store_true', default=False, help='Do not build the compiled code')
     parser.add_argument('--source-dir', help='The compiled cpp code output directory.')
     parser.add_argument('--project-archive', default='project-source.zip', help='Archive the project directory')
@@ -439,6 +516,7 @@ def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, sour
     filtercfg = args['filter']
     do_compile = not args['no_build']
     source_archive = args['project_archive']
+    dynamic_register = args['dynamic_register']
 
     if args['source_dir']:
         project_dir = args['source_dir']
@@ -460,7 +538,7 @@ def dcc_main(apkfile, filtercfg, outapk, do_compile=True, project_dir=None, sour
         APKTOOL = dcc_cfg['apktool']
 
     try:
-        dcc_main(infile, filtercfg, outapk, do_compile, project_dir, source_archive)
+        dcc_main(infile, filtercfg, outapk, do_compile, project_dir, source_archive, dynamic_register)
     except Exception as e:
         logger.error("Compile %s failed!" % infile, exc_info=True)
     finally:

dex2c/compiler.py

@@ -101,13 +101,18 @@ def get_source(self):
         if self.writer:
             return str(self.writer)
         return ''
+    
+    def get_prototype(self):
+        if self.writer:
+            return self.writer.get_prototype()
+        return ''
 
     def __repr__(self):
         return 'class IrMethod(object): %s' % self.name
 
 
 class IrBuilder(object):
-    def __init__(self, methanalysis):
+    def __init__(self, methanalysis, dynamic_register):
         method = methanalysis.get_method()
         self.method = method
         self.irmethod = None
@@ -118,6 +123,7 @@ def __init__(self, methanalysis):
         self.var_to_name = defaultdict()
         self.offset_to_node = {}
         self.graph = None
+        self.dynamic_register = dynamic_register
 
         self.access = util.get_access_method(method.get_access_flags())
 
@@ -176,7 +182,7 @@ def process(self):
         irmethod.params = self.lparams
         irmethod.params_type = self.params_type
 
-        writer = Writer(irmethod)
+        writer = Writer(irmethod, self.dynamic_register)
         writer.write_method()
         irmethod.writer = writer
         return irmethod
@@ -547,18 +553,19 @@ def process_and_show(self):
 
 
 class Dex2C:
-    def __init__(self, vm, vmx):
+    def __init__(self, vm, vmx, dynamic_register):
         self.vm = vm
         self.vmx = vmx
+        self.dynamic_register = dynamic_register
 
     def get_source_method(self, m):
         mx = self.vmx.get_method(m)
-        z = IrBuilder(mx)
+        z = IrBuilder(mx, self.dynamic_register)
         irmethod = z.process()
         if irmethod:
-            return irmethod.get_source()
+            return (irmethod.get_source(), irmethod.get_prototype())
         else:
-            return None
+            return (None, None)
 
     def get_source_class(self, _class):
         c = DvClass(_class, self.vmx)

dex2c/writer.py

@@ -52,12 +52,14 @@ def get_name(self, item):
 
 
 class Writer(object):
-    def __init__(self, irmethod):
+    def __init__(self, irmethod, dynamic_register):
         self.graph = irmethod.graph
         self.method = irmethod.method
         self.irmethod = irmethod
         self.visited_nodes = set()
         self.buffer = []
+        self.dynamic_register = dynamic_register
+        self.prototype = []
 
         entry = irmethod.entry
         self.ra = irmethod.ra
@@ -67,6 +69,9 @@ def __init__(self, irmethod):
 
     def __str__(self):
         return ''.join(self.buffer)
+    
+    def get_prototype(self):
+        return ''.join(self.prototype)
 
     def write_trace(self, ins):
         s = ins.dump()
@@ -98,7 +103,13 @@ def write_method(self):
 
         self.write("\n/* %s->%s%s */\n" % (class_name, name, proto))
 
-        self.write('extern "C" JNIEXPORT %s JNICALL\n' % get_native_type(self.irmethod.rtype))
+        if self.dynamic_register:
+            self.write(get_native_type(self.irmethod.rtype) + ' ')
+            self.prototype.append(get_native_type(self.irmethod.rtype) + ' ')
+            self.prototype.append(jni_name)
+        else:
+            self.write('extern "C" JNIEXPORT %s JNICALL\n' % get_native_type(self.irmethod.rtype))
+            
         self.write(jni_name)
         params = self.irmethod.params
         if 'static' not in access:
@@ -110,8 +121,12 @@ def write_method(self):
                                zip(self.irmethod.params_type, params)])
         if proto:
             self.write('(JNIEnv *env, jobject thiz, %s)' % proto)
+            if self.dynamic_register:
+                self.prototype.append('(JNIEnv *env, jobject thiz, %s)' % proto)
         else:
             self.write('(JNIEnv *env, jobject thiz)')
+            if self.dynamic_register:
+                self.prototype.append('(JNIEnv *env, jobject thiz)')
         self.write('{\n')
         nodes = self.irmethod.irblocks
         for node in nodes:
@@ -145,7 +160,7 @@ def visit_node(self, node):
             var_declared.add(r)
             self.write('%s v%s' % (get_cdecl_type(var_type), r))
             if util.is_ref(var_type):
-                self.write(' = NULL');
+                self.write(' = NULL')
             self.write(";\n")
 
         if node.var_to_declare and self.irmethod.landing_pads:

project/jni/nc/Dex2C.cpp

@@ -6,6 +6,7 @@
 #include "ScopedLocalRef.h"
 #include "ScopedPthreadMutexLock.h"
 #include "well_known_classes.h"
+#include "DynamicRegister.h"
 
 struct MemberTriple {
     MemberTriple(const char *cls_name, const char *name, const char *sig):class_name_(cls_name), member_name_(name), signautre_(sig) {}
@@ -257,5 +258,11 @@ JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *reserved) {
         return JNI_ERR;
     }
     cache_well_known_classes(env);
+    const char *result = dynamic_register_compile_methods(env);
+    if (result != nullptr)
+    {
+        LOGD("d2c_throw_exception %s", result);
+        return JNI_ERR;
+    }
     return JNI_VERSION_1_6;
 }

project/jni/nc/DynamicRegister.h

@@ -0,0 +1,8 @@
+#ifndef _DYNAMIC_REGISTER_H_ // !_DYNAMIC_REGISTER_H_
+#define _DYNAMIC_REGISTER_H_
+
+#include <jni.h>
+
+const char *dynamic_register_compile_methods(JNIEnv *env);
+
+#endif // !_DYNAMIC_REGISTER_H_