c_cpp 控制平面：推送命令和获取结果

Posted 2021-05-14

/* For sockaddr_in */
#include <netinet/in.h>
/* For socket functions */
#include <sys/socket.h>
/* For fcntl */
#include <fcntl.h>
/* for select */
#include <sys/select.h>

#include <assert.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <errno.h>

#include "live-cmd.h"

#define MAX_LINE 16384

#ifdef DEBUG
#define COMMENT(x) printf("%s\n", x)
#else
#define COMMENT(x) {}
#endif

static struct fd_state *states[FD_SETSIZE];
static int log[LIVE_TYPE_MAX]; /* store activated fd */

struct fd_state {
    char buffer[MAX_LINE];
    size_t buffer_used;

	FILE *fp;
	int done;

    int writing;
    size_t n_written;
    size_t write_upto;
};

struct fd_state *
alloc_fd_state(void)
{
    struct fd_state *state = malloc(sizeof(struct fd_state));
    if (!state)
        return NULL;
    state->buffer_used = state->n_written = state->writing =
        state->write_upto = state->done = 0;
    return state;
}

void
free_fd_state(struct fd_state *state)
{
    free(state);
}

void
make_nonblocking(int fd)
{
    fcntl(fd, F_SETFL, O_NONBLOCK);
}

int
do_read(int fd, struct fd_state *state)
{
    int i;
    ssize_t result;

	if(state->fp)
	{
		int ret;
		
		COMMENT("Data plane");
		ret = read(fd, state->buffer+state->buffer_used, sizeof(state->buffer)-state->buffer_used);
		
		if(ret > 0)
		{
			state->buffer_used += ret;
		}
		else if(ret == 0)
		{
			state->done = 1;
#if 1
			result = sprintf(state->buffer+state->buffer_used, "\n\nFinished\n\0");
			state->buffer_used += result;
#endif
		}
	}
	else
	{
		LiveCmd cmd;
		
		COMMENT("Ctrl plane");
		result = recv(fd, &cmd, sizeof(cmd), MSG_WAITALL);

		if(result > 0)
		{
			if(cmd.op == LIVE_SET)
			{
				char cmd_string[1024]; /* FIXME */
				int act_fd = log[cmd.type];
				FILE *fp = NULL;
				
				*cmd_string = '\0';

				if(act_fd != -1)
				{
					COMMENT("First, stop existing process");
					COMMENT("Close data plane");
					if(states[act_fd]->fp)
						pclose(states[act_fd]->fp);
					free_fd_state(states[act_fd]);
					states[act_fd] = NULL;
				}

				if(cmd.type == LIVE_PING)
				{
					char count_opt[128];
					char interval_opt[128];
					char timeout_opt[128];

					*count_opt = *interval_opt = *timeout_opt = '\0';

					if(cmd.ping.count)
					{
						sprintf(count_opt, " -c %d ", cmd.ping.count);
					}
					if(cmd.ping.interval)
					{
#if LINUX
						sprintf(interval_opt, " -i %d ", cmd.ping.interval);
#endif
					}
					if(cmd.ping.timeout)
					{
#if LINUX
						sprintf(timeout_opt, " -t %d ", cmd.ping.timeout);
#else
						sprintf(timeout_opt, " -i %d ", cmd.ping.timeout);
#endif
					}
					if(strlen(cmd.target))
					{
						sprintf(cmd_string, "ping %s %s %s %s", count_opt, interval_opt, timeout_opt, cmd.target);
					}
				}
				else if(cmd.type == LIVE_TRACERT)
				{
					char wait_opt[128];
					char ttl_opt[128];

					*wait_opt = *ttl_opt = '\0';

					if(cmd.tracert.wait)
					{
						sprintf(wait_opt, " -w %d ", cmd.tracert.wait);
					}
					if(cmd.tracert.ttl)
					{
						sprintf(ttl_opt, " -m %d ", cmd.tracert.ttl);
					}
					if(strlen(cmd.target))
					{
						sprintf(cmd_string, "traceroute %s %s %s", ttl_opt, wait_opt, cmd.target);
					}
				}

				if(strlen(cmd_string))
				{
					COMMENT(cmd_string);
					fp = popen(cmd_string, "r");
				}

				if(fp)
				{
					int popen_fd;

					popen_fd = fileno(fp);
					states[popen_fd] = alloc_fd_state();
					states[popen_fd]->fp = fp;

					log[cmd.type] = popen_fd;
				}
			}
			else if(cmd.op == LIVE_GET)
			{
				int act_fd = log[cmd.type];

				if(act_fd != -1)
				{
					COMMENT("Prepare data to write");
					memcpy(state->buffer, states[act_fd], states[act_fd]->buffer_used);
					state->writing = 1;
					state->write_upto = states[act_fd]->buffer_used;

					COMMENT("Clear existing data");
					states[act_fd]->buffer_used = 0;

					COMMENT("Close finished data plane");
					if(states[act_fd]->done)
					{
						if(states[act_fd]->fp)
							pclose(states[act_fd]->fp);
						free_fd_state(states[act_fd]);
						states[act_fd] = NULL;
						log[cmd.type] = -1;
					}
				}
				else
				{
					COMMENT("There is no existing process");
					return -1;
				}
			}
		}
		else if(result == 0)
		{
			COMMENT("End of receiving data");
			return 1;
		}
	}

    return 0;
}

