#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <netinet/in.h>
#include <linux/types.h>
#include <linux/nbd.h>

#include <tlmplib.h>
#include <tlmpnet.h>

static DEBUG_KEY DRTEST("drtest","Test program");

#ifdef WORDS_BIGENDIAN
static unsigned long long ntohll(unsigned long long a) {
        return a;
}
#else
static unsigned long long ntohll(unsigned long long a) {
        unsigned lo = a & 0xffffffff;
        unsigned hi = a >> 32U;
        lo = ntohl(lo);
        hi = ntohl(hi);
        return ((unsigned long long) lo) << 32U | hi;
}
#endif
#define htonll ntohll


#define INIT_PASSWD "NBDMAGIC"
#define CLISERV_MAGIC 0x00420281861253LL


struct nbd_init{
	char passwd[8];
	unsigned long long magic;
	unsigned long long size;
	char stuff[128];
};

union U_INIT{
	char binit[sizeof(nbd_init)];
	nbd_init init;
};

union U_REPLY{
	char breply[sizeof(nbd_reply)];
	nbd_reply reply;
};

static int bufsize = 100*1024;

static int nbd_read (
	_F_tcpconnect *c,
	unsigned long long &offset,
	unsigned long long size,
	unsigned long &readsize)
{
	readsize = size - offset;
	if (readsize > (unsigned)bufsize){
		readsize = bufsize;
	}
	nbd_request req;
	req.magic = ntohl(NBD_REQUEST_MAGIC);
	req.type = ntohl(NBD_CMD_READ);
	req.from = ntohll(offset);
	req.len  = ntohl(readsize);
	*((unsigned long long*)&req.handle) = offset;
	c->send (&req,sizeof(req));
	offset += readsize;
	debug_printf (DRTEST,"Reading %lu\n",readsize);
	return 0;
}
/*
	Write from stdin
*/
static int nbd_write (
	_F_tcpconnect *c,
	unsigned long long offset)
{
	int ret = -1;
	char buffer[100*1024+10];
	int len = fread (buffer,1,sizeof(buffer),stdin);
	if (len <= 0){
		tlmp_error ("Nothing read on stdin, nothing to write");
	}else if (len > 100*1024){
		tlmp_error ("Maximum input size: 100k");
	}else{
		nbd_request req;
		req.magic = ntohl(NBD_REQUEST_MAGIC);
		req.type = ntohl(NBD_CMD_WRITE);
		req.from = ntohll(offset);
		req.len  = ntohl(len);
		*((unsigned long long*)&req.handle) = offset;
		c->send (&req,sizeof(req));
		c->send (buffer,len);
		debug_printf (DRTEST,"Writing %d\n",len);
		ret = 0;
	}
	return ret;
}

<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal const char *host = NULL;
	glocal const char *port = NULL;
	glocal bool read = false;
	glocal int write = -1;
	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		extern const char *version;
		setproginfo ("drtest",version
			,"Test program for drsnapd\n"
			 "\n"
			 "Implement the NBD protocol. It can be used\n"
			 "to talk to an nbd-server.\n");
		setarg ('b',"bufsize","max read request",bufsize,false);
		setarg ('h',"host","Host running the nbd-server",glocal.host,true);
		setarg ('p',"port","TCP port to connect to",glocal.port,true);
		setarg ('r',"read","Read the device",glocal.read,false);
		setarg ('w',"write","Write (offset) to the device",glocal.write,false);
	</f>
	<f showerror>
		fprintf (stderr,"%s\n",msg);
	</f>
	<f main_noarg>
		int ret = -1;
		if (glocal.read && glocal.write != -1){
			usage();
		}else if (glocal.read || glocal.write != -1){
			glocal FILE *fout = stdout;
			glocal bool initseen = false;
			glocal unsigned initsofar = 0;
			glocal unsigned long readsize = 0;
			glocal unsigned readsofar = 0;
			glocal unsigned long long offset = 0;
			glocal U_INIT u;
			glocal U_REPLY r;
			<call tcpconnect>(glocal.host,glocal.port,5);
			<f init>
				// Nothing to do
				setrawmode (true);
			</f>
			<f oneline>
				if (glocal.initseen){
					unsigned len = info.linelen;
					debug_printf (DRTEST,"reading %u sofar %u\n"
						,len,glocal.readsofar);
					if (glocal.readsofar + len > sizeof(U_REPLY)+glocal.readsize){
						tlmp_error ("Too much data received");
						end = true;
					}else if (glocal.readsofar < sizeof(U_REPLY)){
						while (len > 0
							&& glocal.readsofar < sizeof(U_REPLY)){
							glocal.r.breply[glocal.readsofar++] = *line++;
							len--;
						}
						if (glocal.readsofar == sizeof(U_REPLY)){
							glocal.r.reply.magic = ntohl(glocal.r.reply.magic);
							glocal.r.reply.error = ntohl(glocal.r.reply.error);
							if (glocal.r.reply.error){
								tlmp_error("Some read or write error");
								end = true;
							}else if (glocal.write != -1){
								end = true;
							}
						}
					}
					if (len > 0){
						fwrite (line,1,len,glocal.fout);
						glocal.readsofar += len;
						if (glocal.readsofar == glocal.readsize
							+ sizeof(U_REPLY)){
							if (glocal.offset == glocal.u.init.size){
								debug_printf (DRTEST,"End of file\n");
								end = true;
							}else{
								nbd_read (this,glocal.offset,glocal.u.init.size,glocal.readsize);
								glocal.readsofar = 0;
							}
						}
					}
				}else{
					int needed = sizeof(U_INIT)-glocal.initsofar;
					debug_printf (DRTEST,"REC %d %d\n",info.linelen,needed);
					if (info.linelen <= needed){
						memcpy (glocal.u.binit+glocal.initsofar,line,info.linelen);
						glocal.initsofar += info.linelen;
						if (glocal.initsofar == sizeof(U_INIT)){
							glocal.initseen = true;
							glocal.u.init.magic = ntohll(glocal.u.init.magic);
							glocal.u.init.size = ntohll(glocal.u.init.size);
							if (glocal.u.init.magic != CLISERV_MAGIC){
								tlmp_error ("Bad magic");
								end = true;
							}else{
								debug_printf (DRTEST
									,"device size is %Ld\n"
									,glocal.u.init.size);
								if (glocal.read){
									nbd_read (this,glocal.offset,glocal.u.init.size,glocal.readsize);
								}else{
									nbd_write (this,glocal.write);
									glocal.readsize = 0;
								}
								glocal.readsofar = 0;
								glocal.initseen = true;
							}
						}
					}else{
						tlmp_error ("Too much init data");
						end = true;
					}
				}
			</f>
			</call>
		}else{
			tlmp_error ("Nothing to do for now");
		}
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>