TCP echo client and server

tcpev.c

/*
 * -------------------------------  tcpev.c  ---------------------------------
 *
 * Copyright (c) 2015, Yorick de Wid <yorick17 at outlook dot com>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   * Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *   * Neither the name of Redis nor the names of its contributors may be used
 *     to endorse or promote products derived from this software without
 *     specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * Simple tree structure. It is a btree but not a binary search tree.
 *
 * Compile as:
 *     cc -std=gnu99 -Wall tcpev.c -o tcpev
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <ev.h>

#define PORT_NO 3033
#define BUFFER_SIZE 1024

static int total_clients = 0;  // Total number of connected clients

EV_P;
char *line = NULL;
size_t len = 0;

struct sock_ev_client {
	ev_io io;
	int fd;
	int index;
};

static struct sock_ev_client *conn_client = NULL; // In client mode there is one persistent connection

int setnonblock(int fd) {
	int flags;

	flags = fcntl(fd, F_GETFL);
	flags |= O_NONBLOCK;
	return fcntl(fd, F_SETFL, flags);
}

/* Read client message */
void read_cb(EV_P_ struct ev_io *watcher, int revents){
	char buffer[BUFFER_SIZE];
	ssize_t read;

	if (EV_ERROR & revents) {
		perror("got invalid event");
		return;
	}

	struct sock_ev_client *client = (struct sock_ev_client *)watcher;

	// Receive message from client socket
	read = recv(client->fd, buffer, BUFFER_SIZE, 0);
	if (read < 0) {
		if (EAGAIN == errno) {
			puts("Undefined state");
		} else {
			perror("read error");
		}
		return;
	}

	if (read == 0) {
		// Stop and free watchet if client socket is closing
		puts("Client disconnected");
		ev_io_stop(EV_A_ &client->io);
		close(client->fd);
		free(client);
		total_clients--; // Decrement total_clients count
		printf("%d client(s) connected.\n", total_clients);
	} else {
		printf("Client: %.*s\n", (int)read, buffer);
	}

	// Send message bach to the client
	send(client->fd, buffer, read, 0);
	memset(buffer, 0, read);
}

/* Accept client requests */
void accept_cb(EV_P_ struct ev_io *watcher, int revents) {
	struct sockaddr_in client_addr;
	socklen_t client_len = sizeof(client_addr);
	int sd;
	struct sock_ev_client *client = (struct sock_ev_client *)malloc(sizeof(struct sock_ev_client));
	
	if (EV_ERROR & revents) {
		perror("got invalid event");
		return;
	}

	// Accept client request
	sd = accept(watcher->fd, (struct sockaddr *)&client_addr, &client_len);
	if (sd < 0) {
		if (errno != EAGAIN && errno != EWOULDBLOCK) {
			perror("accept error");
			return;
		}
	}

	// Set it non-blocking
	if (setnonblock(sd)<0) {
		perror("echo server socket nonblock");
		return;
	}

	client->fd = sd;
	client->index = ++total_clients; // Increment total_clients count
	printf("Successfully connected with client.\n");
	printf("%d client(s) connected.\n", total_clients);

	// Initialize and start watcher to read client requests
	ev_io_init(&client->io, read_cb, sd, EV_READ);
	ev_io_start(EV_A_ &client->io);
}

void send_cb(EV_P_ ev_io *watcher, int revents) {
	struct sock_ev_client *client = (struct sock_ev_client *)watcher;
	if (revents & EV_WRITE) {
		if (send(client->fd, line, len, 0)<0) {
			perror("echo send");
			return;
		}

		// Now wait on socket response
		ev_io_stop(EV_A_ &client->io);
		ev_io_set(&client->io, client->fd, EV_READ);
		ev_io_start(EV_A_ &client->io);
	} else if (revents & EV_READ) {
		ssize_t read;
		char buffer[BUFFER_SIZE];

		read = recv(client->fd, buffer, BUFFER_SIZE, 0);
		if (read < 0) {
			if (EAGAIN == errno) {
				puts("Undefined state");
			} else {
				perror("read error");
			}
			return;
		}

		if (read == 0) {
			// Stop and free watchet if client socket is closing
			puts("Orderly disconnect");
			ev_io_stop(EV_A_ &client->io);
			close(client->fd);
			free(client);
			client = NULL;
			return;
		} else {
			printf("Server: %.*s\n", (int)read, buffer);
		}
		memset(buffer, 0, read);
	} else if (EV_ERROR & revents) {
		perror("got invalid event");
		return;
	}
}

