pd/pd.asm

; vim:ft=fasm:

include 'format/format.inc'

include 'utils.inc'
include 'struct.inc'

format ELF64 executable 3

segment readable writeable
envp dq ?
print_buf.len = 30
print_buf rb print_buf.len
print_buf.end = $


segment readable
nl                static_string 10
usage             static_string ' <pid> [cmd]', 10,\
                                9, 'pid', 9, 'process id to wait for', 10,\
                                9, 'cmd', 9, "optionall command to run with '/bin/sh' after <pid> exits", 10
no_prog           static_string 'ERROR: failed to get program name (arg0)', 10
ston_failed       static_string 'ERROR: failed to convert string to number: "'
ston_failed.end   static_string '"', 10
pidfd_open_failed static_string 'ERROR: pidfd_open failed: error code: '
execve_failed     static_string 'ERROR: execve failed', 10

bin_sh            static_string '/bin/sh', 0
bin_sh.cmd        static_string '-c', 0

repeat 1, max:PID_MAX
pid_too_big static_string 'ERROR: pid too big (max is ', `max, ')', 10
end repeat


segment readable executable
include 'stdlib.inc'
entry0:
	entry .
	prep_args
	mov [envp], rdx
	main rdi, rsi
	exit rax

macro main argc*, argv*
	call_args main_impl, argc, argv
end macro
proc main_impl
	.prog = r15
	.pid  = r14
	.fd = r13d
	.cmd = r12

	cmp rdi, 1
	jl .no_prog
	mov .prog, [rsi + (0 * qword)]

	cmp rdi, 2
	jl .no_pid
	mov .pid, [rsi + (1 * qword)]

	mov .cmd, [rsi + (2 * qword)]
	cmp rdi, 3
	jge .cmd_done
.no_cmd:
	xor .cmd, .cmd
.cmd_done:

	stonZ u10, .pid
	jz .ston_failed
	mov .pid, rax
	
	cmp .pid, PID_MAX
	jg .pid_too_big
	
	signal SIG.S_HUP, SIG.H_IGN

	pidfd_open .pid, 0
	mov .fd, eax
	cmp .fd, 0
	jl .pidfd_open_failed
	using_regs rax
		stack_frame sizeof.pollfd
			.pollfd = stack_frame.use(sizeof.pollfd)
			lea rax, [.pollfd]
			mov dword [rax + pollfd.fd],      .fd
			mov  word [rax + pollfd.events],  pool.IN
			mov  word [rax + pollfd.revents], 0
			; TODO: add option to wait many pids (any or all)
			poll rax, 1, -1
		end stack_frame
	end using_regs
	close rax

	test .cmd, .cmd
	jnz .run_cmd
	xor rax, rax
	ret

.run_cmd:
	stack_frame (4 * qword)
		.args = stack_frame.use(4 * qword)
		lea rax, [.args]
		mov qword [rax + (0 * qword)], bin_sh
		mov qword [rax + (1 * qword)], bin_sh.cmd
		mov qword [rax + (2 * qword)], .cmd
		mov qword [rax + (3 * qword)], 0
		execve bin_sh, rax, [envp]
	end stack_frame

	; normally unreachable
	write STDERR_FILENO, execve_failed, execve_failed.len
	jmp .err

.no_pid:
	writeZ STDERR_FILENO, .prog
	write  STDERR_FILENO, usage, usage.len
	xor rax, rax
	ret

.pidfd_open_failed:
	write STDERR_FILENO, pidfd_open_failed, pidfd_open_failed.len
	neg .fd
	print i, .fd, print_buf.end
	write STDERR_FILENO, rax, rcx
	write STDERR_FILENO, nl, nl.len
	jmp .err

.pid_too_big:
	write STDERR_FILENO, pid_too_big, pid_too_big.len
	jmp .err

.ston_failed:
	write STDERR_FILENO, ston_failed, ston_failed.len
	writeZ STDERR_FILENO, .pid
	write STDERR_FILENO, ston_failed.end, ston_failed.end.len
	jmp .err

.no_prog:
	write STDERR_FILENO, no_prog, no_prog.len

.err:
	mov rax, 1
end proc