From 3c75202f5f4e9c1c61fead33f0f08e742d143223 Mon Sep 17 00:00:00 2001
From: cbdev <cb@cbcdn.com>
Date: Thu, 20 Jun 2019 12:44:23 +0200
Subject: Implement keep-alive pings (Fixes #4)

---
 README.md   |  2 ++
 config.c    |  6 ++++++
 config.h    |  1 +
 websocket.c |  6 +++---
 websocket.h |  2 +-
 websocksy.c | 40 ++++++++++++++++++++++++++++++++--------
 websocksy.h |  1 +
 7 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index b1decc9..369f44d 100644
--- a/README.md
+++ b/README.md
@@ -104,6 +104,7 @@ at the time.
 
 * `-p <port>`: Set the listen port for incoming WebSocket connections (Default: `8001`)
 * `-l <host>`: Set the host for listening for incoming WebSocket connections (Default: `::`)
+* `-k <seconds>`: Set the inactivity timeout for sending WebSocket keep-alive pings
 * `-b <backend>`: Select external backend
 * `-c <option>=<value>`: Pass configuration option to backend
 
@@ -117,6 +118,7 @@ In the `[core]` section, the following options are recognized:
 
 * `port`: Listen port for incoming WebSocket connections
 * `listen`: Host for incoming WebSocket connections
+* `ping`: Inactivity timeout for WebSocket keep-alive pings in seconds
 * `backend`: External backend selection
 
 In the `[backend]` section, all options are passed directly to the backend and thus are dependent on the specific
diff --git a/config.c b/config.c
index 2d50f3e..6998cc7 100644
--- a/config.c
+++ b/config.c
@@ -22,6 +22,9 @@ static int config_file_line(ws_config* config, char* key, char* value, size_t li
 		free(config->host);
 		config->host = strdup(value);
 	}
+	else if(!strcmp(key, "ping")){
+		config->ping_interval = strtoul(value, NULL, 10);
+	}
 	else if(!strcmp(key, "backend")){
 		//clean up the previously registered backend
 		if(config->backend.cleanup){
@@ -136,6 +139,9 @@ int config_parse_arguments(ws_config* config, int argc, char** argv){
 			case 'b':
 				config_file_line(config, "backend", argv[u + 1], 0);
 				break;
+			case 'k':
+				config_file_line(config, "ping", argv[u + 1], 0);
+				break;
 			case 'c':
 				if(!strchr(argv[u + 1], '=')){
 					return 1;
diff --git a/config.h b/config.h
index e39f13e..b151486 100644
--- a/config.h
+++ b/config.h
@@ -2,6 +2,7 @@
 typedef struct /*_websocksy_config*/ {
 	char* host;
 	char* port;
+	time_t ping_interval;
 	ws_backend backend;
 } ws_config;
 
diff --git a/websocket.c b/websocket.c
index 710178f..1fff538 100644
--- a/websocket.c
+++ b/websocket.c
@@ -95,10 +95,11 @@ int ws_close(websocket* ws, ws_close_reason code, char* reason){
 }
 
 /* Accept a new WebSocket connection */
-int ws_accept(int listen_fd){
+int ws_accept(int listen_fd, time_t current_time){
 	websocket ws = {
 		.ws_fd = accept(listen_fd, NULL, NULL),
-		.peer_fd = -1
+		.peer_fd = -1,
+		.last_event = current_time
 	};
 
 	return client_register(&ws);
@@ -417,7 +418,6 @@ int ws_send_frame(websocket* ws, ws_operation opcode, uint8_t* data, size_t len)
 	uint16_t* payload_len16 = (uint16_t*) (frame_header + 2);
 	uint64_t* payload_len64 = (uint64_t*) (frame_header + 2);
 
-	//FIXME might want to support segmented transfers here
 	//set up the basic frame header
 	frame_header[0] = WS_FLAG_FIN | opcode;
 	if(len <= 125){
diff --git a/websocket.h b/websocket.h
index f57fccc..011e63e 100644
--- a/websocket.h
+++ b/websocket.h
@@ -2,6 +2,6 @@
 
 /* WebSocket connection handling functions */
 int ws_close(websocket* ws, ws_close_reason code, char* reason);
-int ws_accept(int listen_fd);
+int ws_accept(int listen_fd, time_t current_time);
 int ws_send_frame(websocket* ws, ws_operation opcode, uint8_t* data, size_t len);
 int ws_data(websocket* ws);
diff --git a/websocksy.c b/websocksy.c
index c2f0631..8b37cec 100644
--- a/websocksy.c
+++ b/websocksy.c
@@ -9,6 +9,7 @@
 #include <unistd.h>
 #include <fcntl.h>
 #include <ctype.h>
+#include <time.h>
 
 #include "websocksy.h"
 #include "builtins.h"
@@ -22,7 +23,6 @@
 
 /* TODO
  * - TLS
- * - pings
  * - continuation
  */
 
@@ -46,6 +46,7 @@ char* xstr_lower(char* in){
 static ws_config config = {
 	.host = NULL,
 	.port = NULL,
+	.ping_interval = 30,
 	/* Assign the built-in defaultpeer backend by default */
 	.backend.init = backend_defaultpeer_init,
 	.backend.config = backend_defaultpeer_configure,
@@ -205,8 +206,9 @@ static int usage(char* fn){
 	fprintf(stderr, "\t%s <configuration file>\n", fn);
 	fprintf(stderr, "\t%s [-p <port>] [-l <listen address>] [-b <discovery backend>] [-c <option>=<value>]\n", fn);
 	fprintf(stderr, "Arguments:\n");
-	fprintf(stderr, "\t-p <port>\t\tWebSocket listen port (Current: %s, Default: %s)\n", config.port, DEFAULT_PORT);
-	fprintf(stderr, "\t-l <address>\t\tWebSocket listen address (Current: %s, Default: %s)\n", config.host, DEFAULT_HOST);
+	fprintf(stderr, "\t-p <port>\t\tWebSocket listen port (Current: %s, Default: %s)\n", config.port ? config.port : DEFAULT_PORT, DEFAULT_PORT);
+	fprintf(stderr, "\t-l <address>\t\tWebSocket listen address (Current: %s, Default: %s)\n", config.host ? config.host : DEFAULT_HOST, DEFAULT_HOST);
+	fprintf(stderr, "\t-k <seconds>\t\tKeepalive ping interval (Current: %lu)\n", config.ping_interval);
 	fprintf(stderr, "\t-b <backend>\t\tPeer discovery backend (Default: built-in 'defaultpeer')\n");
 	fprintf(stderr, "\t-c <option>=<value>\tPass configuration options to the peer discovery backend\n");
 	return EXIT_FAILURE;
@@ -273,6 +275,10 @@ int main(int argc, char** argv){
 	fd_set read_fds;
 	size_t n;
 	int listen_fd = -1, status, max_fd;
+	struct timespec current_time;
+	struct timeval select_timeout = {
+		0
+	};
 
 	//register default framing functions before parsing arguments, as they may be assigned within a backend configuration
 	if(plugin_register_framing("auto", framing_auto)
@@ -313,8 +319,14 @@ int main(int argc, char** argv){
 
 	//core loop
 	while(!shutdown_requested){
+		//reset the select timeout
+		select_timeout.tv_sec = config.ping_interval / 2;
+		select_timeout.tv_usec = 0;
+
+		//clear the select set
 		FD_ZERO(&read_fds);
 
+		//push the websocket fd
 		FD_SET(listen_fd, &read_fds);
 		max_fd = listen_fd;
 
@@ -336,18 +348,21 @@ int main(int argc, char** argv){
 		}
 
 		//block until something happens
-		status = select(max_fd + 1, &read_fds, NULL, NULL, NULL);
+		status = select(max_fd + 1, &read_fds, NULL, NULL, config.ping_interval ? &select_timeout : NULL);
 		if(status < 0){
 			fprintf(stderr, "Failed to select: %s\n", strerror(errno));
 			break;
 		}
-		else if(status == 0){
-			//timeout in select - ignore for now
-		}
 		else{
+			//update current timestamp
+			if(clock_gettime(CLOCK_MONOTONIC_COARSE, &current_time) < 0){
+				fprintf(stderr, "Failed to update current timestamp\n");
+				break;
+			}
+
 			//new websocket client
 			if(FD_ISSET(listen_fd, &read_fds)){
-				if(ws_accept(listen_fd)){
+				if(ws_accept(listen_fd, current_time.tv_sec)){
 					break;
 				}
 			}
@@ -356,6 +371,7 @@ int main(int argc, char** argv){
 			for(n = 0; n < socks; n++){
 				if(sock[n].ws_fd >= 0){
 					if(FD_ISSET(sock[n].ws_fd, &read_fds)){
+						sock[n].last_event = current_time.tv_sec;
 						if(ws_data(sock + n)){
 							break;
 						}
@@ -366,6 +382,14 @@ int main(int argc, char** argv){
 							break;
 						}
 					}
+
+					//send keep-alive frames
+					if(config.ping_interval &&
+							current_time.tv_sec - sock[n].last_event > config.ping_interval){
+						if(ws_send_frame(sock + n, ws_frame_ping, (uint8_t*) "PING", 4)){
+							ws_close(sock + n, ws_close_unexpected, NULL);
+						}
+					}
 				}
 			}
 		}
diff --git a/websocksy.h b/websocksy.h
index 9d9faf1..6524977 100644
--- a/websocksy.h
+++ b/websocksy.h
@@ -121,6 +121,7 @@ typedef struct /*_web_socket*/ {
 	uint8_t read_buffer[WS_MAX_LINE];
 	size_t read_buffer_offset;
 	ws_state state;
+	time_t last_event;
 
 	/* HTTP request headers */
 	size_t headers;
-- 
cgit v1.2.3