Added argument to set limit on number of concurrent connections

This commit is contained in:
Christopher Ramey 2015-03-09 01:39:25 +00:00 committed by cdramey
parent 5ec1783299
commit 5071a2fcd4

View File

@ -17,6 +17,7 @@
#include <getopt.h> #include <getopt.h>
#include <poll.h> #include <poll.h>
#include <syslog.h> #include <syslog.h>
#include <signal.h>
#include "wstationd.h" #include "wstationd.h"
@ -27,16 +28,18 @@
// dictates the memory cost per connection. // dictates the memory cost per connection.
#define BUFFERSIZE 65539 #define BUFFERSIZE 65539
// Hard limit on number of concurrent connections
#define SLOTLIMIT 128
// Timeout for each connection, in seconds // Timeout for each connection, in seconds
#define TIMEOUT 30 #define TIMEOUT 30
static char *exec_name; static char *exec_name;
// Should this application daemonize?
static int daemonize = 1; static int daemonize = 1;
// Hard limit on number of concurrent connections
static int slotlimit = 128;
// Used to handle signaled shutdowns
static int keep_polling = 1;
struct connection struct connection
{ {
@ -50,10 +53,16 @@ struct connection
static int write_connection(struct connection *); static int write_connection(struct connection *);
void shutdown_handler(int sig)
{
keep_polling = 0;
}
static void print_help() static void print_help()
{ {
fprintf(stdout, "usage: %s [-p port] [-b addr] directory\n", exec_name); fprintf(stdout, "usage: %s [-d] [-l limit] [-p port] [-b addr] directory\n", exec_name);
fprintf(stdout, " -d, --nodaemon Do not detach and daemonize\n"); fprintf(stdout, " -d, --nodaemon Do not detach and daemonize\n");
fprintf(stdout, " -l, --limit <limit> Set concurrent connection limit (default: 128)\n");
fprintf(stdout, " -p, --port <port> Port to listen on (default: 10800)\n"); fprintf(stdout, " -p, --port <port> Port to listen on (default: 10800)\n");
fprintf(stdout, " -b, --bind <address> Address to bind to (default: 0.0.0.0)\n"); fprintf(stdout, " -b, --bind <address> Address to bind to (default: 0.0.0.0)\n");
fprintf(stdout, " -h, --help Print this message and exit\n"); fprintf(stdout, " -h, --help Print this message and exit\n");
@ -64,6 +73,7 @@ static void print_help()
int main(int argc, char* argv[]) int main(int argc, char* argv[])
{ {
const struct option longopts[] = { const struct option longopts[] = {
{ "limit", required_argument, NULL, 'l' },
{ "nodaemon", required_argument, NULL, 'd' }, { "nodaemon", required_argument, NULL, 'd' },
{ "port", required_argument, NULL, 'p' }, { "port", required_argument, NULL, 'p' },
{ "bind", required_argument, NULL, 'b' }, { "bind", required_argument, NULL, 'b' },
@ -77,8 +87,8 @@ int main(int argc, char* argv[])
int port = 10800; // Default port int port = 10800; // Default port
uint32_t address = INADDR_ANY; // Default listen address uint32_t address = INADDR_ANY; // Default listen address
struct connection* conns[SLOTLIMIT]; struct connection** conns = NULL;
memset(conns, 0, sizeof(struct connection*) * SLOTLIMIT); struct pollfd* fds = NULL;
if((exec_name = basename(argv[0])) == NULL){ if((exec_name = basename(argv[0])) == NULL){
fprintf(stderr, "%s: cannot get basename - %s\n", fprintf(stderr, "%s: cannot get basename - %s\n",
@ -93,6 +103,15 @@ int main(int argc, char* argv[])
daemonize = 0; daemonize = 0;
break; break;
case 'l':
slotlimit = (int)strtol(optarg, NULL, 10);
if(slotlimit < 1 || slotlimit > 1024){
fprintf(stderr, "%s: invalid limit (must be between 1 and 1024)\n", exec_name);
print_help();
goto shutdown_error;
}
break;
case 'p': case 'p':
port = (int)strtol(optarg, NULL, 10); port = (int)strtol(optarg, NULL, 10);
if(port < 1 || port > 65535){ if(port < 1 || port > 65535){
@ -135,6 +154,9 @@ int main(int argc, char* argv[])
goto shutdown_error; goto shutdown_error;
} }
// Handle signals
signal(SIGINT, shutdown_handler);
// Daemonize if requested // Daemonize if requested
if(daemonize){ if(daemonize){
pid_t pid = fork(); pid_t pid = fork();
@ -228,16 +250,25 @@ int main(int argc, char* argv[])
goto shutdown_error; goto shutdown_error;
} }
struct pollfd fds[SLOTLIMIT+1]; fds = malloc(sizeof(struct pollfd) * (slotlimit + 1));
memset(&fds, 0, sizeof(struct pollfd) * (SLOTLIMIT + 1)); conns = malloc(sizeof(struct connection*) * slotlimit);
if(conns == NULL || fds == NULL){
if(daemonize){ syslog(LOG_ERR, "out of memory"); }
else { fprintf(stderr, "%s: out of memory\n", exec_name); }
}
memset(conns, 0, sizeof(struct connection*) * slotlimit);
memset(fds, 0, sizeof(struct pollfd) * (slotlimit + 1));
// Add sockfd // Add sockfd
fds[0].fd = sockfd; fds[0].fd = sockfd;
fds[0].events = POLLIN; fds[0].events = POLLIN;
fds[0].revents = 0; fds[0].revents = 0;
while(1){ while(keep_polling){
if(poll(fds, (SLOTLIMIT + 1), 1000) == -1){ if(poll(fds, (slotlimit + 1), 1000) == -1){
// Ignore interrupt generated errors
if(errno != EINTR){
if(daemonize){ if(daemonize){
syslog(LOG_ERR, "socket poll error - %s", syslog(LOG_ERR, "socket poll error - %s",
strerror(errno) strerror(errno)
@ -249,10 +280,11 @@ int main(int argc, char* argv[])
} }
goto shutdown_error; goto shutdown_error;
} }
}
// Handle existing connections // Handle existing connections
int slotfirst = 0, slotfree = 0; int slotfirst = 0, slotfree = 0;
for(int i = 1; i <= SLOTLIMIT; i++){ for(int i = 1; i <= slotlimit; i++){
// Handle event // Handle event
if(fds[i].revents & POLLIN){ if(fds[i].revents & POLLIN){
ssize_t sz = recv( ssize_t sz = recv(
@ -390,12 +422,24 @@ int main(int argc, char* argv[])
shutdown_clean: shutdown_clean:
if(sockfd > 0) close(sockfd); if(sockfd > 0) close(sockfd);
for(int i = 0; i < SLOTLIMIT; i++){ if(conns[i] != NULL) free(conns[i]); } if(fds != NULL){ free(fds); fds = NULL; }
if(conns != NULL){
for(int i = 0; i < slotlimit; i++){
if(conns[i] != NULL) free(conns[i]);
}
free(conns); conns = NULL;
}
return 0; return 0;
shutdown_error: shutdown_error:
if(sockfd > 0) close(sockfd); if(sockfd > 0) close(sockfd);
for(int i = 0; i < SLOTLIMIT; i++){ if(conns[i] != NULL) free(conns[i]); } if(fds != NULL){ free(fds); fds = NULL; }
if(conns != NULL){
for(int i = 0; i < slotlimit; i++){
if(conns[i] != NULL) free(conns[i]);
}
free(conns); conns = NULL;
}
return 1; return 1;
} }