From 792c1c14d77dc6013ab730fc8d5267ac07864916 Mon Sep 17 00:00:00 2001
From: Jiao Fenfang <[email protected]>
Date: Wed, 26 Jan 2022 14:37:36 +0800
Subject: [PATCH] support arm64

---
 src/toa.c | 93 ++++++++++++++++++++++++++++++++++++++++++++++++++-----
 src/toa.h |  2 ++
 2 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/src/toa.c b/src/toa.c
index e02ef24..2942674 100644
--- a/src/toa.c
+++ b/src/toa.c
@@ -23,6 +23,50 @@ struct toa_stats_entry toa_stats[] = {
 
 struct toa_stat_mib *ext_stats;
 
+#if  defined(CONFIG_ARM64)
+/*
+  * ARM64 interface
+  * Lookup the page table entry for a virtual address. Return a pointer
+  * to the entry and the level of the mapping.
+  *
+  * Note: We return pud and pmd either when the entry is marked large
+  * or when the present bit is not set. Otherwise we would return a
+  * pointer to a nonexisting mapping.
+  */
+static pte_t *lookup_address(unsigned long address, unsigned int *level)
+{
+	pgd_t *pgdp;
+	pud_t *pudp, pud;
+	pmd_t *pmdp, pmd;
+	pte_t *ptep;
+	unsigned long addr = address;
+
+	unsigned long init_mm = kallsyms_lookup_name("init_mm");
+	if (!init_mm) {
+		TOA_INFO("CPU [%u] lookup_address init_mm null\n", smp_processor_id());
+		return NULL;
+	}
+
+	/*pgdp = pgd_offset_k(addr);*/
+	pgdp = pgd_offset((struct mm_struct *)init_mm, addr);
+	if (pgd_none(READ_ONCE(*pgdp)))
+        	return NULL;
+
+	pudp = pud_offset(pgdp, addr);
+	pud = READ_ONCE(*pudp);
+	if (pud_none(pud))
+        	return NULL;
+
+	pmdp = pmd_offset(pudp, addr);
+	pmd = READ_ONCE(*pmdp);
+	if (pmd_none(pmd))
+        	return NULL;
+
+	ptep = pte_offset_kernel(pmdp, addr);
+	return ptep;
+}
+#endif
+
 /*
  * Funcs for toa hooks
  */
@@ -304,6 +348,7 @@ hook_toa_functions(void)
 	/* hook tcp_v4_syn_recv_sock for ipv4 */
 	struct inet_connection_sock_af_ops *ipv4_specific_p =
 			(struct inet_connection_sock_af_ops *)&ipv4_specific;
+
 #ifdef CONFIG_IP_VS_TOA_IPV6
 	/* hook inet6_getname for ipv6 */
 	struct proto_ops *inet6_stream_ops_p =
@@ -316,9 +361,18 @@ hook_toa_functions(void)
 	pte = lookup_address((unsigned long )inet_stream_ops_p, &level);
 	if (pte == NULL)
 		return 1;
-	if (pte->pte & ~_PAGE_RW) {
-		pte->pte |= _PAGE_RW;
-	}
+
+#if  defined(CONFIG_ARM64)
+        *pte = pte_mkwrite(*pte);
+	__sync_icache_dcache(*pte);
+	flush_tlb_all();
+#else
+        if (pte->pte & ~_PAGE_RW) {
+                pte->pte |= _PAGE_RW;
+        }
+#endif
+	/*TOA_INFO("CPU [%u] hooked before inet_getname <%p> --> <%p>\n",
+                smp_processor_id(), inet_getname, inet_getname_toa);*/
 
 	inet_stream_ops_p->getname = inet_getname_toa;
 	TOA_INFO("CPU [%u] hooked inet_getname <%p> --> <%p>\n",
@@ -345,7 +399,14 @@ hook_toa_functions(void)
 	pte = lookup_address((unsigned long )inet_stream_ops_p, &level);
 	if (pte == NULL)
 		return 1;
-	pte->pte |= pte->pte &~_PAGE_RW;
+
+#if  defined(CONFIG_ARM64)
+        *pte = pte_wrprotect(*pte);
+	__sync_icache_dcache(*pte);
+	flush_tlb_all();
+#else
+        pte->pte |= pte->pte &~_PAGE_RW;
+#endif
 
 	return 0;
 }
@@ -376,9 +437,18 @@ unhook_toa_functions(void)
 	pte = lookup_address((unsigned long )inet_stream_ops_p, &level);
 	if (pte == NULL)
 		return 1;
-	if (pte->pte & ~_PAGE_RW) {
-		pte->pte |= _PAGE_RW;
-	}
+
+#if  defined(CONFIG_ARM64)
+        *pte = pte_mkwrite(*pte);
+	__sync_icache_dcache(*pte);
+	flush_tlb_all();
+#else
+        if (pte->pte & ~_PAGE_RW) {
+                pte->pte |= _PAGE_RW;
+        }
+#endif
+	/*TOA_INFO("CPU [%u] unhooked before inet_getname: %p\n",
+                smp_processor_id(), inet_getname);*/
 
 	inet_stream_ops_p->getname = inet_getname;
 	TOA_INFO("CPU [%u] unhooked inet_getname\n",
@@ -403,7 +473,14 @@ unhook_toa_functions(void)
 	pte = lookup_address((unsigned long )inet_stream_ops_p, &level);
 	if (pte == NULL)
 		return 1;
-	pte->pte |= pte->pte &~_PAGE_RW;
+
+#if  defined(CONFIG_ARM64)
+        *pte = pte_wrprotect(*pte);
+	__sync_icache_dcache(*pte);
+	flush_tlb_all();
+#else
+        pte->pte |= pte->pte &~_PAGE_RW;
+#endif
 
 	return 0;
 }
diff --git a/src/toa.h b/src/toa.h
index b47fac3..a46f82d 100644
--- a/src/toa.h
+++ b/src/toa.h
@@ -20,6 +20,8 @@
 #include <net/ipv6.h>
 #include <net/transp_v6.h>
 
+#include <asm/tlbflush.h>
+
 #define TOA_VERSION "1.0.0.0"
 
 #ifdef TOA_DBG_MSG
-- 
2.20.1