path: root/controller/main.c

#include <stdio.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/select.h>
#include "global.h"
#include "line_reader.h"
#include "unzbuffer.h"


static bool writeall(int sock,const char *data,i64 len){
	i64 cursor=0;
	while(cursor<len){
		i64 nwr=write(sock,data+cursor,len-cursor);
		if(nwr<=0)return false;
		cursor+=nwr;
	}
	return true;
}

FILE *logfile=NULL;

struct frame_process {
	pid_t pid;
	i64 namelen;
	char *name;
	i64 usertime,systime,rss;
};

struct data_frame {
	struct timeval stamp;
	i64 numpids;
	struct frame_process *procs;
};

static i64 deserialise8(const u8 *buf){
	u64 value=0;
	for(i64 i=0;i<8;i++)value|=((u64)buf[i])<<(8*i);
	return (i64)value;
}

static i32 deserialise4(const u8 *buf){
	u32 value=0;
	for(i64 i=0;i<4;i++)value|=((u32)buf[i])<<(8*i);
	return (i32)value;
}

static void process_frame(const struct data_frame *frame){
	fprintf(logfile,"{\"stamp\":[%" PRIi64 ",%d],\"procs\":[",
			(i64)frame->stamp.tv_sec,frame->stamp.tv_usec);
	for(i64 i=0;i<frame->numpids;i++){
		const struct frame_process *p=&frame->procs[i];
		if(i!=0)fputc(',',logfile);
		fprintf(logfile,"{\"pid\":%d,\"name\":\"",p->pid);
		for(i64 j=0;j<p->namelen;j++){
			if(p->name[j]>=32&&p->name[j]<=126&&p->name[j]!='"'&&p->name[j]!='\\'){
				fputc(p->name[j],logfile);
			} else {
				fprintf(logfile,"\\x%c%c",
						"0123456789abcdef"[(unsigned char)p->name[j]/16],
						"0123456789abcdef"[(unsigned char)p->name[j]%16]);
			}
		}
		fprintf(logfile,"\",\"usertime\":%" PRIi64 ",\"systime\":%" PRIi64 ",\"rss\":%" PRIi64 "}",
				p->usertime,p->systime,p->rss);
	}
	fprintf(logfile,"]}\n");

	/*char buf[26];
	memcpy(buf,ctime(&frame->stamp.tv_sec),26);
	buf[24]='\0';
	fprintf(stderr,"%s: (%" PRIi64 ")\n",buf,frame->numpids);
	for(i64 i=0;i<frame->numpids;i++){
		const struct frame_process *p=&frame->procs[i];
		fprintf(stderr,"  (%d) \"%s\" u=%.2f s=%.2f rss=%" PRIi64 "\n",
				p->pid,p->name,p->usertime/1000000000.0,p->systime/1000000000.0,p->rss);
	}*/
}

