Permalink
Branch: master
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
426 lines (370 sloc) 10.4 KB
#define _GNU_SOURCE
#include <arpa/inet.h>
#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <linux/nbd.h>
#include <linux/types.h>
#include <signal.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/wait.h>
#include <unistd.h>
#include <uuid/uuid.h>
#include "config.h"
#include "aio.h"
#include "block-swap.h"
#include "ioh.h"
extern void dump_swapstat(void);
static FILE *tracefile = NULL;
struct sock_info {
int sock;
};
struct BlockDriverState;
struct client_info {
struct BlockDriverState *bs;
int sock;
};
struct read_info {
struct client_info *ci;
uint8_t *buffer;
int len;
struct nbd_reply reply;
};
static void got_data(void *opaque);
static void shell(char *arg, ...) {
char *list[16];
int l = 0;
va_list ap;
va_start(ap, arg);
char *a = arg;
while (a) {
list[l++] = a;
a = va_arg(ap, char *);
}
list[l] = NULL;
if (!fork()) {
if (execvp(list[0], list)) {
err(1, "execvp error");
}
va_end(ap);
exit(1);
}
}
static inline int safe_read(int fd, void *buf, size_t sz)
{
uint8_t *b = buf;
size_t left = sz;
while (left) {
ssize_t r;
do {
r = read(fd, b, left);
} while (r < 0 && errno == EINTR);
if (r < 0) {
err(1, "nbd read failed");
} else if (r == 0) {
break;
}
left -= r;
b += r;
}
return sz;
}
static inline int safe_write(int fd, const void *buf, size_t sz)
{
const uint8_t *b = buf;
size_t left = sz;
while (left) {
ssize_t r;
do {
r = write(fd, b, left);
} while (r < 0 && errno == EINTR);
if (r < 0) {
err(1, "nbd write failed");
} else if (r == 0) {
break;
}
left -= r;
b += r;
}
return sz;
}
static void nbd_read_done(void *opaque, int ret) {
int r;
struct read_info *ri = opaque;
struct client_info *ci = ri->ci;
r = safe_write(ci->sock, &ri->reply, sizeof(ri->reply));
if (r != sizeof(ri->reply)) {
err(1, "sock write (a) failed");
}
r = safe_write(ci->sock, ri->buffer, ri->len);
if (r != ri->len) {
err(1, "sock write (b) failed");
}
free(ri->buffer);
free(ri);
swap_aio_add_wait_object(ci->sock, got_data, ci);
}
static void nbd_write_done(void *opaque, int ret) {
uint8_t *buffer = opaque;
free(buffer);
}
static void got_data(void *opaque)
{
struct client_info *ci = opaque;
struct nbd_request request;
struct nbd_reply reply = {};
int r;
r = safe_read(ci->sock, &request, sizeof(request));
if (r == sizeof(request)) {
memcpy(reply.handle, request.handle, sizeof(reply.handle));
reply.magic = htonl(NBD_REPLY_MAGIC);
reply.error = htonl(0);
assert(request.magic == htonl(NBD_REQUEST_MAGIC));
int len = ntohl(request.len);
switch(ntohl(request.type)) {
case NBD_CMD_FLUSH: {
printf("got flush\n");
swap_flush(ci->bs);
r = safe_write(ci->sock, &reply, sizeof(reply));
if (r != sizeof(reply)) {
err(1, "sock write (c) failed");
}
break;
}
case NBD_CMD_DISC: {
printf("got disc\n");
break;
}
case NBD_CMD_READ: {
int len = ntohl(request.len);
uint64_t offset = be64toh(request.from);
struct read_info *ri = malloc(sizeof(struct read_info));
ri->ci = ci;
ri->buffer = malloc(len);
ri->len = len;
ri->reply = reply;
if (tracefile) {
fprintf(tracefile, "%lx %x\n", offset / 512, len / 512);
}
swap_aio_read(ci->bs, offset / 512, ri->buffer, len / 512, nbd_read_done, ri);
break;
}
case NBD_CMD_WRITE: {
assert(!(len % 4096));
r = safe_write(ci->sock, &reply, sizeof(reply));
uint8_t *buffer = malloc(len);
int len = ntohl(request.len);
uint64_t offset = be64toh(request.from);
uint8_t *b;
int left;
for (left = len, b = buffer; left > 0; b += r, left -= r) {
r = safe_read(ci->sock, b, left);
if (r < 0) {
err(1, "sock read failed");
}
}
swap_aio_write(ci->bs, offset / 512, buffer, len / 512, nbd_write_done, buffer);
swap_aio_add_wait_object(ci->sock, got_data, ci);
break;
}
case NBD_CMD_TRIM: {
printf("got TRIM!\n");
assert(!(len % 4096));
r = safe_write(ci->sock, &reply, sizeof(reply));
int len = ntohl(request.len);
uint64_t offset = be64toh(request.from);
uint8_t *zero = calloc(1, len);
swap_aio_write(ci->bs, offset / 512, zero, len / 512, nbd_write_done, zero);
swap_aio_add_wait_object(ci->sock, got_data, ci);
break;
}
default: {
printf("default %x\n", ntohl(request.type));
r = safe_write(ci->sock, &reply, sizeof(reply));
if (r != sizeof(reply)) {
err(1, "sock write (d) failed");
}
swap_aio_add_wait_object(ci->sock, got_data, ci);
break;
}
};
}
}
volatile int should_exit = 0;
volatile int should_close = 0;
static BlockDriverState bs;
void signal_handler(int s)
{
if (s == SIGINT) {
should_exit = 1;
} else if (s == SIGHUP) {
should_close = 1;
}
}
int main(int argc, char **argv)
{
int r;
if (argc < 3) {
fprintf(stderr, "Ussage: %s [filename.swap] [statechange-script]\n", argv[0]);
exit(1);
}
--argc;
++argv;
if (argc >= 4 && !strcmp(argv[0], "-t")) {
printf("writing read-trace to %s\n", argv[1]);
tracefile = fopen(argv[1], "wb");
argc -= 2;
argv += 2;
}
char *script = argv[1];
shell("/sbin/modprobe", "nbd", NULL);
r = mlockall(MCL_CURRENT | MCL_FUTURE);
if (r < 0) {
err(1, "mlockall failed");
}
char *fn = argv[0];
struct sigaction sig;
sig.sa_handler = signal_handler;
sigemptyset(&sig.sa_mask);
sig.sa_flags = 0;
for (int i = 1; i < 31; ++i) {
switch (i) {
case SIGHUP:
case SIGINT:
sigaction(i, &sig, NULL);
break;
default:
/* some versions of linux will disconnect nbd
* if getting signalled, so ignore as many as we can. */
signal(i, SIG_IGN);
break;
}
}
printf("opening swapimage %s...\n", fn);
ioh_init();
swap_aio_init();
int needs_format = 0;
if (!file_exists(fn)) {
r = swap_create(fn, 100 << 20, 0);
if (r < 0) {
printf("error creating %s\n", fn);
exit(1);
}
needs_format = 1;
}
r = swap_open(&bs, fn, 0);
if (r < 0) {
printf("error opening %s\n", fn);
exit(1);
}
int sp[2];
int sp2[2];
r = socketpair(AF_UNIX, SOCK_STREAM, 0, sp);
if (r) {
err(1, "socketpair (a) failed");
}
r = socketpair(AF_UNIX, SOCK_STREAM, 0, sp2);
if (r) {
err(1, "socketpair (b) failed");
}
char dev[32];
int device;
for (int i = 0; i < 16; ++i) {
sprintf(dev, "/dev/nbd%u", i);
device = open(dev, O_RDWR);
if (device < 0) {
if (errno == EBUSY) {
continue;
} else {
err(1, "nbd device open failed");
}
}
break;
}
printf("connecting to NBD...\n");
r = ioctl(device, NBD_SET_SIZE_BLOCKS, 1000 << 20);
if (r) {
err(1, "device ioctl (a) failed");
}
r = ioctl(device, NBD_SET_BLKSIZE, 0x1000);
if (r) {
err(1, "device ioctl (b) failed");
}
r = ioctl(device, NBD_SET_TIMEOUT, 120);
if (r) {
err(1, "device ioctl (c) failed");
}
r = ioctl(device, NBD_CLEAR_SOCK);
if (r) {
err(1, "device ioctl (d) failed");
}
int ok = 0;
int child = fork();
if (child == 0) {
close(sp[0]);
if(ioctl(device, NBD_SET_SOCK, sp[1]) == -1){
fprintf(stderr, "NBD_SET_SOCK failed %s\n", strerror(errno));
ok = 0;
r = safe_write(sp2[1], &ok, sizeof(ok));
} else {
ok = 1;
r = safe_write(sp2[1], &ok, sizeof(ok));
r = ioctl(device, NBD_DO_IT);
fprintf(stderr, "nbd device terminated %d\n", r);
if (r == -1)
fprintf(stderr, "%s\n", strerror(errno));
}
ioctl(device, NBD_CLEAR_QUE);
ioctl(device, NBD_CLEAR_SOCK);
exit(0);
}
r = safe_read(sp2[0], &ok, sizeof(ok));
if (r != sizeof(ok)) {
err(1, "socket pair read failed");
}
if (!ok) {
fprintf(stderr, "failed to init nbd, exiting!\n");
exit(1);
}
printf("configuring device using %s\n", script);
struct client_info *ci = malloc(sizeof(struct client_info));
ci->sock = sp[0];
ci->bs = &bs;
swap_aio_add_wait_object(ci->sock, got_data, ci);
uuid_t uuid;
swap_ioctl(&bs, 0, uuid);
char uuid_str[37];
uuid_unparse_lower(uuid, uuid_str);
char pid_str[16];
sprintf(pid_str, "%d", getpid());
setenv("DEVICE", dev, 1);
setenv("UUID", uuid_str, 1);
setenv("PID", pid_str, 1);
if (needs_format) {
shell(script, "create", NULL);
} else {
shell(script, "open", NULL);
}
while (!should_exit) {
swap_aio_wait();
if (should_close) {
shell(script, "close", NULL);
should_close = 0;
}
}
ioctl(device, NBD_DISCONNECT);
ioctl(device, NBD_CLEAR_SOCK);
int wstatus;
waitpid(child, &wstatus, 0);
swap_flush(&bs);
dump_swapstat();
swap_close(&bs);
swap_aio_close();
return 0;
}