diff --git a/Userland/Utilities/nc.cpp b/Userland/Utilities/nc.cpp index 50a5cd1ddf..4d776449a6 100644 --- a/Userland/Utilities/nc.cpp +++ b/Userland/Utilities/nc.cpp @@ -5,6 +5,7 @@ */ #include +#include #include #include #include @@ -69,10 +70,11 @@ ErrorOr serenity_main(Main::Arguments arguments) } } - int fd; + int fd = -1; + int listen_fd = -1; if (should_listen) { - int listen_fd = TRY(Core::System::socket(AF_INET, SOCK_STREAM, 0)); + listen_fd = TRY(Core::System::socket(AF_INET, SOCK_STREAM, 0)); sockaddr_in sa {}; sa.sin_family = AF_INET; @@ -98,13 +100,6 @@ ErrorOr serenity_main(Main::Arguments arguments) if (verbose) warnln("waiting for a connection on {}:{}", inet_ntop(sin.sin_family, &sin.sin_addr, addr_str, sizeof(addr_str) - 1), ntohs(sin.sin_port)); - len = sizeof(sin); - TRY(Core::System::accept(listen_fd, (struct sockaddr*)&sin, &len)); - - if (verbose) - warnln("got connection from {}:{}", inet_ntop(sin.sin_family, &sin.sin_addr, addr_str, sizeof(addr_str) - 1), ntohs(sin.sin_port)); - - TRY(Core::System::close(listen_fd)); } else { fd = TRY(Core::System::socket(AF_INET, SOCK_STREAM, 0)); @@ -135,12 +130,15 @@ ErrorOr serenity_main(Main::Arguments arguments) warnln("connected!"); } + HashTable connected_clients; + bool stdin_closed = false; bool fd_closed = false; + bool listen_fd_closed = false; fd_set readfds, writefds, exceptfds; - while (!stdin_closed || !fd_closed) { + while (!stdin_closed || !fd_closed || !listen_fd_closed) { FD_ZERO(&readfds); FD_ZERO(&writefds); FD_ZERO(&exceptfds); @@ -152,12 +150,27 @@ ErrorOr serenity_main(Main::Arguments arguments) FD_SET(STDIN_FILENO, &exceptfds); highest_fd = max(highest_fd, STDIN_FILENO); } - if (!fd_closed) { + if (!fd_closed && fd) { FD_SET(fd, &readfds); FD_SET(fd, &exceptfds); highest_fd = max(highest_fd, fd); } + if (!listen_fd_closed && listen_fd) { + FD_SET(listen_fd, &readfds); + FD_SET(listen_fd, &exceptfds); + highest_fd = max(highest_fd, listen_fd); + } + + bool has_clients = (should_listen && !connected_clients.is_empty()); + if (has_clients) { + for (auto const& client_fd : connected_clients) { + FD_SET(client_fd, &readfds); + FD_SET(client_fd, &exceptfds); + highest_fd = max(highest_fd, client_fd); + } + } + int ready = select(highest_fd + 1, &readfds, &writefds, &exceptfds, nullptr); if (ready == -1) { if (errno == EINTR) @@ -179,11 +192,21 @@ ErrorOr serenity_main(Main::Arguments arguments) if (verbose) warnln("stdin closed"); if (should_close) { - TRY(Core::System::close(fd)); - fd_closed = true; + if (should_listen) { + TRY(Core::System::close(listen_fd)); + listen_fd_closed = true; + } else { + TRY(Core::System::close(fd)); + fd_closed = true; + } } } else { - TRY(Core::System::write(fd, buffer_span)); + if (should_listen && has_clients) { + for (auto const& client_fd : connected_clients) + TRY(Core::System::write(client_fd, buffer_span)); + } else { + TRY(Core::System::write(fd, buffer_span)); + } } } @@ -204,6 +227,44 @@ ErrorOr serenity_main(Main::Arguments arguments) TRY(Core::System::write(STDOUT_FILENO, buffer_span)); } } + + if (!listen_fd_closed && FD_ISSET(listen_fd, &readfds)) { + char client_str[INET_ADDRSTRLEN]; + sockaddr_in client; + socklen_t clientlen = sizeof(client); + + int new_client = TRY(Core::System::accept(listen_fd, (struct sockaddr*)&client, &clientlen)); + connected_clients.set(new_client); + + if (verbose) + warnln("got connection from {}:{}", inet_ntop(client.sin_family, &client.sin_addr, client_str, sizeof(client_str) - 1), ntohs(client.sin_port)); + } + + if (has_clients) { + for (auto const client_fd : connected_clients) { + if (FD_ISSET(client_fd, &readfds)) { + Array buffer; + Bytes buffer_span = buffer.span(); + auto nread = TRY(Core::System::read(client_fd, buffer_span)); + buffer_span = buffer_span.trim(nread); + + if (nread == 0) { + if (verbose) { + struct sockaddr_in client; + socklen_t clientlen = sizeof(client); + TRY(Core::System::getpeername(client_fd, (struct sockaddr*)&client, &clientlen)); + warnln("remote connection closed {}:{}", inet_ntoa(client.sin_addr), ntohs(client.sin_port)); + } + connected_clients.remove(client_fd); + close(client_fd); + FD_CLR(client_fd, &readfds); + FD_CLR(client_fd, &exceptfds); + } else { + TRY(Core::System::write(STDOUT_FILENO, buffer_span)); + } + } + } + } } return 0;