void stdin_cb(EV_P_ ev_io *watcher, int revents) {
	getline(&line, &len, stdin);
	// Allow to write to the socket
	ev_io_stop(EV_A_ &conn_client->io);
	ev_io_set(&conn_client->io, conn_client->fd, EV_READ | EV_WRITE);
	ev_io_start(EV_A_ &conn_client->io);
}

int client_init() {
	int sd;

	// Create server socket
	if ((sd = fileno(stdin))<0){
		perror("socket error");
		return -1;
	}

	// Set it non-blocking
	if (setnonblock(sd)<0) {
		perror("echo server socket nonblock");
		return -1;
	}

	return sd;
}

int client_connect(EV_P_ char *remote_addr) {
	int sd;
 	struct sockaddr_in remote;
 	conn_client = (struct sock_ev_client *)malloc(sizeof(struct sock_ev_client));

 	// Create client socket
	if ((sd = socket(AF_INET, SOCK_STREAM, 0))<0){
		perror("socket error");
		return -1;
	}

	// Set it non-blocking
	if (setnonblock(sd)<0) {
		perror("echo server socket nonblock");
		return -1;
	}

	conn_client->fd = sd;
	conn_client->index = 0;

	// initialize the send callback, but wait to start until there is data to write
	ev_io_init(&conn_client->io, send_cb, sd, EV_READ);
	ev_io_start(EV_A_ &conn_client->io);

	memset(&remote, 0, sizeof(remote));
	remote.sin_family = AF_INET;
	remote.sin_port = htons(PORT_NO);
	remote.sin_addr.s_addr = inet_addr(remote_addr);

	int res = connect(sd, (struct sockaddr *)&remote, sizeof(remote));
	if (res < 0) {
		if (errno != EINPROGRESS) {
			perror("connect error");
			return -1;
		}
	}
	puts("Connected");

  	return sd;
}

int server_init(int max_queue) {
	int sd;
	struct sockaddr_in addr;

	// Create server socket
	if ((sd = socket(AF_INET, SOCK_STREAM, 0))<0){
		perror("socket error");
		return -1;
	}

	// Set it non-blocking
	if (setnonblock(sd)<0) {
		perror("echo server socket nonblock");
		return -1;
	}

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = htons(PORT_NO);
	addr.sin_addr.s_addr = INADDR_ANY;

	// Bind socket to address
	if (bind(sd, (struct sockaddr*)&addr, sizeof(addr))<0) {
		perror("bind error");
		return -1;
	}

	// Start listing on the socket
	if (listen(sd, max_queue)<0) {
		perror("listen error");
		return -1;
	}

	return sd;
}

int main(int argc, char *argv[]) {
	int sd;
	int server = 1;
	int max_queue = 128;
	loop = EV_DEFAULT;
	struct ev_io w_accept, w_stdin;
	char remote_addr[INET_ADDRSTRLEN];

	if (argc > 1) {
		server = 0;
		strncpy(remote_addr, argv[1], INET_ADDRSTRLEN);
	}

	// client mode
	if (!server) {
		client_connect(EV_A_ remote_addr);

		// Initialize client
		sd = client_init();

		// Initialize and start a watcher on the input stream
		ev_io_init(&w_stdin, stdin_cb, sd, EV_READ);
		ev_io_start(EV_A_ &w_stdin);
	} else {
		// Initialize daemon
		sd = server_init(max_queue);

		// Initialize and start a watcher to accepts client requests
		ev_io_init(&w_accept, accept_cb, sd, EV_READ);
		ev_io_start(EV_A_ &w_accept);
	}

	// Start infinite loop
	puts("Starting events");
	ev_loop(EV_A_ 0);

	ev_loop_destroy(loop);

	return 0;
}