global _main

extern _printf, _puts, _getaddrinfo, _exit, _socket, _perror, _connect, _close, _send, _recv, _putchar, _fflush, _inet_ntop, _ntohs, _poll, _read, _write

default rel

section .text


_main: ;int(int argc,char **argv)
	enter 32, 0
	mov [rsp], rdi  ;argc
	mov [rsp+8], rsi  ;argv

	cmp rdi, 2
	jl .invalidusage
	cmp rdi, 3
	jg .invalidusage

	mov rax, [rsi+8]
	mov [.hoststrp], rax
	lea rax, [.str80]
	cmp rdi, 2
	cmovg rax, [rsi+16]
	mov [.portstrp], rax

	mov rdi, [.hoststrp]
	mov rsi, [.portstrp]
	lea rdx, [.hints_addrinfo]
	lea rcx, [.addrinfo_chain]
	; int 3
	call _getaddrinfo
	cmp rax, 0
	je .getaddrinfo_ok
	lea rdi, [.strGetaddrinfo]
	call _perror
	mov rax, 1
	leave
	ret
.getaddrinfo_ok:
	mov rdi, [.addrinfo_chain]
	call connectsocketfromaichain
	mov [.socket], eax
	cmp rax, -1
	jne .connectsuccess
	lea rdi, [.strfCouldntconnect]
	mov rsi, [.hoststrp]
	mov rdx, [.portstrp]
	mov rax, 0
	call _printf
	mov rax, 1
	leave
	ret
.connectsuccess:
	mov dword [.pollfdarr], 0
	mov word [.pollfdarr+4], POLLIN
	mov dword [.pollfdarr+8], eax
	mov word [.pollfdarr+12], POLLIN
	; mov rdi, rax
	; lea rsi, [.strHoi]
	; mov rdx, .strHoiLen
	; mov rcx, 0
	; call _send
.ioloop:
	mov word [.pollfdarr+6], 0
	mov word [.pollfdarr+14], 0
	lea rdi, [.pollfdarr]
	mov rsi, 2
	mov rdx, -1
	call _poll
	cmp rax, 0
	jle .pollerror  ;if equal, supposedly timeout expired
	mov ax, [.pollfdarr+6]
	cmp ax, 0
	jnz .stdinevent
	mov ax, [.pollfdarr+14]
	cmp ax, 0
	jnz .socketevent
	jmp .ioloop  ;poll didn't set any events; what?
.stdinevent:
	and ax, POLLIN
	jz .stdineventerror
	mov rdi, 0
	lea rsi, [.readbuf]
	mov rdx, .readbuf.len
	call _read
	cmp rax, 0
	je .stdineof
	jl .stdinreaderror
	mov rdi, [.socket]
	lea rsi, [.readbuf]
	mov rdx, rax
	call writefull
	cmp rax, 0
	jne .socketwriteerror
	jmp .ioloop
.socketevent:
	and ax, POLLIN
	jz .socketeventerror
	mov rdi, [.socket]
	lea rsi, [.readbuf]
	mov rdx, .readbuf.len
	call _read
	cmp rax, 0
	je .socketeof
	jl .socketreaderror
	mov rdi, 1
	lea rsi, [.readbuf]
	mov rdx, rax
	call writefull
	cmp rax, 0
	jne .stdoutwriteerror
	jmp .ioloop
.stdineventerror:
	lea rdi, [.strStdinerror]
	call _puts
	jmp .closeexit1
.socketeventerror:
	lea rdi, [.strSocketerror]
	call _puts
	jmp .closeexit1
.stdinreaderror:
.socketreaderror:
	lea rdi, [.strRead]
	call _perror
	jmp .closeexit1
.socketwriteerror:
.stdoutwriteerror:
	lea rdi, [.strWrite]
	call _perror
	jmp .closeexit1
.stdineof:
.socketeof:
.closeandexit:
	mov rdi, [.socket]
	call _close
	mov rax, 0
	leave
	ret