int
do_write(int fd, struct fd_state *state)
{
    while (state->n_written < state->write_upto) {
        ssize_t result = send(fd, state->buffer + state->n_written,
                              state->write_upto - state->n_written, 0);
        if (result < 0) {
            if (errno == EAGAIN)
                return 0;
            return -1;
        }
        assert(result != 0);

        state->n_written += result;
    }

    if (state->n_written == state->buffer_used)
        state->n_written = state->write_upto = state->buffer_used = 0;

    state->writing = 0;

	COMMENT("End of transfering data");
    return 1;
}

void
run(void)
{
    int listener;
    struct sockaddr_in sin;
    int i, maxfd;
    fd_set readset, writeset, exset;

    sin.sin_family = AF_INET;
    sin.sin_addr.s_addr = 0;
    sin.sin_port = htons(40713);

    for (i = 0; i < FD_SETSIZE; ++i)
        states[i] = NULL;
	for (i = 0; i < LIVE_TYPE_MAX; ++i)
		log[i] = -1;

    listener = socket(AF_INET, SOCK_STREAM, 0);
    make_nonblocking(listener);

#ifndef WIN32
    {
        int one = 1;
        setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
    }
#endif

    if (bind(listener, (struct sockaddr*)&sin, sizeof(sin)) < 0) {
        perror("bind");
        return;
    }

    if (listen(listener, 16)<0) {
        perror("listen");
        return;
    }

    FD_ZERO(&readset);
    FD_ZERO(&writeset);
    FD_ZERO(&exset);

    while (1) {
        maxfd = listener;

        FD_ZERO(&readset);
        FD_ZERO(&writeset);
        FD_ZERO(&exset);

        FD_SET(listener, &readset);

        for (i=0; i < FD_SETSIZE; ++i) {
            if (states[i]) {
                if (i > maxfd)
                    maxfd = i;
				if(!states[i]->done) {
                	FD_SET(i, &readset);
				}
                if (states[i]->writing) {
                    FD_SET(i, &writeset);
                }
            }
        }

        if (select(maxfd+1, &readset, &writeset, &exset, NULL) < 0) {
            perror("select");
            return;
        }

        if (FD_ISSET(listener, &readset)) {
            struct sockaddr_storage ss;
            socklen_t slen = sizeof(ss);
            int fd = accept(listener, (struct sockaddr*)&ss, &slen);
            if (fd < 0) {
                perror("accept");
            } else if (fd > FD_SETSIZE) {
                close(fd);
            } else {
                make_nonblocking(fd);
                states[fd] = alloc_fd_state();
                assert(states[fd]);/*XXX*/
            }
        }

        for (i=0; i < maxfd+1; ++i) {
            int r = 0;
            if (i == listener)
                continue;

            if (FD_ISSET(i, &readset)) {
                r = do_read(i, states[i]);
            }
            if (FD_ISSET(i, &writeset)) {
                r = do_write(i, states[i]);
            }
			if(r)
			{
				COMMENT("Close ctrl plane");
				free_fd_state(states[i]);
				states[i] = NULL;
				close(i);
			}
        }
    }
}

int
main(int c, char **v)
{
    setvbuf(stdout, NULL, _IONBF, 0);

    run();
    return 0;
}