diff --git a/nbd.c b/nbd.c index 82eb98db27..4822843b8a 100644 --- a/nbd.c +++ b/nbd.c @@ -600,6 +600,37 @@ struct NBDExport { QSIMPLEQ_HEAD(, NBDRequest) requests; }; +struct NBDClient { + int refcount; + void (*close)(NBDClient *client); + + NBDExport *exp; + int sock; +}; + +static void nbd_client_get(NBDClient *client) +{ + client->refcount++; +} + +static void nbd_client_put(NBDClient *client) +{ + if (--client->refcount == 0) { + g_free(client); + } +} + +static void nbd_client_close(NBDClient *client) +{ + qemu_set_fd_handler2(client->sock, NULL, NULL, NULL, NULL); + close(client->sock); + client->sock = -1; + if (client->close) { + client->close(client); + } + nbd_client_put(client); +} + static NBDRequest *nbd_request_get(NBDExport *exp) { NBDRequest *req; @@ -712,9 +743,11 @@ out: return rc; } -int nbd_trip(NBDExport *exp, int csock) +static int nbd_trip(NBDClient *client) { + NBDExport *exp = client->exp; NBDRequest *req = nbd_request_get(exp); + int csock = client->sock; struct nbd_request request; struct nbd_reply reply; int rc = -1; @@ -836,7 +869,30 @@ out: return rc; } -int nbd_negotiate(NBDExport *exp, int csock) +static void nbd_read(void *opaque) { - return nbd_send_negotiate(csock, exp->size, exp->nbdflags); + NBDClient *client = opaque; + + nbd_client_get(client); + if (nbd_trip(client) != 0) { + nbd_client_close(client); + } + + nbd_client_put(client); +} + +NBDClient *nbd_client_new(NBDExport *exp, int csock, + void (*close)(NBDClient *)) +{ + NBDClient *client; + if (nbd_send_negotiate(csock, exp->size, exp->nbdflags) == -1) { + return NULL; + } + client = g_malloc0(sizeof(NBDClient)); + client->refcount = 1; + client->exp = exp; + client->sock = csock; + client->close = close; + qemu_set_fd_handler2(csock, NULL, nbd_read, NULL, client); + return client; } diff --git a/nbd.h b/nbd.h index c77c2fd116..a8382f096c 100644 --- a/nbd.h +++ b/nbd.h @@ -76,11 +76,12 @@ int nbd_client(int fd); int nbd_disconnect(int fd); typedef struct NBDExport NBDExport; +typedef struct NBDClient NBDClient; NBDExport *nbd_export_new(BlockDriverState *bs, off_t dev_offset, off_t size, uint32_t nbdflags); void nbd_export_close(NBDExport *exp); -int nbd_negotiate(NBDExport *exp, int csock); -int nbd_trip(NBDExport *exp, int csock); +NBDClient *nbd_client_new(NBDExport *exp, int csock, + void (*close)(NBDClient *)); #endif diff --git a/qemu-nbd.c b/qemu-nbd.c index 347c776ab9..155b05840b 100644 --- a/qemu-nbd.c +++ b/qemu-nbd.c @@ -249,15 +249,10 @@ static int nbd_can_accept(void *opaque) return nb_fds < shared; } -static void nbd_read(void *opaque) +static void nbd_client_closed(NBDClient *client) { - int fd = (uintptr_t) opaque; - - if (nbd_trip(exp, fd) != 0) { - qemu_set_fd_handler2(fd, NULL, NULL, NULL, NULL); - close(fd); - nb_fds--; - } + nb_fds--; + qemu_notify_event(); } static void nbd_accept(void *opaque) @@ -268,8 +263,7 @@ static void nbd_accept(void *opaque) int fd = accept(server_fd, (struct sockaddr *)&addr, &addr_len); nbd_started = true; - if (fd != -1 && nbd_negotiate(exp, fd) != -1) { - qemu_set_fd_handler2(fd, NULL, nbd_read, NULL, (void *) (intptr_t) fd); + if (fd != -1 && nbd_client_new(exp, fd, nbd_client_closed)) { nb_fds++; } }