diff --git a/src/wstationd.c b/src/wstationd.c index 7f65113..e5989bf 100644 --- a/src/wstationd.c +++ b/src/wstationd.c @@ -14,6 +14,8 @@ #include #include "wstationd.h" +#define BUFFERSIZE 4096 + static char *exec_name; @@ -42,10 +44,12 @@ int main(int argc, char *argv[]) int port = 10800; // Default port uint32_t address = INADDR_ANY; // Default listen address - int sockflag = 1, sockfd = 0, connfd = 0; + int connflag, sockflag = 1; + int sockfd = 0, connfd = 0; struct sockaddr_in sock, client; socklen_t client_sz = sizeof(client); - char timebuff[15], fnbuff[35]; + char connbuff[BUFFERSIZE], timebuff[15], fnbuff[35]; + ssize_t connbuff_sz; time_t timeraw; struct tm* timeptr; @@ -155,6 +159,25 @@ int main(int argc, char *argv[]) fprintf(stderr, "%s: accept failed - %s\n", exec_name, strerror(errno) ); + continue; + } + + // Get current flags before setting nonblock + if((connflag = fcntl(connfd, F_GETFL)) == -1){ + fprintf(stderr, "%s: cannot get connection flags - %s\n", + exec_name, strerror(errno) + ); + close(connfd); + continue; + } + + // Set nonblocking on socket + if(fcntl(connfd, F_SETFL, connflag | O_NONBLOCK) == -1){ + fprintf(stderr, "%s: cannot set connection flags - %s\n", + exec_name, strerror(errno) + ); + close(connfd); + continue; } // Collect the time of the connection, process it into @@ -170,8 +193,9 @@ int main(int argc, char *argv[]) timebuff, inet_ntoa(client.sin_addr) ); - if((writefds[i] = open(fnbuff, O_WRONLY | O_CREAT, 0644)) == -1){ - writefds[i] = 0; + // Open a file to dump our data to + if((writefds[connfd] = open(fnbuff, O_WRONLY | O_CREAT, 0644)) == -1){ + writefds[connfd] = 0; fprintf(stderr, "%s: cannot open - %s\n", exec_name, strerror(errno) ); @@ -180,16 +204,36 @@ int main(int argc, char *argv[]) } else { FD_SET(connfd, &masterfds); } - connfd = 0; // Read from existing connections } else { - if(writefds[i] != 0){ + // Do the read + connbuff_sz = recv(i, connbuff, BUFFERSIZE, 0); + if(connbuff_sz > 0){ + // If the read fails, throw an error, and set the + // buffer size to 0, which forces a disconnect. + if(write(writefds[i], connbuff, connbuff_sz) == -1){ + fprintf(stderr, "%s: file write error - %s\n", + exec_name, strerror(errno) + ); + connbuff_sz = 0; + } + } + + // If the buffer is empty (which with select(), will only happen + // on disconnect, disconnect and cleanup. + if(connbuff_sz == 0){ close(writefds[i]); writefds[i] = 0; + FD_CLR(i, &masterfds); + close(i); + } else if(connbuff_sz == -1){ + fprintf(stderr, "%s: recv error - %s\n", + exec_name, strerror(errno) + ); + FD_CLR(i, &masterfds); + close(i); } - FD_CLR(i, &masterfds); - close(i); } } }