vncs.c

这是一个同样来自贝尔实验室的和UNIX有着渊源的操作系统, 其简洁的设计和实现易于我们学习和理解

#define	Image	IMAGE
#include "vnc.h"
#include "vncs.h"
#include "compat.h"
#include <cursor.h>
#include "screen.h"
#include "kbd.h"

#include <mp.h>
#include <libsec.h>

extern	Dev	drawdevtab;
extern	Dev	mousedevtab;
extern	Dev	consdevtab;

Dev	*devtab[] =
{
	&drawdevtab,
	&mousedevtab,
	&consdevtab,
	nil
};

static char *msgname[] = {
	[MPixFmt] = "MPixFmt",
	[MFixCmap] = "MFixCmap",
	[MSetEnc] = "MSetEnc",
	[MFrameReq] = "MFrameReq",
	[MKey] = "MKey",
	[MMouse] = "MMouse",
	[MCCut] = "MCCut",
};

static char *encname[] = {
	[EncRaw] = "raw",
	[EncCopyRect] = "copy rect",
	[EncRre] = "rre",
	[EncCorre] = "corre",
	[EncHextile] = "hextile",
	[EncZlib]	= "zlib",
	[EncTight]	= "zlibtight",
	[EncZHextile]	= "zhextile",
	[EncMouseWarp]	= "mousewarp",

};

/* 
 * list head. used to hold the list, the lock, dim, and pixelfmt
 */
struct {
	QLock;
	Vncs *head;
} clients;

int	shared;
int	sleeptime = 5;
int	verbose = 0;
char *cert;
char *pixchan = "r5g6b5";
static int	cmdpid;
static int	srvfd;
static int	exportfd;
static Vncs	**vncpriv;

static int parsedisplay(char*);
static void vnckill(char*, int, int);
static int vncannounce(char *net, int display, char *adir, int base);
static void noteshutdown(void*, char*);
static void vncaccept(Vncs*);
static int vncsfmt(Fmt*);
static void getremote(char*, char*);
static void vncname(char*, ...);
#pragma varargck argpos vncname 1

#pragma varargck type "V" Vncs*

void
usage(void)
{
	fprint(2, "usage: vncs [-v] [-c cert] [-d :display] [-g widthXheight] [-p pixelfmt] [cmd [args]...]\n");
	fprint(2, "\tto kill a server: vncs [-v] -k :display\n");
	exits("usage");
}

void
main(int argc, char **argv)
{
	int altnet, baseport, cfd, display, exnum, fd, h, killing, w;
	char adir[NETPATHLEN], ldir[NETPATHLEN];
	char net[NETPATHLEN], *p;
	char *rc[] = { "/bin/rc", "-i", nil };
	Vncs *v;

	fmtinstall('V', vncsfmt);
	display = -1;
	killing = 0;
	altnet = 0;
	w = 1024;
	h = 768;
	baseport = 5900;
	setnetmtpt(net, sizeof net, nil);
	ARGBEGIN{
	default:
		usage();
	case 'c':
		cert = EARGF(usage());
		baseport = 35729;
		break;
	case 'd':
		if(display != -1)
			usage();
		display = parsedisplay(EARGF(usage()));
		break;
	case 'g':
		p = EARGF(usage());
		w = strtol(p, &p, 10);
		if(*p != 'x' && *p != 'X' && *p != ' ' && *p != '	')
			usage();
		h = strtol(p+1, &p, 10);
		if(*p != 0)
			usage();
		break;
	case 'k':
		if(display != -1)
			usage();
		display = parsedisplay(EARGF(usage()));
		killing = 1;
		break;
	case 'p':
		pixchan = EARGF(usage());
		break;
/* DEBUGGING
	case 's':
		sleeptime = atoi(EARGF(usage()));
		break;
*/
	case 'v':
		verbose++;
		break;
	case 'x':
		p = EARGF(usage());
		setnetmtpt(net, sizeof net, p);
		altnet = 1;
		break;
	}ARGEND

	if(killing){
		vnckill(net, display, baseport);
		exits(nil);
	}

	if(altnet && !cert)
		sysfatal("announcing on alternate network requires TLS (-c)");

	if(argc == 0)
		argv = rc;

	/* easy exit */
	if(access(argv[0], AEXEC) < 0)
		sysfatal("access %s for exec: %r", argv[0]);

	/* background ourselves */
	switch(rfork(RFPROC|RFNAMEG|RFFDG|RFNOTEG)){
	case -1:
		sysfatal("rfork: %r");
	default:
		exits(nil);
	case 0:
		break;
	}

	vncpriv = privalloc();
	if(vncpriv == nil)
		sysfatal("privalloc: %r");

	/* start screen */
	initcompat();
	if(waserror())
		sysfatal("screeninit %dx%d %s: %s\n", w, h, pixchan, up->error);
	if(verbose)
		fprint(2, "geometry is %dx%d\n", w, h);
	screeninit(w, h, pixchan);
	poperror();

	/* start file system device slaves */
	exnum = exporter(devtab, &fd, &exportfd);

	/* rebuild /dev because the underlying connection might go away (ick) */
	unmount(nil, "/dev");
	bind("#c", "/dev", MREPL);

	/* run the command */
	switch(cmdpid = rfork(RFPROC|RFFDG|RFNOTEG|RFNAMEG|RFREND)){
	case -1:
		sysfatal("rfork: %r");
		break;
	case 0:
		if(mounter("/dev", MBEFORE, fd, exnum) < 0)
			sysfatal("mounter: %r");
		close(exportfd);
		close(0);
		close(1);
		close(2);
		open("/dev/cons", OREAD);
		open("/dev/cons", OWRITE);
		open("/dev/cons", OWRITE);
		exec(argv[0], argv);
		fprint(2, "exec %s: %r\n", argv[0]);
		_exits(nil);
	default:
		close(fd);
		break;
	}

	/* run the service */
	srvfd = vncannounce(net, display, adir, baseport);
	if(srvfd < 0)
		sysfatal("announce failed");
	if(verbose)
		fprint(2, "announced in %s\n", adir);

	atexit(shutdown);
	notify(noteshutdown);
	for(;;){
		vncname("listener");
		cfd = listen(adir, ldir);
		if(cfd < 0)
			break;
		if(verbose)
			fprint(2, "call in %s\n", ldir);
		fd = accept(cfd, ldir);
		if(fd < 0){
			close(cfd);
			continue;
		}
		v = mallocz(sizeof(Vncs), 1);
		if(v == nil){
			close(cfd);
			close(fd);
			continue;
		}
		v->ctlfd = cfd;
		v->datafd = fd;
		v->nproc = 1;
		v->ndead = 0;
		getremote(ldir, v->remote);
		strcpy(v->netpath, ldir);
		qlock(&clients);
		v->next = clients.head;
		clients.head = v;
		qunlock(&clients);
		vncaccept(v);
	}
	exits(0);
}

