diff --git a/src/main/pumpnet/prinet-proxy/client.c b/src/main/pumpnet/prinet-proxy/client.c index b7e32ae..62dab25 100644 --- a/src/main/pumpnet/prinet-proxy/client.c +++ b/src/main/pumpnet/prinet-proxy/client.c @@ -1,19 +1,34 @@ +#include + #include "pumpnet/prinet-proxy/client.h" #include "util/log.h" #include "util/mem.h" #include "util/sock-tcp.h" -static bool _pumpnet_prinet_proxy_client_recv_packet_length(int handle, uint32_t* length) +struct pumpnet_prinet_proxy_client_connection { + int handle; + pthread_mutex_t mutex; +}; + +static const uint32_t _pumpnet_prinet_proxy_client_recv_length_timeout_ms = 200; + +static bool _pumpnet_prinet_proxy_client_recv_packet_length(struct pumpnet_prinet_proxy_client_connection* connection, uint32_t* length) { uint32_t packet_length; while (true) { - ssize_t read = util_sock_tcp_recv_block(handle, &packet_length, sizeof(uint32_t)); + ssize_t read; + + pthread_mutex_lock(&connection->mutex); + + read = util_sock_tcp_recv(connection->handle, &packet_length, sizeof(uint32_t), _pumpnet_prinet_proxy_client_recv_length_timeout_ms); + + pthread_mutex_unlock(&connection->mutex); if (read == 0) { - log_error("Unexpected remote close and no data"); - return NULL; + // no data, repeat + continue; } if (read != sizeof(uint32_t)) { @@ -28,16 +43,107 @@ static bool _pumpnet_prinet_proxy_client_recv_packet_length(int handle, uint32_t return true; } -struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(int handle) +struct pumpnet_prinet_proxy_client_connection* pumpnet_prinet_proxy_client_connection_alloc() { + struct pumpnet_prinet_proxy_client_connection* connection; + + connection = util_xmalloc(sizeof(struct pumpnet_prinet_proxy_client_connection)); + + connection->handle = INVALID_SOCK_HANDLE; + pthread_mutex_init(&connection->mutex, NULL); + + return connection; +} + +bool pumpnet_prinet_proxy_client_connection_accept(int socket, struct pumpnet_prinet_proxy_client_connection* connection) +{ + int handle; + + log_debug("Waiting for incoming connection..."); + + handle = util_sock_tcp_wait_and_accept_remote_connection(socket); + + log_debug("Received connection: %X", socket); + + if (socket == INVALID_SOCK_HANDLE) { + log_error("Waiting and accepting source connection failed"); + return NULL; + } + + pthread_mutex_lock(&connection->mutex); + + connection->handle = handle; + + log_debug("Received connection: %X", handle); + + pthread_mutex_unlock(&connection->mutex); + + return connection; +} + +bool pumpnet_prinet_proxy_client_connection_is_active(struct pumpnet_prinet_proxy_client_connection* connection) +{ + log_assert(connection); + + bool is_active; + + pthread_mutex_lock(&connection->mutex); + + is_active = connection->handle != INVALID_SOCK_HANDLE; + + pthread_mutex_unlock(&connection->mutex); + + return is_active; +} + +void pumpnet_prinet_proxy_client_connection_close(struct pumpnet_prinet_proxy_client_connection* connection) +{ + log_assert(connection); + + pthread_mutex_lock(&connection->mutex); + + if (connection->handle != INVALID_SOCK_HANDLE) { + util_sock_tcp_close(connection->handle); + + connection->handle = INVALID_SOCK_HANDLE; + } + + pthread_mutex_unlock(&connection->mutex); +} + +void pumpnet_prinet_proxy_client_connection_free(struct pumpnet_prinet_proxy_client_connection** connection) +{ + log_assert(connection); + + struct pumpnet_prinet_proxy_client_connection* con = *connection; + + pthread_mutex_lock(&con->mutex); + + if (con->handle != INVALID_SOCK_HANDLE) { + log_error("Freeing connection %X which is not closed", con->handle); + } + + pthread_mutex_unlock(&con->mutex); + + pthread_mutex_destroy(&con->mutex); + util_xfree((void**) connection); +} + +struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(struct pumpnet_prinet_proxy_client_connection* connection) +{ + log_assert(connection); + uint32_t packet_length; - if (!_pumpnet_prinet_proxy_client_recv_packet_length(handle, &packet_length)) { + // locking is handled inside this function + if (!_pumpnet_prinet_proxy_client_recv_packet_length(connection, &packet_length)) { log_error("Receiving length field for request from source failed"); return NULL; } - log_debug("Received length source request (%X): %d", handle, packet_length); + pthread_mutex_lock(&connection->mutex); + + log_debug("Received length source request (%X): %d", connection->handle, packet_length); struct pumpnet_prinet_proxy_packet* packet = util_xmalloc(packet_length); packet->length = packet_length; @@ -50,18 +156,20 @@ struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(int size_t remaining_size = packet->length - sizeof(uint32_t) - read_pos; ssize_t read = util_sock_tcp_recv_block( - handle, + connection->handle, ((uint8_t*) packet) + sizeof(uint32_t) + read_pos, remaining_size); if (read == 0) { log_error("Unexpected remote close and no data"); - return NULL; + util_xfree((void**) &packet); + break; } if (read == -1) { log_error("Receiving length field for request from source failed: %d", read); - return NULL; + util_xfree((void**) &packet); + break; } read_pos += read; @@ -71,23 +179,39 @@ struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(int } } + pthread_mutex_unlock(&connection->mutex); + return packet; } -bool pumpnet_prinet_proxy_client_send_response(int handle, const struct pumpnet_prinet_proxy_packet* packet) +bool pumpnet_prinet_proxy_client_send_response(struct pumpnet_prinet_proxy_client_connection* connection, const struct pumpnet_prinet_proxy_packet* packet) { - log_debug("Sending source response (%X): %d", handle, packet->length); + log_assert(connection); + log_assert(packet); - if (util_sock_tcp_send_block(handle, packet, packet->length) != packet->length) { + bool result; + + result = true; + + pthread_mutex_lock(&connection->mutex); + + log_debug("Sending source response (%X): %d", connection->handle, packet->length); + + if (util_sock_tcp_send_block(connection->handle, packet, packet->length) != packet->length) { log_error("Sending response, len %d, to source failed", packet->length); - return false; - } else { - return true; + result = false; } + + pthread_mutex_unlock(&connection->mutex); + + return result; } bool pumpnet_prinet_proxy_client_unpack_data_request(const struct pumpnet_prinet_proxy_packet* packet, struct util_iobuf* dec_data) { + log_assert(packet); + log_assert(dec_data); + size_t enc_data_len = pumpnet_prinet_proxy_packet_get_data_len(packet); size_t dec_data_len = sec_prinet_decrypt( @@ -109,6 +233,10 @@ bool pumpnet_prinet_proxy_client_unpack_data_request(const struct pumpnet_prinet bool pumpnet_prinet_proxy_client_pack_data_response(const struct pumpnet_prinet_proxy_packet* req_packet, const struct util_iobuf* dec_data, struct pumpnet_prinet_proxy_packet* resp_packet) { + log_assert(req_packet); + log_assert(dec_data); + log_assert(resp_packet); + size_t enc_data_len = sec_prinet_encrypt( req_packet->nounce, sizeof(req_packet->nounce), @@ -126,4 +254,4 @@ bool pumpnet_prinet_proxy_client_pack_data_response(const struct pumpnet_prinet_ resp_packet->length = pumpnet_prinet_proxy_packet_get_header_len() + enc_data_len; return true; -} \ No newline at end of file +} diff --git a/src/main/pumpnet/prinet-proxy/client.h b/src/main/pumpnet/prinet-proxy/client.h index 50b9a3f..cc84a79 100644 --- a/src/main/pumpnet/prinet-proxy/client.h +++ b/src/main/pumpnet/prinet-proxy/client.h @@ -7,9 +7,21 @@ #include "util/iobuf.h" -struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(int handle); +struct pumpnet_prinet_proxy_client_connection; -bool pumpnet_prinet_proxy_client_send_response(int handle, const struct pumpnet_prinet_proxy_packet* packet); +struct pumpnet_prinet_proxy_client_connection* pumpnet_prinet_proxy_client_connection_alloc(); + +bool pumpnet_prinet_proxy_client_connection_accept(int socket, struct pumpnet_prinet_proxy_client_connection* connection); + +bool pumpnet_prinet_proxy_client_connection_is_active(struct pumpnet_prinet_proxy_client_connection* connection); + +void pumpnet_prinet_proxy_client_connection_close(struct pumpnet_prinet_proxy_client_connection* connection); + +void pumpnet_prinet_proxy_client_connection_free(struct pumpnet_prinet_proxy_client_connection** connection); + +struct pumpnet_prinet_proxy_packet* pumpnet_prinet_proxy_client_recv_request(struct pumpnet_prinet_proxy_client_connection* connection); + +bool pumpnet_prinet_proxy_client_send_response(struct pumpnet_prinet_proxy_client_connection* connection, const struct pumpnet_prinet_proxy_packet* packet); bool pumpnet_prinet_proxy_client_unpack_data_request(const struct pumpnet_prinet_proxy_packet* packet, struct util_iobuf* dec_data); diff --git a/src/main/pumpnet/prinet-proxy/main.c b/src/main/pumpnet/prinet-proxy/main.c index 76f1656..20e33a8 100644 --- a/src/main/pumpnet/prinet-proxy/main.c +++ b/src/main/pumpnet/prinet-proxy/main.c @@ -1,4 +1,3 @@ -#include #include #include @@ -20,7 +19,12 @@ static const uint32_t NUM_CONNECTIONS = 1; static const size_t PUMPNET_MAX_RESP_SIZE = 1024 * 1024; -static atomic_int _source_sock; +// Client closes their sending socket after 60 sec of inactivity...jeez, why +// can't you implement a damn session layer correctly?! +static const uint32_t KEEPALIVE_POLL_MS = 30000; + +static int _source_sock; +struct pumpnet_prinet_proxy_client_connection* _source_con; /* Compiled binary data from data folder. Symbols are defined by compiler */ extern const uint8_t _binary_prime_private_key_start[]; @@ -32,6 +36,11 @@ static void _sigint_handler(int sig) { log_info("SIGINT, exiting"); + pumpnet_prinet_proxy_keepalive_shutdown(); + + pumpnet_prinet_proxy_client_connection_close(_source_con); + pumpnet_prinet_proxy_client_connection_free(&_source_con); + util_sock_tcp_close(_source_sock); sec_prinet_finit(); @@ -84,6 +93,11 @@ int main(int argc, char** argv) return -3; } + _source_con = pumpnet_prinet_proxy_client_connection_alloc(); + + // See comment of KEEPALIVE_POLL_MS + pumpnet_prinet_proxy_keepalive_init(_source_con, KEEPALIVE_POLL_MS); + log_info("Running proxy loop"); signal(SIGINT, _sigint_handler); @@ -101,15 +115,9 @@ int main(int argc, char** argv) // active at the same time~ while (true) { - int source_con = INVALID_SOCK_HANDLE; - log_debug("Waiting for incoming connection..."); - source_con = util_sock_tcp_wait_and_accept_remote_connection(_source_sock); - - log_debug("Received connection: %X", source_con); - - if (source_con == INVALID_SOCK_HANDLE) { + if (!pumpnet_prinet_proxy_client_connection_accept(_source_sock, _source_con)) { log_error("Waiting and accepting source connection failed"); continue; } @@ -124,7 +132,7 @@ int main(int argc, char** argv) struct util_iobuf pumpnet_data_resp; struct pumpnet_prinet_proxy_packet* source_resp = NULL; - source_req = pumpnet_prinet_proxy_client_recv_request(source_con); + source_req = pumpnet_prinet_proxy_client_recv_request(_source_con); if (!source_req) { log_error("Receiving request from source failed"); @@ -173,7 +181,7 @@ int main(int argc, char** argv) goto cleanup_iteration; } - if (!pumpnet_prinet_proxy_client_send_response(source_con, source_resp)) { + if (!pumpnet_prinet_proxy_client_send_response(_source_con, source_resp)) { log_error("Sending response to source failed"); inner_loop = false; goto cleanup_iteration; @@ -197,6 +205,6 @@ int main(int argc, char** argv) } } while (inner_loop); - util_sock_tcp_close(source_con); + pumpnet_prinet_proxy_client_connection_close(_source_con); } } \ No newline at end of file