uni/inf1060/exam/src/client.c

#include "job.h"
#include "tcp.h"
#include "argparser.h"

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <netdb.h>
#include <signal.h>

#ifndef NOREADLINE
#include <readline/readline.h>
#include <readline/history.h>
#endif

#define BUF_SIZE 5

static int sockfd = -1;
pid_t pid_a;
pid_t pid_b;
static int pipe_a[2];
static int pipe_b[2];

// input_str will be a pointer to the string containing user input,
// or NULL if there's no user input to be freed.
static char *input_str = NULL;

/*
 * Terminate the program, closing sockets and sending close messages.
 *
 * Input:
 *     code: The exit code.
 *     err: The string to send to perror.
 *          If NULL, perror won't be called.
 */
void term(int code, char *err)
{
	if (err != NULL)
	{
		perror(err);
	}

	if (sockfd != -1)
	{
		char buf[BUF_SIZE];
		memset(buf, 0, sizeof(buf));
		if (code == 0)
			buf[0] = 'T';
		else
			buf[0] = 'E';
		write(sockfd, buf, sizeof(buf));

		close(sockfd);
	}

	free(input_str);
	kill(pid_a, SIGTERM);
	kill(pid_b, SIGTERM);

	close(pipe_a[0]);
	close(pipe_a[1]);
	close(pipe_b[0]);
	close(pipe_b[1]);

#ifndef NOREADLINE
	clear_history();
#endif

	exit(code);
}

/*
 * Like term, but with herror instead of perror.
 */
void hterm(int code, char *err)
{
	if (err != NULL)
	{
		herror(err);
	}

	term(code, NULL);
}

/*
 * Help text
 */
void cmd_help()
{
	printf("\thelp:\n\t\tShow this help text.\n\n");
	printf("\tget [number | all]:\n\t\tGet [number] jobs, defaults to 1.\n\n");
	printf("\tq:\n\t\nExit.\n");
}

/*
 * Get jobs and send them to a child process.
 *
 * Input:
 *     count:
 *         if NULL: get 1 job
 *         if number: get <number> jobs
 *         if "all": get all jobs
 */
void cmd_get(char *count)
{
	uint32_t c = 1;
	if (count != NULL && strcmp(count, "all") == 0)
		c = ~((uint32_t)0);
	else if (count != NULL)
		c = atoi(count);

	char buf[BUF_SIZE];
	memset(buf, 0, 5);
	buf[0] = 'G';
	buf[1] = c;
	buf[2] = c >> 8;
	buf[3] = c >> 16;
	buf[4] = c >> 24;

	if (write(sockfd, buf, sizeof(buf)) < 0)
		term(1, "Couldn't send data");

	char res[sizeof(job)];
	for (int i = 0; i < c; ++i)
	{
		memset(res, 0, sizeof(res));
		if (read(sockfd, res, sizeof(res)) < 0)
			term(1, "Couldn't read data");

		int written = 0;

		// Send to child A if O or Q
		if (res[0] == 'O' || res[0] == 'Q')
		{
			written = 1;
			if (write(pipe_a[1], res, sizeof(job)) < 0)
				term(1, "Couldn't write to pipe");
		}

		// Send to child B if E or Q
		if (res[0] == 'E' || res[0] == 'Q')
		{
			written = 1;
			if (write(pipe_b[1], res, sizeof(job)) < 0)
				term(1, "Couldn't write to pipe");
		}

		// Quit if Q
		if (res[0] == 'Q')
		{
			usleep(50 * 1000);
			// Children should have exited by now
			term(0, NULL);
		}

		// Server died if 0
		if (res[0] == 0)
		{
			fprintf(stderr, "Server disconnected!\n");
			term(1, NULL);
		}

		// Error if we received an invalid type
		if (!written)
		{
			fprintf(stderr, "Received invalid job type: %c\n", res[0]);
			term(1, NULL);
		}
	}
}

/*
 * Exit gracefully.
 */
void cmd_exit()
{
	term(0, NULL);
}

/*
 * Evaluate an input string, do the desired action.
 *
 * Input:
 *     str: The input string.
 */
int eval(char *str)
{
	// Parse input, we won't need more than 2 arguments
	char *argv[2];
	argparser_parse(argv, 2, str);

	char *cmd = argv[0];
	if (argv[0] == NULL)
		return 0;

	if (strcmp(cmd, "help") == 0)
		cmd_help();
	else if (strcmp(cmd, "get") == 0)
		cmd_get(argv[1]);
	else if (strcmp(cmd, "q") == 0)
		cmd_exit(str);
	else
		printf("Unknown command: '%s'. Type 'help' for help.\n", cmd);

	return 1;
}

/*
 * SIGINT/SIGTERM handler
 */
void onterm(int signum)
{
	printf("\n");
	term(1, NULL);
}

// Code shared between both the children.
//
// Input:
//     fd: The file descriptor to read from.
//     out: the FILE* to write to.
void child_shared(int fd, FILE *out)
{
	char buf[sizeof(job)];
	job j;
	while (1)
	{
		if (read(fd, buf, sizeof(buf)) < 0)
			term(1, "Couldn't read from pipe");

		job_parse(buf, &j);
		if (j.type == 'Q')
		{
			exit(0);
		}
		else
		{
			fprintf(out, "%s\n", j.text);
			fflush(out);
		}
	}
}

void child_a()
{
	child_shared(pipe_a[0], stdout);
}

void child_b()
{
	child_shared(pipe_b[0], stderr);
}

int main(int argc, char *argv[])
{
	if (argc < 2)
	{
		printf("Usage: %s <host> <port>\n", argv[0]);
		return 1;
	}

	// Create child A
	if (pipe(pipe_a) < 0)
		term(1, "Couldn't create pipe");
	pid_a = fork();
	if (pid_a < 0)
		term(1, "Fork failed");
	else if (pid_a == 0)
		child_a();

	// Create child B
	if (pipe(pipe_b) < 0)
		term(1, "Couldn't create pipe");
	pid_b = fork();
	if (pid_b < 0)
		term(1, "Fork failed");
	else if (pid_b == 0)
		child_b();

	// SIGTERM an SIGINT handler
	struct sigaction action;
	memset(&action, 0, sizeof(action));
	action.sa_handler = onterm;
	sigaction(SIGTERM, &action, NULL);
	sigaction(SIGINT, &action, NULL);

	char *host = argv[1];
	int port = atoi(argv[2]);

	// Create socket
	sockfd = tcp_create_sock(host, port);
	if (sockfd == -1)
		term(1, "Couldn't create socket");
	else if (sockfd == -2)
		hterm(1, "Couldn't create socket");

	printf("Connected to server on %s:%i\n", host, port);
	printf("Type 'help' for help.\n");

	// Read user input from readline
	while (1)
	{
#ifndef NOREADLINE
		input_str = readline("> ");
		if (!input_str)
		{
			fprintf(stderr, "Failed to read user input.\n");
			term(1, NULL);
		}
		add_history(input_str);
#else
		input_str = NULL;
		size_t n = 0;
		if (getline(&input_str, &n, stdin) < 0)
		{
			free(input_str);
			fprintf(stderr, "Failed to read user input.\n");
			term(1, NULL);
		}
#endif
		eval(input_str);
		free(input_str);
		input_str = NULL;

		// There's a race condition; the child processes might print
		// their output after we print the "> " prompt, which isn't good.
		// This is an imperfect fix.
		usleep(100 * 1000);
	}

	return 0;
}