.pollerror:
	lea rdi, [.strPoll]
	call _perror
.closeexit1:
	mov rdi, [.socket]
	call _close
	mov rax, 1
	leave
	ret

.invalidusage:
	call printusage
	mov rax, 1
	leave
	ret

section .data
.strHoi: db "hoi!", 10, 0
.strHoiLen: equ $-.strHoi
.strPoll: db "poll", 0
.strRead: db "read", 0
.strWrite: db "write", 0
.strStdinerror: db "Error reading from stdin", 0
.strSocketerror: db "Error reading from socket", 0
.str80: db "80", 0
.strGetaddrinfo: db "getaddrinfo", 0
.strfCouldntconnect: db "Couldn't connect to '%s' port %s", 10, 0
.hints_addrinfo:                     ;struct addrinfo: //44
	dd AI_ADDRCONFIG|AI_NUMERICSERV  ;    int ai_flags; //4
	dd PF_UNSPEC                     ;    int ai_family; //4
	dd SOCK_STREAM                   ;    int ai_socktype; //4
	dd IPPROTO_TCP                   ;    int ai_protocol; //4
	dd 0                             ;    socklen_t ai_addrlen; //4
	dd 0                             ;     (padding) //4
	dq 0                             ;    char *ai_canonname; //8
	dq 0                             ;    struct sockaddr *ai_addr; //8
	dq 0                             ;    struct addrinfo *ai_next; //8
.addrinfo_chain: dq 0
section .bss
.socket: resd 1
.hoststrp: resq 1
.portstrp: resq 1
.pollfdarr: resb 2*8
.readbuf: resb 1024
.readbuf.len: equ $-.readbuf

section .text


printusage: ;void(int argc,char **argv)
	enter 0, 0
	lea rdi, [.formatstr]
	mov rsi, [rsi]
	mov rax, 0
	call _printf
	leave
	ret

section .data
.formatstr: db "nc clone in asm.", 10, \
               "Usage: %s <address> [port]", 10, \
               "Port is assumed 80 if not specified.", 10, 0

section .text


printargs: ;void(int argc,char **argv)
	enter 32, 0
	mov [rsp], rdi
	mov [rsp+8], rsi
	mov rsi, rdi
	lea rdi, [.formatstr1]
	mov rax, 0
	call _printf

	mov rcx, 0
.loop:
	lea rdi, [.formatstr2]
	mov rsi, rcx
	mov rdx, [rsp+8]
	mov rdx, [rdx+8*rcx]
	mov [rsp+16], rcx
	mov rax, 0
	call _printf
	mov rcx, [rsp+16]
	inc rcx
	cmp rcx, [rsp]
	jl .loop
	leave
	ret

section .data
.formatstr1: db "argc=%d", 10, 0
.formatstr2: db "argv[%d]=%s", 10, 0

section .text


connectsocketfromaichain: ;int(struct addrinfo *chain)
	enter 16, 0
	mov [rsp+8], r12

	mov r12, rdi  ;r12 contains the chain arg this whole function
	mov rsi, [r12+32]  ;.ai_addr
	mov rdi, 0
	mov di, [rsi+2]  ;sin[6]_port
	call _ntohs
	mov [rsp], ax  ;port in [rsp] (2 bytes)

	mov edi, [r12+4]  ;.ai_family
	mov rsi, [r12+32]  ;.ai_addr
	cmp rdi, AF_INET
	je .ipv4
	cmp rdi, AF_INET6
	je .ipv6
	;we won't do anything else than ipv{4,6}
	mov rsi, rdi
	lea rdi, [.strfUnknownAF]
	mov rax, 0
	call _printf
	jmp .trynextinchain
.ipv6:
	lea rsi, [rsi+8]  ;&(((struct sockaddr_in6*)ai_addr)->sin6_addr)
	jmp .callntop
