Restoring stdin fcntl flags before exit

wtfstopw.asm

 	lapcountlen: equ $ - lapcount
 	laplen: dq 2
 
+	fcntl_flag: dq 0
+
 section .text
 global _start
 _start:
 mov rax, 72 ; fcntl
 mov rsi, 3 ; F_GETFL
 syscall
+mov [fcntl_flag], rax
 mov rdx, rax
 or rdx, 0x800 ; O_NONBLOCK
 mov rax, 72 ; fcntl
 	je newline_exit
 	jg flush_stdin
 
+mov rdi, 0 ; EXIT OK
+;
+; Expect rdi to be the return code
+;
 exit:
+	push rdi
+	; restoring stdin state
+	mov rax, 72 ; fcntl
+	xor rdi, rdi
+	mov rsi, 4 ; F_SETFL
+	mov rdx, [fcntl_flag]
+	syscall
+	cmp rax, 0
+	je exit_end
+	pop rdi ; failed to restore
+	push 1 ; exit FAIL
+
+	exit_end:
 	mov rax, 60 ; sys_exit
-	mov rdi, 0 ; OK
+	pop rdi ; return code
 	syscall
 fault:
 	mov rax, 1 ; write
 	mov rdx, faultmsglen
 	syscall
 
-	mov rax, 60 ; sys_exit
 	mov rdi, 1 ; failure
-	syscall
+	jmp exit
 
 newline_exit:
 	mov rax, 1
 	mov rsi, nl
 	mov rdx, 1
 	syscall
+
+	mov rdi, 0 ; exit OK
 	jmp exit
 
 ;