Added module randomization

selfpatch/include/spslr.h

@@ -17,6 +17,9 @@ struct spslr_module {
 	const void* dpins;
 };
 
+void spslr_init(void);
 void spslr_selfpatch(void);
+void spslr_patch_module(const struct spslr_module* m);
+void spslr_cleanup(void);
 
 #endif

selfpatch/src/spslr.c

@@ -6,19 +6,25 @@
 
 #define SPSLR_SANITY_CHECK
 
-static void spslr_selfpatch_patch_dpins(void);
+static void spslr_patch_dpins(const struct spslr_dpin* dpins, spslr_u32 cnt);
-static void spslr_selfpatch_patch_dpin(void* addr, spslr_u32 target);
+static void spslr_patch_dpin(void* addr, spslr_u32 target);
-static void spslr_selfpatch_patch_ipins(void);
+static void spslr_patch_ipins(const struct spslr_ipin* ipins, spslr_u32 cnt,
+	const struct spslr_ipin_op* ipin_ops, spslr_u32 op_cnt);
 
 static void reorder_object(void* dst, const void* src, spslr_u32 target);
-static spslr_s64 spslr_calculate_ipin_value(spslr_u32 start);
+static spslr_s64 spslr_calculate_ipin_value(const struct spslr_ipin_op* ipin_ops, spslr_u32 op_cnt, spslr_u32 start);
 
 static void* reorder_buffer = NULL;
 
 static void allocate_reorder_buffer(void);
 static void release_reorder_buffer(void);
 
