jit_tls_relax_test.c (8847B)
1 /* In-process JIT TLS Local-Exec relaxation for the Windows (COFF) TEB idiom. 2 * 3 * The single-threaded JIT rewrites the PE TLS-access idiom (read TEB 4 * ThreadLocalStoragePointer, index by `_tls_index`, then `+ SECREL(var)`) to 5 * address the in-image .tls instance directly, dropping the TEB/`_tls_index` 6 * indirection (see src/link/link_jit.c and the per-arch jit_tls_le_relax). The 7 * aarch64 path is verified end-to-end on the Windows VM; this pins the 8 * byte-level rewrite for BOTH arches, and is the only automated guard for the 9 * x86-64 idiom (x64-windows self-host can't run end-to-end yet). 10 * 11 * Each case hand-encodes the exact idiom the codegen emits 12 * (x64_tls_addr_of_win64 / aa_tls_addr_of_win), drives the relax, and asserts 13 * the result is the in-image address materialization. The x64 builder is pinned 14 * to a golden captured from real `kit cc -c` disassembly so it can't drift away 15 * from the codegen it must match. 16 * 17 * Exit 0 = pass; non-zero = fail. */ 18 19 #include <kit/core.h> 20 #include <stdint.h> 21 #include <string.h> 22 23 #include "core/core.h" 24 #include "lib/kit_unit.h" 25 #include "obj/obj.h" 26 27 /* Reached via LinkArchDesc.jit_tls_le_relax; not exposed in a header. */ 28 void x64_jit_tls_le_relax(Compiler* c, RelocKind k, u8* site, u64 storage, 29 u64 site_pc); 30 void aa64_jit_tls_le_relax(Compiler* c, RelocKind k, u8* site, u64 storage, 31 u64 site_pc); 32 33 static KitUnit g_u; 34 #define EXPECT(cond, ...) CU_EXPECT(&g_u, cond, __VA_ARGS__) 35 36 static KitCompiler* compiler_for(KitArchKind arch) { 37 static KitCompiler* aa64 = NULL; 38 static KitCompiler* x64 = NULL; 39 KitCompiler** slot = arch == KIT_ARCH_ARM_64 ? &aa64 : &x64; 40 if (!*slot) { 41 /* Windows/COFF so the compiler matches the idiom's origin; only the panic 42 * path reads it, so the exact spec barely matters. */ 43 KitTargetSpec t = kit_unit_target(arch, KIT_OS_WINDOWS, KIT_OBJ_COFF); 44 if (kit_unit_compiler_new(&g_u, t, slot) != KIT_OK || !*slot) { 45 fprintf(stderr, "compiler_new failed for arch=%d\n", (int)arch); 46 exit(2); 47 } 48 } 49 return *slot; 50 } 51 52 static u32 rd32(const u8* p) { 53 return (u32)p[0] | ((u32)p[1] << 8) | ((u32)p[2] << 16) | ((u32)p[3] << 24); 54 } 55 56 /* ---- x86-64 -------------------------------------------------------------- */ 57 58 /* Emit the 4-instruction Win64 TLS idiom for destination register `rd` into 59 * `b`, mirroring x64_tls_addr_of_win64 byte-for-byte. Returns total length; 60 * *site_off receives the offset of the trailing lea's SECREL disp32. */ 61 static u32 x64_emit_idiom(u8* b, u32 rd, u32* site_off) { 62 u32 n = 0; 63 /* (1) mov rd, gs:[0x58] (9 bytes) */ 64 b[n++] = 0x65; 65 b[n++] = (u8)(0x48u | ((rd & 8u) ? 0x04u : 0u)); 66 b[n++] = 0x8B; 67 b[n++] = (u8)(((rd & 7u) << 3) | 4u); 68 b[n++] = 0x25; 69 b[n++] = 0x58; 70 b[n++] = 0; 71 b[n++] = 0; 72 b[n++] = 0; 73 /* (2) mov r11d, [rip+_tls_index] (7 bytes) */ 74 b[n++] = 0x44; 75 b[n++] = 0x8B; 76 b[n++] = 0x1D; 77 b[n++] = 0; 78 b[n++] = 0; 79 b[n++] = 0; 80 b[n++] = 0; 81 /* (3) mov rd, [rd+r11*8] (4 or 5 bytes) */ 82 b[n++] = (u8)(0x4Au | ((rd & 8u) ? 0x05u : 0u)); 83 b[n++] = 0x8B; 84 if ((rd & 7u) == 5u) { 85 b[n++] = (u8)((1u << 6) | ((rd & 7u) << 3) | 4u); 86 b[n++] = (u8)(0xD8u | (rd & 7u)); 87 b[n++] = 0; 88 } else { 89 b[n++] = (u8)(((rd & 7u) << 3) | 4u); 90 b[n++] = (u8)(0xD8u | (rd & 7u)); 91 } 92 /* (4) lea rd, [rd + sym@SECREL] (7 or 8 bytes) */ 93 b[n++] = (u8)(0x48u | ((rd & 8u) ? 0x05u : 0u)); 94 b[n++] = 0x8D; 95 if ((rd & 7u) == 4u) { 96 b[n++] = (u8)((2u << 6) | ((rd & 7u) << 3) | 4u); 97 b[n++] = (u8)((4u << 3) | (rd & 7u)); 98 } else { 99 b[n++] = (u8)((2u << 6) | ((rd & 7u) << 3) | (rd & 7u)); 100 } 101 *site_off = n; 102 b[n++] = 0; 103 b[n++] = 0; 104 b[n++] = 0; 105 b[n++] = 0; 106 return n; 107 } 108 109 static void x64_check(u32 rd) { 110 u8 buf[40]; 111 u32 site_off; 112 u32 total = x64_emit_idiom(buf, rd, &site_off); 113 /* Distinct write/runtime "addresses": disp must use site_pc, not &buf. */ 114 const u64 site_pc = 0x140005000ull + site_off; 115 const u64 storage = 0x140009123ull; 116 u32 lea = total - 7u; /* the rewritten 7-byte rip-lea sits at the block end */ 117 i64 want_disp = (i64)storage - (i64)(site_pc + 4u); 118 u32 i; 119 x64_jit_tls_le_relax(compiler_for(KIT_ARCH_X86_64), R_COFF_SECREL, 120 &buf[site_off], storage, site_pc); 121 for (i = 0; i < lea; ++i) 122 EXPECT(buf[i] == 0x90u, "x64 rd=%u: byte %u not NOP (0x%02x)", rd, i, 123 buf[i]); 124 EXPECT(buf[lea] == (u8)(0x48u | ((rd >= 8u) ? 0x04u : 0u)), 125 "x64 rd=%u: lea REX 0x%02x", rd, buf[lea]); 126 EXPECT(buf[lea + 1] == 0x8Du, "x64 rd=%u: lea opcode 0x%02x", rd, 127 buf[lea + 1]); 128 EXPECT(buf[lea + 2] == (u8)(((rd & 7u) << 3) | 5u), 129 "x64 rd=%u: lea modrm 0x%02x", rd, buf[lea + 2]); 130 EXPECT((i32)rd32(&buf[lea + 3]) == (i32)want_disp, 131 "x64 rd=%u: disp 0x%08x want 0x%08x", rd, rd32(&buf[lea + 3]), 132 (u32)want_disp); 133 } 134 135 /* Pin x64_emit_idiom to real codegen: the rd=r8 idiom captured from 136 * `kit cc -target x86_64-windows -c` disassembly. A drift in the builder (and 137 * thus a relax tested against a fictional idiom) turns this red. */ 138 static void x64_golden(void) { 139 static const u8 want[] = {0x65, 0x4c, 0x8b, 0x04, 0x25, 0x58, 0x00, 140 0x00, 0x00, 0x44, 0x8b, 0x1d, 0x00, 0x00, 141 0x00, 0x00, 0x4f, 0x8b, 0x04, 0xd8, 0x4d, 142 0x8d, 0x80, 0x00, 0x00, 0x00, 0x00}; 143 u8 buf[40]; 144 u32 site_off; 145 u32 total = x64_emit_idiom(buf, 8u /* r8 */, &site_off); 146 EXPECT(total == sizeof want, "x64 golden length %u != %zu", total, 147 sizeof want); 148 EXPECT(memcmp(buf, want, sizeof want) == 0, "x64 golden idiom mismatch"); 149 } 150 151 /* ---- aarch64 ------------------------------------------------------------- */ 152 153 #define AA_NOP 0xd503201fu 154 155 static u32 aa_add_imm(u32 rd, u32 rn, u32 imm12, u32 sh) { 156 return 0x91000000u | (sh << 22) | ((imm12 & 0xfffu) << 10) | 157 ((rn & 0x1fu) << 5) | (rd & 0x1fu); 158 } 159 160 /* Emit the 7-instruction Win64/aarch64 TLS idiom for `rd` (aa_tls_addr_of_win). 161 * Only instruction (6) — the HIGH12A add the relax keys on — needs a faithful 162 * encoding; the rest are placeholders the relax overwrites unconditionally. */ 163 static u32 aa_emit_idiom(u32* w, u32 rd) { 164 w[0] = 0xf9400240u | (18u << 5) | rd; /* ldr rd,[x18,#0x58] (placeholder) */ 165 w[1] = 0x90000010u; /* adrp x16,_tls_index */ 166 w[2] = aa_add_imm(16u, 16u, 0u, 0u); /* add x16,x16,:lo12: */ 167 w[3] = 0xb9400210u; /* ldr w16,[x16] */ 168 w[4] = 0xf8607a00u | rd; /* ldr rd,[rd,x16,lsl#3] */ 169 w[5] = aa_add_imm(rd, rd, 0u, 1u); /* add rd,rd,:secrel_hi12: (HIGH12A) */ 170 w[6] = aa_add_imm(rd, rd, 0u, 0u); /* add rd,rd,:secrel_lo12: (LOW12A) */ 171 return 5u; /* word index of the HIGH12A site the relax keys on */ 172 } 173 174 /* Decode ADRP+ADD (the relax's output) back to the absolute address it 175 * materializes, so we confirm it equals `storage`. */ 176 static u64 aa_decode_adrp_add(u32 adrp, u32 add, u64 adrp_pc) { 177 u32 immlo = (adrp >> 29) & 0x3u; 178 u32 immhi = (adrp >> 5) & 0x7ffffu; 179 i64 imm = (i64)(((u64)immhi << 2) | immlo); 180 if (imm & (1ll << 20)) imm -= (1ll << 21); /* sign-extend 21 bits */ 181 u64 page = (adrp_pc & ~0xfffull) + ((u64)imm << 12); 182 return page + ((add >> 10) & 0xfffu); 183 } 184 185 static void aa_check(u32 rd) { 186 u32 w[7]; 187 u32 site_i = aa_emit_idiom(w, rd); 188 const u64 site_pc = 0x140005000ull + site_i * 4u; 189 const u64 storage = 0x140009123ull; 190 u8* site = (u8*)&w[site_i]; 191 u32 i; 192 aa64_jit_tls_le_relax(compiler_for(KIT_ARCH_ARM_64), 193 R_COFF_AARCH64_SECREL_HIGH12A, site, storage, site_pc); 194 for (i = 0; i < 4u; ++i) /* instructions (1)-(4) -> NOP */ 195 EXPECT(w[i] == AA_NOP, "aa64 rd=%u: insn %u not NOP (0x%08x)", rd, i, w[i]); 196 EXPECT((w[4] & 0x9f000000u) == 0x90000000u && (w[4] & 0x1fu) == rd, 197 "aa64 rd=%u: insn(5) not ADRP rd (0x%08x)", rd, w[4]); 198 EXPECT((w[5] & 0xff800000u) == 0x91000000u && (w[5] & 0x1fu) == rd && 199 ((w[5] >> 5) & 0x1fu) == rd, 200 "aa64 rd=%u: insn(6) not ADD rd,rd,#imm (0x%08x)", rd, w[5]); 201 EXPECT(w[6] == AA_NOP, "aa64 rd=%u: insn(7) not NOP (0x%08x)", rd, w[6]); 202 EXPECT(aa_decode_adrp_add(w[4], w[5], site_pc - 4u) == storage, 203 "aa64 rd=%u: ADRP+ADD != storage", rd); 204 } 205 206 int main(void) { 207 /* Cover the x64 length-variant branches: rd&7==4 (rsp/r12 -> lea sib), 208 * rd&7==5 (rbp/r13 -> mov disp8), and rd>=8 (REX.R/B). */ 209 static const u32 x64_rds[] = {0, 1, 4, 5, 8, 12, 13}; 210 static const u32 aa_rds[] = {0, 1, 9, 20}; 211 size_t i; 212 kit_unit_init(&g_u); 213 x64_golden(); 214 for (i = 0; i < sizeof x64_rds / sizeof x64_rds[0]; ++i) 215 x64_check(x64_rds[i]); 216 for (i = 0; i < sizeof aa_rds / sizeof aa_rds[0]; ++i) aa_check(aa_rds[i]); 217 kit_unit_summary(&g_u, "jit_tls_relax_test"); 218 return kit_unit_status(&g_u); 219 }