.ipv4:
	lea rsi, [rsi+4]  ;&(((struct sockaddr_in*)ai_addr)->sin_addr)
.callntop:
	lea rdx, [.addrStringBuf]
	mov rcx, .addrStringBuf.len
	call _inet_ntop
	lea rdi, [.strfTrying]
	lea rsi, [.addrStringBuf]
	mov rdx, 0
	mov dx, [rsp]
	mov rax, 0
	call _printf

.createsocket:
	mov edi, [r12+4]  ;.ai_family
	mov esi, [r12+8]  ;.ai_socktype
	mov edx, [r12+12]  ;.ai_protocol
	call _socket
	mov [.socket], eax
	cmp rax, 0
	jl .socketfailed
	mov rdi, rax
	mov rsi, [r12+32]  ;.ai_addr
	mov edx, [r12+16]  ;.ai_addrlen
	call _connect
	cmp rax, 0
	jge .success
.connectfailed:
	mov edi, [.socket]
	call _close
	;int 3
	lea rdi, [.strConnectFailed]
	call _puts
	jmp .trynextinchain

.socketfailed:
	;int 3
	lea rdi, [.strSocketFailed]
	call _puts

.trynextinchain:
	mov rdi, r12
	mov rdi, [rdi+40]  ;.ai_next
	cmp rdi, 0
	jne .havenextinchain
	mov rax, -1
	mov r12, [rsp+8]
	leave
	ret
.havenextinchain:
	mov r12, [rsp+8]
	leave
	jmp connectsocketfromaichain  ;tail call

.success:
	mov eax, [.socket]
	mov r12, [rsp+8]
	leave
	ret

section .data
.socket: dd 0
.strfTrying: db "Trying [%s]:%hu...", 10, 0
.strfUnknownAF: db "Unknown address family %d!", 10, 0
.strConnectFailed: db "Connect failed", 0
.strSocketFailed: db "Socket failed", 0
section .bss
.addrStringBuf: resb 128
.addrStringBuf.len: equ $-.addrStringBuf

section .text


;returns 0 on success, <0 on error
writefull: ;int(int fd,const char *buf,int bufsz)
	enter 32, 0
	cmp rdx, 0
	jle .leaveret
	mov [rsp], rdi  ;fd
	mov [rsp+8], rsi  ;buf
	mov [rsp+16], rdx  ;bufsz

	mov rcx, 0  ;cursor
	mov [rsp+24], rcx

.wrloop:
	mov rdi, [rsp]
	mov rsi, [rsp+8]
	mov rdx, [rsp+16]
	sub rdx, rcx
	call _write
	cmp rax, -1
	jle .leaveret
	add rcx, rax
	cmp rcx, [rsp+16]
	jl .wrloop
	mov rax, 0
.leaveret:
	leave
	ret


section .data

AF_UNSPEC: equ 0  ;socket.h
AF_INET: equ 2  ;socket.h
AF_INET6: equ 30  ;socket.h
PF_UNSPEC: equ AF_UNSPEC  ;socket.h
SOCK_STREAM: equ 1  ;socket.h
IPPROTO_TCP: equ 6  ;netinet/in.h
AI_ADDRCONFIG: equ 0x400  ;netdb.h
AI_NUMERICSERV: equ 0x1000  ;netdb.h
AI_PASSIVE: equ 0x1  ;netdb.h (only for bind()-ready sockets)
POLLIN: equ 0x0001  ;poll.h
POLLPRI: equ 0x0002  ;poll.h
POLLOUT: equ 0x0004  ;poll.h
POLLRDNORM: equ 0x0040  ;poll.h
POLLWRNORM: equ POLLOUT  ;poll.h
POLLRDBAND: equ 0x0080  ;poll.h
POLLWRBAND: equ 0x0100  ;poll.h
POLLERR: equ 0x0008  ;poll.h
POLLHUP: equ 0x0010  ;poll.h
POLLNVAL: equ 0x0020  ;poll.h