Merge pull request #28 from byeongkeunahn/short-4

Shorten the generated code

Author: byeongkeunahn

scripts/base91.py

@@ -13,11 +13,11 @@ def encode(x_in, use_rle=False):
                 i += 1
             if len(x) > 0 and x[-1] == 0:
                 zeros_cnt = 1
-                while i - 1 + zeros_cnt < len(x_in) and zeros_cnt < 255 and x_in[i - 1 + zeros_cnt] == 0:
+                while i - 1 + zeros_cnt < len(x_in) and zeros_cnt < 256 and x_in[i - 1 + zeros_cnt] == 0:
                     zeros_cnt += 1
                 if zeros_cnt >= 2:
                     x.pop()
-                    x.append(zeros_cnt)
+                    x.append(zeros_cnt - 1)
                     sharp_insertion_points.append((current_bits // 13 * 2) + len(sharp_insertion_points))
                     i += zeros_cnt - 1
         sharp_insertion_points = list(reversed(sharp_insertion_points))

scripts/static-pie-elf2bin.py

@@ -160,7 +160,28 @@ def load_elf64(elf):
             continue        # since bytearray is zero-initialized
 
         dst_off, src_off, cnt = sh_dict['sh_addr'], sh_dict['sh_offset'], sh_dict['sh_size']
-        memory_bin[dst_off:dst_off+cnt] = elf[src_off:src_off+cnt]
+        blob = elf[src_off:src_off+cnt]
+
+        if sh_dict['sh_type'] == SHT_DYNAMIC:
+            # Trim the DYNAMIC section, leaving only relocation-related entries
+            # 16 == sizeof(Elf64_Dyn)
+            dst = 0 
+            for src in range(0, len(blob), 16):
+                # Included entries:
+                #   DT_PLTRELSZ = 2, DT_RELA = 7, DT_RELASZ = 8, DT_RELAENT = 9,
+                #   DT_REL = 17, DT_RELSZ = 18, DT_RELENT = 19, DT_PLTREL = 20,
+                #   DT_TEXT_REL = 22, DT_JMPREL = 23.
+                #
+                # Note: DT_RELACOUNT = 0x6fff_fff9 and DT_RELCOUNT = 0x6fff_fffa
+                #   are not included since they are redundant since
+                #   DT_RELACOUNT = DT_RELASZ/DT_RELAENT and
+                #   DT_RELCOUNT = DT_RELSZ/DT_RELENT.
+                if b2i(blob[src:src+8]) in [2, 7, 8, 9, 17, 18, 19, 20, 22, 23]:
+                    blob[dst:dst+16] = blob[src:src+16]
+                    dst += 16
+            blob[dst:] = bytearray(len(blob[dst:])) # fill remaining part with zeros
+
+        memory_bin[dst_off:dst_off+cnt] = blob
 
     entrypoint_offset = b2i(elf[24:32])
     return memory_bin, pos_begin, entrypoint_offset

scripts/static-pie-gen.py

@@ -68,6 +68,7 @@
 sol = "".join(sol)
 
 # binary (raw)
+# Since we append a little-endian 8-byte nonnegative integer, we can practically ensure that the last byte is zero.
 code_raw = memory_bin[:-8]
 code_raw += (len(code_raw) + 8 - loader_fdict['entrypoint_offset']).to_bytes(8, byteorder='little')
 code_raw_b91 = base91.encode(code_raw, use_rle=True).decode('ascii')
@@ -123,7 +124,7 @@
 
 # template
 template_candidates = [template_path]
-if lang_name == "Rust" and "x86_64" in target_name and "short" in template_path and len(code_raw) <= 4096:
+if lang_name == "Rust" and "x86_64" in target_name and "short" in template_path and len(code_raw) <= 4096 - 256:
     template_candidates.append(template_path.replace("short", "shorter"))
 
 out = None

scripts/static-pie-prestub-amd64-print.py

@@ -22,9 +22,14 @@
 prestub = prestub[:j]
 
 # settings
-SPECIFIER = ".quad"
-CHUNK_SIZE = 8
-ENTRIES_PER_LINE = 4
+if "--octa" in sys.argv:
+    SPECIFIER = ".octa"
+    CHUNK_SIZE = 16
+    ENTRIES_PER_LINE = 10
+else:
+    SPECIFIER = ".quad"
+    CHUNK_SIZE = 8
+    ENTRIES_PER_LINE = 4
 
 # pad to align at `CHUNK_SIZE`-byte boundary
 while len(prestub) % CHUNK_SIZE != 0:

scripts/static-pie-prestub-amd64-shorter.asm

@@ -19,8 +19,7 @@ _svc_alloc_rwx:
     xor     r9d, r9d                ; offset
     push    rsi                     ; save rsi
     xor     edi, edi                ; rdi=0
-    push    1
-    pop     rsi                     ; size
+    mov     esi, eax                ; size (anything in [1, 4096])
     mov     dl, 7                   ; protect (safe since we have ensured rdx=0)
     push    0x22
     pop     r10                     ; flags
@@ -39,27 +38,25 @@ _decode_loop:
     shl     eax, 13
 _decode_loop_2:
     lodsb
+    xor     ecx, ecx                ; ecx = 0
     sub     al, 0x23
-    cdq
-    jc      _jump_to_entrypoint
-    jz      _decode_zeros
+    jbe     _decode_zeros
     dec     al
-    xchg    eax, edx
+    xchg    eax, ecx
     lodsb
     sub     al, 0x24
     imul    eax, eax, 91
-    add     eax, edx
+    add     eax, ecx
 _decode_output:
     stosb
     shr     eax, 8
     test    ah, 16
     jnz     _decode_output
     jmp     _decode_loop
 _decode_zeros:
-    dec     rdi
-    movzx   ecx, byte [rdi]
-    rep     stosb                   ; the fact we jumped to here ensures al=0
-    jmp     _decode_loop_2
+    xchg    byte [rdi-1], cl        ; ecx = cl = ((number of zeros) - 1), byte [rdi-1] = 0
+    rep     stosb                   ; we have made sure the last byte is zero (in the packager)
+    jz      _decode_loop_2
 
 ; Jump to entrypoint
 _jump_to_entrypoint:

scripts/static-pie-prestub-amd64-shorter.bin

@@ -1 +1 @@
-j	X�E1�V1�j
^�j"AZj�AX^H�����,#�r"t�Ȓ�,$k�[
����u���H������H+���
+j	X�E1�V1��Ʋj"AZj�AX^H�����1�,#v�ȑ�,$k�[
Ȫ����u����O��t�H+���

scripts/static-pie-template-amd64-shorter.rs

@@ -5,4 +5,4 @@
 $$$$solution_src$$$$
 }
 // SOLUTION END
-#[no_link]extern crate std;#[no_mangle]unsafe fn _start(){std::arch::asm!(".octa 6a07b25e016aff3156c931459958096ah,0e0c11fb097485e050f5841ff6a5a4122h,6b242cac92c8fe1674227299232cac0dh,48dfebf77510c4f608e8c1aad0015bc0h,0d7fff87f2b48d8ebaaf30fb60fcfffh",in("rsi")r$$$$binary_raw_base91$$$$.as_ptr())}
+#[no_link]extern crate std;#[no_mangle]unsafe fn _start(){std::arch::asm!(".octa 226a07b2c689ff3156c931459958096ah,0de0c11fb097485e050f5841ff6a5a41h,5bc06b242cac91c8fe1676232cc931ach,0ff4f86e0ebf77510c4f608e8c1aac801h,0d7fff87f2b48dc74aaf3h",in("rsi")r$$$$binary_raw_base91$$$$.as_ptr())}