static int
parsedisplay(char *p)
{
	int n;

	if(*p != ':')
		usage();
	if(*p == 0)
		usage();
	n = strtol(p+1, &p, 10);
	if(*p != 0)
		usage();
	return n;
}

static void
getremote(char *ldir, char *remote)
{
	char buf[NETPATHLEN];
	int fd, n;

	snprint(buf, sizeof buf, "%s/remote", ldir);
	strcpy(remote, "<none>");
	if((fd = open(buf, OREAD)) < 0)
		return;
	n = readn(fd, remote, NETPATHLEN-1);
	close(fd);
	if(n < 0)
		return;
	remote[n] = 0;
	if(n>0 && remote[n-1] == '\n')
		remote[n-1] = 0;
}

static int
vncsfmt(Fmt *fmt)
{
	Vncs *v;

	v = va_arg(fmt->args, Vncs*);
	return fmtprint(fmt, "[%d] %s %s", getpid(), v->remote, v->netpath);
}

/*
 * We register exiting as an atexit handler in each proc, so that 
 * client procs need merely exit when something goes wrong.
 */
static void
vncclose(Vncs *v)
{
	Vncs **l;

	/* remove self from client list if there */
	qlock(&clients);
	for(l=&clients.head; *l; l=&(*l)->next)
		if(*l == v){
			*l = v->next;
			break;
		}
	qunlock(&clients);

	/* if last proc, free v */
	vnclock(v);
	if(++v->ndead < v->nproc){
		vncunlock(v);
		return;
	}

	freerlist(&v->rlist);
	vncterm(v);
	if(v->ctlfd >= 0)
		close(v->ctlfd);
	if(v->datafd >= 0)
		close(v->datafd);
	if(v->image)
		freememimage(v->image);
	free(v);
}

static void
exiting(void)
{
	vncclose(*vncpriv);
}

void
vnchungup(Vnc *v)
{
	if(verbose)
		fprint(2, "%V: hangup\n", (Vncs*)v);
	exits(0);	/* atexit and exiting() will take care of everything */
}

/*
 * Kill all clients except safe.
 * Used to start a non-shared client and at shutdown. 
 */
static void
killclients(Vncs *safe)
{
	Vncs *v;

	qlock(&clients);
	for(v=clients.head; v; v=v->next){
		if(v == safe)
			continue;
		if(v->ctlfd >= 0){
			hangup(v->ctlfd);
			close(v->ctlfd);
			v->ctlfd = -1;
			close(v->datafd);
			v->datafd = -1;
		}
	}
	qunlock(&clients);
}