static i64 read_frame(const u8 *buffer,i64 len){
	if(len<20)return 0;
	struct data_frame *frame=malloc(1,struct data_frame);
	frame->stamp.tv_sec=deserialise8(buffer);
	frame->stamp.tv_usec=deserialise4(buffer+8);
	i64 numpids=frame->numpids=deserialise8(buffer+12);
	i64 cursor=20;
	// fprintf(stderr,"numpids=%" PRIi64 "\n",numpids);
	if(len-cursor<numpids*36){
		free(frame);
		return 0;
	}
	// calloc to null unfilled pointers in procs
	frame->procs=calloc(numpids,struct frame_process);

	i64 pi;
	for(pi=0;pi<frame->numpids;pi++){
		// fprintf(stderr,"pi=%" PRIi64 "\n",pi);
		if(len-cursor<36){
			// fprintf(stderr,"insuff: pi\n");
			goto insufficient_data;
		}
		frame->procs[pi].pid=deserialise4(buffer+cursor);
		cursor+=4;
		i64 namelen=frame->procs[pi].namelen=deserialise8(buffer+cursor);
		cursor+=8;
		// fprintf(stderr,"  pid=%d  namelen=%" PRIi64 "\n",frame->procs[pi].pid,namelen);
		if(len-cursor<namelen+24){
			// fprintf(stderr,"insuff: namelen\n");
			goto insufficient_data;
		}
		frame->procs[pi].name=malloc(namelen+1,char);
		memcpy(frame->procs[pi].name,buffer+cursor,namelen);
		frame->procs[pi].name[namelen]='\0';
		cursor+=namelen;
		frame->procs[pi].usertime=deserialise8(buffer+cursor);
		cursor+=8;
		frame->procs[pi].systime=deserialise8(buffer+cursor);
		cursor+=8;
		frame->procs[pi].rss=deserialise8(buffer+cursor);
		cursor+=8;
	}

	process_frame(frame);
	for(i64 i=0;i<numpids;i++)free(frame->procs[i].name);
	free(frame->procs);
	free(frame);
	return cursor;

insufficient_data:
	if(frame->procs){
		for(i64 i=0;i<=pi;i++){
			if(frame->procs[i].name)free(frame->procs[i].name);
		}
		free(frame->procs);
	}
	free(frame);
	return 0;
}

static void frame_data_sink(u8 *newdata,i64 newdatalen,void *payload){
	(void)payload;
	static u8 *buffer=NULL;
	static i64 buffersize=0,bufferlen=0;

	// fprintf(stderr,"frame_data_sink(len=%" PRIi64 ")\n",newdatalen);

	if(buffer==NULL){
		buffersize=4096;
		buffer=malloc(buffersize,u8);
	}

	if(bufferlen+newdatalen>buffersize){
		do buffersize*=2; while(bufferlen+newdatalen>buffersize);
		buffer=realloc(buffer,buffersize,u8);
	}
	memcpy(buffer+bufferlen,newdata,newdatalen);
	bufferlen+=newdatalen;

	i64 cursor=0;
	while(cursor<bufferlen){
		i64 nr=read_frame(buffer+cursor,bufferlen-cursor);
		if(nr<=0)break;
		cursor+=nr;
	}
	if(cursor<bufferlen){
		memmove(buffer,buffer+cursor,bufferlen-cursor);
	}
	bufferlen-=cursor;
}

