3169ab17b9e005a9f95885610c14dd35b7d5b0cd
[nbd.git] / nbd-tester-client.c
1 /*
2  * Test client to test the NBD server. Doesn't do anything useful, except
3  * checking that the server does, actually, work.
4  *
5  * Note that the only 'real' test is to check the client against a kernel. If
6  * it works here but does not work in the kernel, then that's most likely a bug
7  * in this program and/or in nbd-server.
8  *
9  * Copyright(c) 2006  Wouter Verhelst
10  *
11  * This program is Free Software; you can redistribute it and/or modify it
12  * under the terms of the GNU General Public License as published by the Free
13  * Software Foundation, in version 2.
14  *
15  * This program is distributed in the hope that it will be useful, but WITHOUT
16  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
17  * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
18  * more details.
19  *
20  * You should have received a copy of the GNU General Public License along with
21  * this program; if not, write to the Free Software Foundation, Inc., 51
22  * Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
23  */
24 #include <stdlib.h>
25 #include <stdio.h>
26 #include <string.h>
27 #include <sys/time.h>
28 #include <sys/types.h>
29 #include <sys/socket.h>
30 #include <syslog.h>
31 #include <unistd.h>
32 #include "config.h"
33 #include "lfs.h"
34 #define MY_NAME "nbd-tester-client"
35 #include "cliserv.h"
36
37 #include <netinet/in.h>
38 #include <glib.h>
39
40 static gchar errstr[1024];
41 const static int errstr_len=1024;
42
43 static uint64_t size;
44
45 typedef enum {
46         CONNECTION_TYPE_NONE,
47         CONNECTION_TYPE_CONNECT,
48         CONNECTION_TYPE_INIT_PASSWD,
49         CONNECTION_TYPE_CLISERV,
50         CONNECTION_TYPE_FULL,
51 } CONNECTION_TYPE;
52
53 typedef enum {
54         CONNECTION_CLOSE_PROPERLY,
55         CONNECTION_CLOSE_FAST,
56 } CLOSE_TYPE;
57
58 inline int read_all(int f, void *buf, size_t len) {
59         ssize_t res;
60         size_t retval=0;
61
62         while(len>0) {
63                 if((res=read(f, buf, len)) <=0) {
64                         snprintf(errstr, errstr_len, "Read failed: %s", strerror(errno));
65                         return -1;
66                 }
67                 len-=res;
68                 buf+=res;
69                 retval+=res;
70         }
71         return retval;
72 }
73
74 int setup_connection(gchar *hostname, int port, gchar* name, CONNECTION_TYPE ctype) {
75         int sock;
76         struct hostent *host;
77         struct sockaddr_in addr;
78         char buf[256];
79         uint64_t mymagic = (name ? opts_magic : cliserv_magic);
80         u64 tmp64;
81         uint32_t tmp32 = 0;
82
83         sock=0;
84         if(ctype<CONNECTION_TYPE_CONNECT)
85                 goto end;
86         if((sock=socket(PF_INET, SOCK_STREAM, IPPROTO_TCP))<0) {
87                 strncpy(errstr, strerror(errno), errstr_len);
88                 goto err;
89         }
90         setmysockopt(sock);
91         if(!(host=gethostbyname(hostname))) {
92                 strncpy(errstr, strerror(errno), errstr_len);
93                 goto err_open;
94         }
95         addr.sin_family=AF_INET;
96         addr.sin_port=htons(port);
97         addr.sin_addr.s_addr=*((int *) host->h_addr);
98         if((connect(sock, (struct sockaddr *)&addr, sizeof(addr))<0)) {
99                 strncpy(errstr, strerror(errno), errstr_len);
100                 goto err_open;
101         }
102         if(ctype<CONNECTION_TYPE_INIT_PASSWD)
103                 goto end;
104         if(read_all(sock, buf, strlen(INIT_PASSWD))<0) {
105                 snprintf(errstr, errstr_len, "Could not read INIT_PASSWD: %s",
106                                 strerror(errno));
107                 goto err_open;
108         }
109         if(strlen(buf)==0) {
110                 snprintf(errstr, errstr_len, "Server closed connection");
111                 goto err_open;
112         }
113         if(strncmp(buf, INIT_PASSWD, strlen(INIT_PASSWD))) {
114                 snprintf(errstr, errstr_len, "INIT_PASSWD does not match");
115                 goto err_open;
116         }
117         if(ctype<CONNECTION_TYPE_CLISERV)
118                 goto end;
119         if(read_all(sock, &tmp64, sizeof(tmp64))<0) {
120                 snprintf(errstr, errstr_len, "Could not read cliserv_magic: %s",
121                                 strerror(errno));
122                 goto err_open;
123         }
124         tmp64=ntohll(tmp64);
125         if(tmp64 != mymagic) {
126                 strncpy(errstr, "mymagic does not match", errstr_len);
127                 goto err_open;
128         }
129         if(ctype<CONNECTION_TYPE_FULL)
130                 goto end;
131         if(!name) {
132                 read_all(sock, &size, sizeof(size));
133                 size=ntohll(size);
134                 read_all(sock, buf, 128);
135                 goto end;
136         }
137         /* flags */
138         read_all(sock, buf, sizeof(uint16_t));
139         /* reserved field */
140         write(sock, &tmp32, sizeof(tmp32));
141         /* magic */
142         tmp64 = htonll(opts_magic);
143         write(sock, &tmp64, sizeof(tmp64));
144         /* name */
145         tmp32 = htonl(NBD_OPT_EXPORT_NAME);
146         write(sock, &tmp32, sizeof(tmp32));
147         tmp32 = htonl((uint32_t)strlen(name));
148         write(sock, &tmp32, sizeof(tmp32));
149         write(sock, name, strlen(name));
150         read_all(sock, &size, sizeof(size));
151         size = ntohll(size);
152         read_all(sock, buf, sizeof(uint16_t)+124);
153         goto end;
154 err_open:
155         close(sock);
156 err:
157         sock=-1;
158 end:
159         return sock;
160 }
161
162 int close_connection(int sock, CLOSE_TYPE type) {
163         struct nbd_request req;
164         u64 counter=0;
165
166         switch(type) {
167                 case CONNECTION_CLOSE_PROPERLY:
168                         req.magic=htonl(NBD_REQUEST_MAGIC);
169                         req.type=htonl(NBD_CMD_DISC);
170                         memcpy(&(req.handle), &(counter), sizeof(counter));
171                         counter++;
172                         req.from=0;
173                         req.len=0;
174                         if(write(sock, &req, sizeof(req))<0) {
175                                 snprintf(errstr, errstr_len, "Could not write to socket: %s", strerror(errno));
176                                 return -1;
177                         }
178                 case CONNECTION_CLOSE_FAST:
179                         if(close(sock)<0) {
180                                 snprintf(errstr, errstr_len, "Could not close socket: %s", strerror(errno));
181                                 return -1;
182                         }
183                         break;
184                 default:
185                         g_critical("Your compiler is on crack!"); /* or I am buggy */
186                         return -1;
187         }
188         return 0;
189 }
190
191 int read_packet_check_header(int sock, size_t datasize, long long int curhandle) {
192         struct nbd_reply rep;
193         int retval=0;
194         char buf[datasize];
195
196         read_all(sock, &rep, sizeof(rep));
197         rep.magic=ntohl(rep.magic);
198         rep.error=ntohl(rep.error);
199         if(rep.magic!=NBD_REPLY_MAGIC) {
200                 snprintf(errstr, errstr_len, "Received package with incorrect reply_magic. Index of sent packages is %lld (0x%llX), received handle is %lld (0x%llX). Received magic 0x%lX, expected 0x%lX", curhandle, curhandle, *((u64*)rep.handle), *((u64*)rep.handle), (long unsigned int)rep.magic, (long unsigned int)NBD_REPLY_MAGIC);
201                 retval=-1;
202                 goto end;
203         }
204         if(rep.error) {
205                 snprintf(errstr, errstr_len, "Received error from server: %ld (0x%lX). Handle is %lld (0x%llX).", (long int)rep.error, (long unsigned int)rep.error, (long long int)(*((u64*)rep.handle)), *((u64*)rep.handle));
206                 retval=-1;
207                 goto end;
208         }
209         read_all(sock, &buf, datasize);
210
211 end:
212         return retval;
213 }
214
215 int throughput_test(gchar* hostname, int port, char* name, int sock, char sock_is_open, char close_sock) {
216         long long int i;
217         char buf[1024];
218         struct nbd_request req;
219         int requests=0;
220         fd_set set;
221         struct timeval tv;
222         struct timeval start;
223         struct timeval stop;
224         float timespan;
225         int speed;
226         char speedchar[2] = { '\0', '\0' };
227         int retval=0;
228         size_t tmp;
229         signed int do_write=TRUE;
230         pid_t mypid = getpid();
231
232         size=0;
233         if(!sock_is_open) {
234                 if((sock=setup_connection(hostname, port, name, CONNECTION_TYPE_FULL))<0) {
235                         g_warning("Could not open socket: %s", errstr);
236                         retval=-1;
237                         goto err;
238                 }
239         }
240         req.magic=htonl(NBD_REQUEST_MAGIC);
241         req.type=htonl(NBD_CMD_READ);
242         req.len=htonl(1024);
243         if(gettimeofday(&start, NULL)<0) {
244                 retval=-1;
245                 snprintf(errstr, errstr_len, "Could not measure start time: %s", strerror(errno));
246                 goto err_open;
247         }
248         for(i=0;i+1024<=size;i+=1024) {
249                 if(do_write) {
250                         memcpy(&(req.handle),&i,sizeof(i));
251                         req.from=htonll(i);
252                         write(sock, &req, sizeof(req));
253                         printf("%d: Requests(+): %d\n", (int)mypid, ++requests);
254                 }
255                 do {
256                         FD_ZERO(&set);
257                         FD_SET(sock, &set);
258                         tv.tv_sec=0;
259                         tv.tv_usec=0;
260                         select(sock+1, &set, NULL, NULL, &tv);
261                         if(FD_ISSET(sock, &set)) {
262                                 /* Okay, there's something ready for
263                                  * reading here */
264                                 if(read_packet_check_header(sock, 1024, i)<0) {
265                                         retval=-1;
266                                         goto err_open;
267                                 }
268                                 printf("%d: Requests(-): %d\n", (int)mypid, --requests);
269                         }
270                 } while FD_ISSET(sock, &set);
271                 /* Now wait until we can write again or until a second have
272                  * passed, whichever comes first*/
273                 FD_ZERO(&set);
274                 FD_SET(sock, &set);
275                 tv.tv_sec=1;
276                 tv.tv_usec=0;
277                 do_write=select(sock+1,NULL,&set,NULL,&tv);
278                 if(!do_write) printf("Select finished\n");
279                 if(do_write<0) {
280                         snprintf(errstr, errstr_len, "select: %s", strerror(errno));
281                         retval=-1;
282                         goto err_open;
283                 }
284         }
285         /* Now empty the read buffer */
286         do {
287                 FD_ZERO(&set);
288                 FD_SET(sock, &set);
289                 tv.tv_sec=0;
290                 tv.tv_usec=0;
291                 select(sock+1, &set, NULL, NULL, &tv);
292                 if(FD_ISSET(sock, &set)) {
293                         /* Okay, there's something ready for
294                          * reading here */
295                         read_packet_check_header(sock, 1024, i);
296                         printf("%d: Requests(-): %d\n", (int)mypid, --requests);
297                 }
298         } while (requests);
299         if(gettimeofday(&stop, NULL)<0) {
300                 retval=-1;
301                 snprintf(errstr, errstr_len, "Could not measure end time: %s", strerror(errno));
302                 goto err_open;
303         }
304         timespan=(float)(stop.tv_sec-start.tv_sec+(stop.tv_usec-start.tv_usec))/(float)1000000;
305         speed=(int)(size/timespan);
306         if(speed>1024) {
307                 speed>>=10;
308                 speedchar[0]='K';
309         }
310         if(speed>1024) {
311                 speed>>=10;
312                 speedchar[0]='M';
313         }
314         if(speed>1024) {
315                 speed>>=10;
316                 speedchar[0]='G';
317         }
318         g_message("%d: Throughput test complete. Took %.3f seconds to complete, %d%sB/s", (int)getpid(), timespan,speed,speedchar);
319
320 err_open:
321         if(close_sock) {
322                 close_connection(sock, CONNECTION_CLOSE_PROPERLY);
323         }
324 err:
325         return retval;
326 }
327
328 int main(int argc, char**argv) {
329         gchar *hostname;
330         long int p = 0;
331         int port;
332         char* name = NULL;
333         int sock=0;
334
335         if(argc<3) {
336                 g_message("%d: Not enough arguments", (int)getpid());
337                 g_message("%d: Usage: %s <hostname> <port>", (int)getpid(), argv[0]);
338                 g_message("%d: Or: %s <hostname> -N <exportname>", (int)getpid(), argv[0]);
339                 exit(EXIT_FAILURE);
340         }
341         logging();
342         hostname=g_strdup(argv[1]);
343         if(!strncmp(argv[2], "-N", 2)) {
344                 name = g_strdup(argv[3]);
345                 p = 10809;
346         } else {
347                 p=(strtol(argv[2], NULL, 0));
348                 if(p==LONG_MIN||p==LONG_MAX) {
349                         g_critical("Could not parse port number: %s", strerror(errno));
350                         exit(EXIT_FAILURE);
351                 }
352         }
353         port=(int)p;
354
355         if(throughput_test(hostname, port, name, sock, FALSE, TRUE)<0) {
356                 g_warning("Could not run throughput test: %s", errstr);
357                 exit(EXIT_FAILURE);
358         }
359
360         return 0;
361 }