-void __init spslr_selfpatch(void) {
+static int initialized = 0;
+
+void __init spslr_init(void) {
+	if (initialized)
+		return;
+
 	if (spslr_randomizer_init() < 0)
 		spslr_env_panic("failed to initialize randomizer");
 
@@ -43,11 +49,43 @@ void __init spslr_selfpatch(void) {
 #endif
 
 	allocate_reorder_buffer();
-	spslr_selfpatch_patch_dpins();
+
+	initialized = 1;
+}
+
+void __init spslr_selfpatch(void) {
+	if (!initialized)
+		spslr_env_panic("tried to selfpatchpatch with uninitialized spslr context");
+
+	spslr_patch_dpins(spslr_dpins, spslr_dpin_cnt);
+	spslr_patch_ipins(spslr_ipins, spslr_ipin_cnt, spslr_ipin_ops, spslr_ipin_op_cnt);
+}
+
+void __init spslr_patch_module(const struct spslr_module* m) {
+	if (!initialized)
+		spslr_env_panic("tried to patch module with uninitialized spslr context");
+
+	if (!m || !m->ipin_cnt || !m->ipins || !m->ipin_op_cnt || !m->ipin_ops || !m->dpin_cnt || !m->dpins)
+		spslr_env_panic("tried to patch module with incomplete spslr data");
+
+	spslr_u32 module_ipin_cnt = *(const spslr_u32*)m->ipin_cnt;
+	const struct spslr_ipin* module_ipins = (const struct spslr_ipin*)m->ipins;
+	spslr_u32 module_ipin_op_cnt = *(const spslr_u32*)m->ipin_op_cnt;
+	const struct spslr_ipin_op* module_ipin_ops = (const struct spslr_ipin_op*)m->ipin_ops;
+	spslr_u32 module_dpin_cnt = *(const spslr_u32*)m->dpin_cnt;
+	const struct spslr_dpin* module_dpins = (const struct spslr_dpin*)m->dpins;
+
+	spslr_patch_dpins(module_dpins, module_dpin_cnt);
+	spslr_patch_ipins(module_ipins, module_ipin_cnt, module_ipin_ops, module_ipin_op_cnt);
+}
+
+void __init spslr_cleanup(void) {
+	if (!initialized)
+		spslr_env_panic("tried to clean up uninitialized spslr context");
+
+	initialized = 0;
+
 	release_reorder_buffer();
-
-	spslr_selfpatch_patch_ipins();
-
 	spslr_randomizer_clear();
 }
 
@@ -74,10 +112,10 @@ static void __init release_reorder_buffer(void) {
 	reorder_buffer = NULL;
 }
 
-static void __init spslr_selfpatch_patch_dpins(void) {
+static void __init spslr_patch_dpins(const struct spslr_dpin* dpins, spslr_u32 cnt) {
-	for (spslr_u32 dpidx = 0; dpidx < spslr_dpin_cnt; dpidx++) {
+	for (spslr_u32 dpidx = 0; dpidx < cnt; dpidx++)  {
-		const struct spslr_dpin* dp = &spslr_dpins[dpidx];
+		const struct spslr_dpin* dp = &dpins[dpidx];
-		spslr_selfpatch_patch_dpin((void*)dp->addr, dp->target);
+		spslr_patch_dpin((void*)dp->addr, dp->target);
 	}
 }
 
@@ -98,7 +136,7 @@ static void __init reorder_object(void* dst, const void* src, spslr_u32 target)
 	}
 }
 
-static void __init spslr_selfpatch_patch_dpin(void* addr, spslr_u32 target) {
+static void __init spslr_patch_dpin(void* addr, spslr_u32 target) {
 #ifdef SPSLR_SANITY_CHECK
 	if (target >= spslr_target_cnt)
 		spslr_env_panic("dpin refers to invalid target %u", target);
@@ -111,11 +149,12 @@ static void __init spslr_selfpatch_patch_dpin(void* addr, spslr_u32 target) {
 	spslr_env_poke_data(addr, reorder_buffer, t->size);
 }
 
-static void __init spslr_selfpatch_patch_ipins(void) {
+static void __init spslr_patch_ipins(const struct spslr_ipin* ipins, spslr_u32 cnt,
-	for (spslr_u32 ipidx = 0; ipidx < spslr_ipin_cnt; ipidx++) {
+		const struct spslr_ipin_op* ipin_ops, spslr_u32 op_cnt) {
-		const struct spslr_ipin* ip = &spslr_ipins[ipidx];
+	for (spslr_u32 ipidx = 0; ipidx < cnt; ipidx++) {
+		const struct spslr_ipin* ip = &ipins[ipidx];
 
-		spslr_s64 value = spslr_calculate_ipin_value(ip->program);
+		spslr_s64 value = spslr_calculate_ipin_value(ipin_ops, op_cnt, ip->program);
 
 		switch (ip->size) {
 			case 1:
@@ -138,19 +177,19 @@ static void __init spslr_selfpatch_patch_ipins(void) {
 	}
 }
 
-static spslr_s64 __init spslr_calculate_ipin_value(spslr_u32 start) {
+static spslr_s64 __init spslr_calculate_ipin_value(const struct spslr_ipin_op* ipin_ops, spslr_u32 op_cnt, spslr_u32 start) {
 	spslr_s64 res = 0;
 
 	spslr_u32 pc = start;
 	while (1) {
 #ifdef SPSLR_SANITY_CHECK
-		if (pc >= spslr_ipin_op_cnt)
+		if (pc >= op_cnt)
 			spslr_env_panic("ipin op %u out of bounds", pc);
 #endif
 
 		int end_flag = 0;
 
-		const struct spslr_ipin_op* op = &spslr_ipin_ops[pc++];
+		const struct spslr_ipin_op* op = &ipin_ops[pc++];
 		switch (op->code) {
 			case SPSLR_IPIN_OP_PATCH:
 				end_flag = 1;

subject/main.c

@@ -71,7 +71,7 @@ static int do_module_test_access_pid(const char *path, const struct task_struct
 		return -1;
 	}
 
-	// TODO -> Patch module
+	spslr_patch_module(&mod);
 
 	module_test_access_fn fn = (module_test_access_fn)dlsym(handle, "module_test_access_pid");
 
@@ -112,7 +112,9 @@ void arr_test() {
 
 int main(void)
 {
+    spslr_init();
     spslr_selfpatch();
+
     print_layout();
 
     struct list_head task_list;
@@ -150,6 +152,7 @@ int main(void)
     int module_pid = do_module_test_access_pid("./subject/libspslr_module.so", &global);
     printf("module_test_access returned %d (should be 42)\n", module_pid);
 
+    spslr_cleanup();
     return 0;
 }