/*
 * Kill the executing command and then kill everyone else.
 * Called to close up shop at the end of the day
 * and also if we get an unexpected note.
 */
static char killkin[] = "die vnc kin";
static void
killall(void)
{
	postnote(PNGROUP, cmdpid, "hangup");
	close(srvfd);
	srvfd = -1;
	close(exportfd);
	exportfd = -1;
	postnote(PNGROUP, getpid(), killkin);
}

void
shutdown(void)
{
	if(verbose)
		fprint(2, "vnc server shutdown\n");
	killall();
}

static void
noteshutdown(void*, char *msg)
{
	if(strcmp(msg, killkin) == 0)	/* already shutting down */
		noted(NDFLT);
	killall();
	noted(NDFLT);
}

/*
 * Kill a specific instance of a server.
 */
static void
vnckill(char *net, int display, int baseport)
{
	int fd, i, n, port;
	char buf[NETPATHLEN], *p;

	for(i=0;; i++){
		snprint(buf, sizeof buf, "%s/tcp/%d/local", net, i);
		if((fd = open(buf, OREAD)) < 0)
			sysfatal("did not find display");
		n = read(fd, buf, sizeof buf-1);
		close(fd);
		if(n <= 0)
			continue;
		buf[n] = 0;
		p = strchr(buf, '!');
		if(p == 0)
			continue;
		port = atoi(p+1);
		if(port != display+baseport)
			continue;
		snprint(buf, sizeof buf, "%s/tcp/%d/ctl", net, i);
		fd = open(buf, OWRITE);
		if(fd < 0)
			sysfatal("cannot open %s: %r", buf);
		if(write(fd, "hangup", 6) != 6)
			sysfatal("cannot hangup %s: %r", buf);
		close(fd);
		break;
	}
}

/*
 * Look for a port on which to announce.
 * If display != -1, we only try that one.
 * Otherwise we hunt.
 *
 * Returns the announce fd.
 */
static int
vncannounce(char *net, int display, char *adir, int base)
{
	int port, eport, fd;
	char addr[NETPATHLEN];

	if(display == -1){
		port = base;
		eport = base+50;
	}else{
		port = base+display;
		eport = port;
	}

	for(; port<=eport; port++){
		snprint(addr, sizeof addr, "%s/tcp!*!%d", net, port);
		if((fd = announce(addr, adir)) >= 0){
			fprint(2, "server started on display :%d\n", port-base);
			return fd;
		}
	}
	if(display == -1)
		fprint(2, "could not find any ports to announce\n");
	else
		fprint(2, "announce: %r\n");
	return -1;
}

/*
 * Handle a new connection.
 */
static void clientreadproc(Vncs*);
static void clientwriteproc(Vncs*);
static void chan2fmt(Pixfmt*, ulong);
static ulong fmt2chan(Pixfmt*);

static void
vncaccept(Vncs *v)
{
	char buf[32];
	int fd;
	TLSconn conn;

	/* caller returns to listen */
	switch(rfork(RFPROC|RFMEM|RFNAMEG)){
	case -1:
		fprint(2, "%V: fork failed: %r\n", v);
		vncclose(v);
		return;
	default:
		return;
	case 0:
		break;
	}
	*vncpriv = v;

	if(!atexit(exiting)){
		fprint(2, "%V: could not register atexit handler: %r; hanging up\n", v);
		exiting();
		exits(nil);
	}

	if(cert != nil){
		memset(&conn, 0, sizeof conn);
		conn.cert = readcert(cert, &conn.certlen);
		if(conn.cert == nil){
			fprint(2, "%V: could not read cert %s: %r; hanging up\n", v, cert);
			exits(nil);
		}
		fd = tlsServer(v->datafd, &conn);
		if(fd < 0){
			fprint(2, "%V: tlsServer: %r; hanging up\n", v);
			free(conn.cert);
			if(conn.sessionID)
				free(conn.sessionID);
			exits(nil);
		}
		close(v->datafd);
		v->datafd = fd;
		free(conn.cert);
		free(conn.sessionID);
	}
	vncinit(v->datafd, v->ctlfd, v);

	if(verbose)
		fprint(2, "%V: handshake\n", v);
	if(vncsrvhandshake(v) < 0){
		fprint(2, "%V: handshake failed; hanging up\n", v);
		exits(0);
	}
	if(verbose)
		fprint(2, "%V: auth\n", v);
	if(vncsrvauth(v) < 0){
		fprint(2, "%V: auth failed; hanging up\n", v);
		exits(0);
	}

	shared = vncrdchar(v);

	if(verbose)
		fprint(2, "%V: %sshared\n", v, shared ? "" : "not ");
	if(!shared)
		killclients(v);