static int connect_server(const char *logfname,const char *ipaddr,const char *port){
	logfile=fopen(logfname,"w");
	if(logfile==NULL){
		fprintf(stderr,"Cannot open log file '%s'\n",logfname);
		return 1;
	}

	struct addrinfo hints,*res;
	memset(&hints,0,sizeof(hints));
	hints.ai_family=AF_UNSPEC;
	hints.ai_socktype=SOCK_STREAM;
	int ret=getaddrinfo(ipaddr,port,&hints,&res);
	if(ret!=0){
		fprintf(stderr,"getaddrinfo: %s\n",gai_strerror(ret));
		return 1;
	}

	int sock=-1;
	for(struct addrinfo *r=res;r!=NULL;r=r->ai_next){
		char buf[256];
		inet_ntop(r->ai_family,&((struct sockaddr_in*)r->ai_addr)->sin_addr,buf,sizeof buf);
		fprintf(stderr,"Trying: %s\n",buf);
		sock=socket(r->ai_family,r->ai_socktype,r->ai_protocol);
		if(sock<0){
			sock=-1;
			perror("socket");
			continue;
		}
		if(connect(sock,r->ai_addr,r->ai_addrlen)<0){
			close(sock);
			sock=-1;
			perror("connect");
			continue;
		}
		break;
	}
	freeaddrinfo(res);
	if(sock==-1)return 1;
	fprintf(stderr,"Connected.\n");

	struct line_reader *liner=line_reader_init(STDIN_FILENO);
	struct unzbuffer *z=unzbuffer_init(frame_data_sink,NULL);

	bool haderror=false;
	while(true){
		fd_set inset;
		FD_ZERO(&inset);
		FD_SET(STDIN_FILENO,&inset);
		FD_SET(sock,&inset);
		int ret=select(sock+1,&inset,NULL,NULL,NULL);
		if(ret<0){
			if(errno==EINTR)continue;
			perror("select");
			break;
		}

		if(FD_ISSET(STDIN_FILENO,&inset)){
			char buf[1024];
			i64 nr=read(STDIN_FILENO,buf,sizeof buf);
			if(nr<0){
				if(errno==EINTR)continue;
				perror("read");
				break;
			}
			if(nr==0)break;  // EOF
			line_reader_supply_data(liner,buf,nr);
			char *line;
			while((line=line_reader_get_line(liner,true))!=NULL){
				if(!writeall(sock,line,strlen(line))){
					fprintf(stderr,"Socket write error\n");
					goto cleanup;
				}
				// printf("Written <%s>\n",line);
				free(line);
			}
		}

		if(FD_ISSET(sock,&inset)){
			u8 buf[1024];
			i64 nr=read(sock,buf,sizeof buf);
			if(nr<0){
				if(errno==EINTR)continue;
				perror("read");
				break;
			}
			if(nr==0){
				fprintf(stderr,"End-of-file on socket\n");
				break;
			}
			// fprintf(stderr,"unzbuffer_write(nr=%" PRIi64 ")\n",nr);
			if(!unzbuffer_write(z,buf,nr)){
				haderror=true;
				break;
			}
		}
	}

cleanup:
	line_reader_destroy(liner);
	if(!unzbuffer_finish_destroy(z))return 1;
	if(haderror)return 1;
	return 0;
}

static int read_zz_file(const char *zzfname,const char *logfname){
	logfile=fopen(logfname,"w");
	if(logfile==NULL){
		fprintf(stderr,"Cannot open log file '%s'\n",logfname);
		return 1;
	}

	FILE *f=fopen(zzfname,"r");
	if(f==NULL){
		fprintf(stderr,"Cannot open zz file '%s'\n",zzfname);
		return 1;
	}

	struct unzbuffer *z=unzbuffer_init(frame_data_sink,NULL);
	u8 buf[4096];
	while(true){
		i64 nr=fread(buf,1,sizeof buf,f);
		if(nr<=0){
			if(ferror(f)){
				fprintf(stderr,"An error occurred reading the zz file '%s'\n",zzfname);
				return 1;
			}
			break;
		}
		if(!unzbuffer_write(z,buf,nr)){
			return 1;
		}
	}
	if(!unzbuffer_finish_destroy(z))return 1;
	return 0;
}

int main(int argc,char **argv){
	if(argc<=1){
		fprintf(stderr,"Usage:\n");
		fprintf(stderr,"  %s connect <logfile> <ip addr> [port=57575]\n",argv[0]);
		fprintf(stderr,"  %s read <zzfile> <logfile>\n",argv[0]);
		fprintf(stderr,"<logfile> will probably contain JSON.\n");
		fprintf(stderr,"<zzfile> should contain zlib-compressed protocol data.\n");
		return 1;
	}

	if(strcmp(argv[1],"connect")==0){
		if(argc<4||argc>5){
			fprintf(stderr,"Invalid number of arguments, run without arguments for usage\n");
			return 1;
		}
		const char *logfname=argv[2];
		const char *ipaddr=argv[3];
		const char *port=argc==5?argv[4]:"57575";
		return connect_server(logfname,ipaddr,port);
	} else if(strcmp(argv[1],"read")==0){
		if(argc!=4){
			fprintf(stderr,"Invalid number of arguments, run without arguments for usage\n");
			return 1;
		}
		const char *zzfname=argv[2];
		const char *logfname=argv[3];
		return read_zz_file(zzfname,logfname);
	} else {
		fprintf(stderr,"Invalid command '%s'; run without arguments for usage\n",argv[1]);
		return 1;
	}
}