Shorten the generated code in short mode

scripts/base91.py

@@ -1,4 +1,29 @@
-def encode(x):
+def encode(x_in, use_rle=False):
+ sharp_insertion_points = []
+ if use_rle:
+ current_bits, current_bytes, i = 0, 0, 0
+ x = bytearray()
+ while i < len(x_in):
+ current_bits += 13
+ while current_bytes < current_bits // 8:
+ if i >= len(x_in):
+ break
+ x.append(x_in[i])
+ current_bytes += 1
+ 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:
+ zeros_cnt += 1
+ if zeros_cnt >= 2:
+ x.pop()
+ x.append(zeros_cnt)
+ sharp_insertion_points.append((current_bits // 13 * 2) + len(sharp_insertion_points))
+ i += zeros_cnt - 1
+ sharp_insertion_points = list(reversed(sharp_insertion_points))
+ else:
+ x = x_in
+
  out = []
  i = 0
  cnt5, stack5 = 0, []
@@ -21,11 +46,14 @@ def encode(x):
  if bits < 13:
  stack5.append((pos, bits))
  cnt5 -= 8
- ret = bytearray(2 * len(out) + 1)
- for i in range(len(out)):
- ret[2 * i + 0] = 0x24 + (out[i] % 91)
- ret[2 * i + 1] = 0x24 + (out[i] // 91)
- ret[-1] = ord('!')
+ ret = bytearray()
+ for pack in out:
+ ret.append(0x24 + (pack % 91))
+ ret.append(0x24 + (pack // 91))
+ if len(sharp_insertion_points) > 0 and len(ret) == sharp_insertion_points[-1]:
+ ret.append(ord(b'#'))
+ sharp_insertion_points.pop()
+ ret.append(ord(b'!'))
  return bytes(ret)
 
 if __name__ == '__main__':

scripts/static-pie-gen.py

@@ -67,6 +67,11 @@
  sol[-1] = sol[-1].rstrip()
 sol = "".join(sol)
 
+# binary (raw)
+code_raw_b91 = base91.encode(memory_bin[:-4], use_rle=True).decode('ascii')
+code_raw_b91_len = len(code_raw_b91)
+code_raw_b91 = '"' + code_raw_b91 + '"'
+
 # binary
 with open(compressed_binary_path, "rb") as f:
  code = f.read()
@@ -115,29 +120,39 @@
 stub_b85 = '"' + stub_b85 + '"'
 
 # template
-with open(template_path, encoding='utf8') as f:
- template = f.read()
-template = template.replace("\ufeff", "")
-
-# putting it all together
-# reference: https://stackoverflow.com/a/15448887
-def multiple_replace(string, rep_dict):
- pattern = re.compile("|".join([re.escape(k) for k in sorted(rep_dict,key=len,reverse=True)]), flags=re.DOTALL)
- return pattern.sub(lambda x: rep_dict[x.group(0)], string)
-
-out = multiple_replace(template, {
- "$$$$solution_src$$$$": sol,
- "$$$$stub_raw$$$$": stub_raw,
- "$$$$stub_base85$$$$": stub_b85,
- "$$$$stub_len$$$$": str(len(stub)),
- "$$$$stub_base85_len$$$$": str(stub_b85_len),
- "$$$$stub_base91$$$$": stub_b91,
- "$$$$stub_base91_len$$$$": str(stub_b91_len),
- "$$$$binary_base85$$$$": r,
- "$$$$binary_base85_len$$$$": str(len(code_b85)),
- "$$$$binary_base91$$$$": code_b91,
- "$$$$binary_base91_len$$$$": str(code_b91_len),
- "$$$$min_len_4096$$$$": str(min(len(code_b85)+1, 4096)),
- "$$$$entrypoint_offset$$$$": str(loader_fdict['entrypoint_offset']),
-})
+template_candidates = [template_path]
+if lang_name == "Rust" and "x86_64" in target_name and "short" in template_path:
+ template_candidates.append(template_path.replace("short", "shorter"))
+
+out = None
+for each_template_path in template_candidates:
+ with open(each_template_path, encoding='utf8') as f:
+ template = f.read()
+ template = template.replace("\ufeff", "")
+
+ # putting it all together
+ # reference: https://stackoverflow.com/a/15448887
+ def multiple_replace(string, rep_dict):
+ pattern = re.compile("|".join([re.escape(k) for k in sorted(rep_dict,key=len,reverse=True)]), flags=re.DOTALL)
+ return pattern.sub(lambda x: rep_dict[x.group(0)], string)
+
+ out_candidate = multiple_replace(template, {
+ "$$$$solution_src$$$$": sol,
+ "$$$$stub_raw$$$$": stub_raw,
+ "$$$$stub_base85$$$$": stub_b85,
+ "$$$$stub_len$$$$": str(len(stub)),
+ "$$$$stub_base85_len$$$$": str(stub_b85_len),
+ "$$$$stub_base91$$$$": stub_b91,
+ "$$$$stub_base91_len$$$$": str(stub_b91_len),
+ "$$$$binary_base85$$$$": r,
+ "$$$$binary_base85_len$$$$": str(len(code_b85)),
+ "$$$$binary_base91$$$$": code_b91,
+ "$$$$binary_base91_len$$$$": str(code_b91_len),
+ "$$$$binary_raw_base91$$$$": code_raw_b91,
+ "$$$$binary_raw_base91_len$$$$": str(code_raw_b91_len),
+ "$$$$min_len_4096$$$$": str(min(len(code_b85)+1, 4096)),
+ "$$$$entrypoint_offset$$$$": str(loader_fdict['entrypoint_offset']),
+ })
+ if out is None or len(out_candidate) < len(out):
+ out = out_candidate
 print(out)

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

@@ -0,0 +1,86 @@
+; -*- tab-width: 4 -*-
+; 
+; The prestub for amd64-rust target
+; (prestub: the code that runs before the stub and sets the stage)
+;
+; build: nasm -f bin -O9 static-pie-prestub-amd64-shorter.asm -o static-pie-prestub-amd64-shorter.bin
+; note: after building with the above command, run static-pie-prestub-amd64-print.py static-pie-prestub-amd64-shorter.bin
+; to obtain the form that can be embedded in Rust as inline assembly.
+
+BITS 64
+ORG 0
+section .text
+
+; Align stack to 16 byte boundary
+; [rsp+ 32, rsp+120): PLATFORM_DATA
+; [rsp+ 0, rsp+ 32): (shadow space for win64 calling convention)
+ enter 56, 0
+
+; svc_alloc_rwx for Linux
+_svc_alloc_rwx:
+ push 9
+ pop rax ; syscall id of x64 mmap
+ cdq ; rdx=0
+ xor r9d, r9d ; offset
+ push rsi ; save rsi
+ xor edi, edi ; rdi=0
+ push 1
+ pop rsi ; size
+ mov dl, 7 ; protect (safe since we have ensured rdx=0)
+ push 0x22
+ pop r10 ; flags
+ push -1
+ pop r8 ; fd
+ syscall
+ pop rsi ; restore rsi
+
+; PLATFORM_DATA
+_t: ; PLATFORM_DATA[32..39] = ptr_alloc_rwx
+ push rdx ; PLATFORM_DATA[24..31] = win_GetProcAddress
+ push rax ; PLATFORM_DATA[16..23] = win_kernel32
+ push 1 ; PLATFORM_DATA[ 8..15] = env_flags (0=None, 1=ENV_FLAGS_LINUX_STYLE_CHKSTK)
+ push 2 ; PLATFORM_DATA[ 0.. 7] = env_id (1=Windows, 2=Linux)
+
+; Current state: rax = new buffer
+ push rax
+ xchg rax, rdi ; rdi = new buffer
+
+; Base91 decoder
+_decode:
+ mov al, 0x1f ; syscall preserves rax; hence at this point rax=9
+_decode_loop:
+ shl eax, 13
+_decode_loop_2:
+ lodsb
+ sub al, 0x23
+ cdq
+ jc _jump_to_entrypoint
+ jz _decode_zeros
+ dec al
+ xchg eax, edx
+ lodsb
+ sub al, 0x24
+ imul eax, eax, 91
+ add eax, edx
+_decode_output:
+ stosb
+ shr eax, 8
+ test ah, 16
+ jnz _decode_output
+ jmp _decode_loop
+_decode_zeros:
+ xchg eax, edx
+ movzx ecx, byte [rdi-1]
+ dec rdi
+ rep stosb
+ xchg eax, edx
+ jmp _decode_loop_2
+
+; Jump to entrypoint
+_jump_to_entrypoint:
+ mov eax, dword [rdi-4]
+ pop rcx
+ add rax, rcx
+ push rsp
+ pop rcx
+ call rax

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

@@ -0,0 +1,8 @@
+// Generated with https://github.com/kiwiyou/basm-rs
+// Learn rust (https://doc.rust-lang.org/book/) and get high performance out of the box!
+// SOLUTION BEGIN
+#![crate_type="cdylib"]#![no_std]#[cfg(any())]mod x{
+$$$$solution_src$$$$
+}
+// SOLUTION END
+#[no_link]extern crate std;#[no_mangle]unsafe fn _start(){std::arch::asm!(".quad 9958096a000038c8h,16aff3156c93145h,6a5a41226a07b25eh,50525e050f5841ffh,0b0974850026a016ah,99232cac0de0c11fh,0ac92c8fe16742572h,0aad0015bc06b242ch,0f77510c4f608e8c1h,48ff4fb60f92dfebh,8bd5eb92aaf3cfffh,5954c8014859fc47h,53503",in("rsi")r$$$$binary_raw_base91$$$$.as_ptr())}