Merge pull request #24 from byeongkeunahn/short-4

Shorten the generated code in short mode

.vscode/tasks.json

@@ -35,6 +35,14 @@
                 "$rustc"
             ],
         },
+        {
+            "label": "build-release-amd64-rs-short-submit",
+            "type": "shell",
+            "command": "mkdir -p ./target/x86_64-unknown-linux-gnu-short/release; ./release-rs.sh --features short > ./target/x86_64-unknown-linux-gnu-short/release/loader.rs; code --reuse-window ./target/x86_64-unknown-linux-gnu-short/release/loader.rs",
+            "problemMatcher": [
+                "$rustc"
+            ],
+        },
         {
             "label": "build-release-i686-submit",
             "type": "shell",

scripts/static-pie-gen.py

@@ -68,7 +68,9 @@
 sol = "".join(sol)
 
 # binary (raw)
-code_raw_b91 = base91.encode(memory_bin[:-4], use_rle=True).decode('ascii')
+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')
 code_raw_b91_len = len(code_raw_b91)
 code_raw_b91 = '"' + code_raw_b91 + '"'
 
@@ -121,7 +123,7 @@
 
 # template
 template_candidates = [template_path]
-if lang_name == "Rust" and "x86_64" in target_name and "short" in template_path:
+if lang_name == "Rust" and "x86_64" in target_name and "short" in template_path and len(code_raw) <= 4096:
     template_candidates.append(template_path.replace("short", "shorter"))
 
 out = None

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

@@ -11,11 +11,6 @@ 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
@@ -34,20 +29,12 @@ _svc_alloc_rwx:
     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
+    mov     al, 0x1f                ; syscall preserves all registers except rcx, r11, rax; hence at this point rax=(previous rdi)=0
 _decode_loop:
     shl     eax, 13
 _decode_loop_2:
@@ -69,18 +56,12 @@ _decode_output:
     jnz     _decode_output
     jmp     _decode_loop
 _decode_zeros:
-    xchg    eax, edx
     movzx   ecx, byte [rdi-1]
     dec     rdi
-    rep     stosb
-    xchg    eax, edx
+    rep     stosb                   ; the fact we jumped to here ensures al=0
     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
+    sub     rdi, qword [rdi-8]
+    call    rdi

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!(".quad 9958096a000038c8h,16aff3156c93145h,6a5a41226a07b25eh,50525e050f5841ffh,0b0974850026a016ah,99232cac0de0c11fh,0ac92c8fe16742572h,0aad0015bc06b242ch,0f77510c4f608e8c1h,48ff4fb60f92dfebh,8bd5eb92aaf3cfffh,5954c8014859fc47h,53503",in("rsi")r$$$$binary_raw_base91$$$$.as_ptr())}
+#[no_link]extern crate std;#[no_mangle]unsafe fn _start(){std::arch::asm!(".quad 56c931459958096ah,6a07b25e016aff31h,0f5841ff6a5a4122h,0e0c11fb097485e05h,74237299232cac0dh,6b242cac92c8fe16h,8e8c1aad0015bc0h,0fdfebf77510c4f6h,0aaf3cfff48ff4fb6h,0d7fff87f2b48d7ebh",in("rsi")r$$$$binary_raw_base91$$$$.as_ptr())}

src/bin/codegen.rs

@@ -83,15 +83,15 @@ unsafe extern "win64" fn _start() -> ! {
     );
     // For "short", we always assume we are running with loader on Linux,
     // since "short" is only meaningful when submitting to online judges (not local test runs).
+    // Note that the stub will ensure that stack is aligned before caling _start.
+    // Also, for "short" on x86_64 Linux, we don't need PLATFORM_DATA, so we don't fabricate it.
     #[cfg(feature = "short")]
     asm!(
         "clc",                              // Not needed but packager wants it
-        "push   rcx",                       // Align stack
-        "mov    rbx, rcx",                  // Save PLATFORM_DATA table
+        "push   rax",                       // Align stack
         "lea    rdi, [rip + __ehdr_start]",
         "lea    rsi, [rip + _DYNAMIC]",
         "call   {0}",
-        "mov    rdi, rbx",
         "call   {1}",                       // This won't return since on Linux we invoke SYS_exitgroup in binary
         sym loader::amd64_elf::relocate,
         sym _start_rust,

src/platform/loader/amd64_elf.rs

@@ -49,6 +49,8 @@ There are currently three files licensed under GPLv2+:
 
 #![allow(clippy::cmp_null)]
 
+use core::mem::MaybeUninit;
+
 // Dynamic section entry types
 const DT_RELA:      u64 = 7;
 const DT_RELASZ:    u64 = 8;
@@ -78,14 +80,14 @@ pub unsafe extern "sysv64" fn relocate(
     ) {
     let mut ptr_dyn: *const Elf64Dyn = addr_dynamic_section as *const Elf64Dyn;
     let mut ptr_rela = 0;
-    let mut relasz = 0;
-    let mut relaent = 0;
+    let mut relasz = MaybeUninit::<u64>::uninit();
+    let mut relaent = MaybeUninit::<u64>::uninit();
     loop {
         match (*ptr_dyn).d_tag {
             0 => { break; }
             DT_RELA => { ptr_rela = addr_image_base + (*ptr_dyn).d_val_or_ptr; },
-            DT_RELASZ => { relasz = (*ptr_dyn).d_val_or_ptr; },
-            DT_RELAENT => { relaent = (*ptr_dyn).d_val_or_ptr; },
+            DT_RELASZ => { relasz.write((*ptr_dyn).d_val_or_ptr); },
+            DT_RELAENT => { relaent.write((*ptr_dyn).d_val_or_ptr); },
             _ => ()
         }
         ptr_dyn = ptr_dyn.add(1);
@@ -99,9 +101,9 @@ pub unsafe extern "sysv64" fn relocate(
     if ptr_rela == 0 {
         return;
     }
-    relasz += ptr_rela;
+    relasz.write(relasz.assume_init() + ptr_rela);
 
-    while ptr_rela < relasz {
+    while ptr_rela < relasz.assume_init() {
         let pst_rela = ptr_rela as *mut Elf64Rela;
         let ul_offset = (*pst_rela).r_offset;
         let ul_info = (*pst_rela).r_info;
@@ -116,6 +118,6 @@ pub unsafe extern "sysv64" fn relocate(
             /* not implemented */
             panic!();
         }
-        ptr_rela += relaent;
+        ptr_rela += relaent.assume_init();
     }
 }

src/platform/mod.rs

@@ -45,9 +45,14 @@ pub fn init(platform_data_by_loader: usize) {
 }
 #[cfg(not(test))]
 pub fn try_exit() {
-    let pd = services::platform_data();
-    if pd.env_id == services::ENV_ID_LINUX {
-        #[cfg(not(target_arch = "wasm32"))]
+    #[cfg(not(all(target_arch = "x86_64", feature = "short")))] {
+        let pd = services::platform_data();
+        if pd.env_id == services::ENV_ID_LINUX {
+            #[cfg(not(target_arch = "wasm32"))]
+            unsafe { os::linux::syscall::exit_group(services::get_exit_status() as usize); }
+        }
+    }
+    #[cfg(all(target_arch = "x86_64", feature = "short"))] {
         unsafe { os::linux::syscall::exit_group(services::get_exit_status() as usize); }
     }
 }