xref: /3.0.3-GA/memcached/programs/mcbasher.cc (revision 464d6f0b)
1/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2// mcbasher is a program that keeps on running commands to a memcached
3// server in order to try to "crash" it. It does <b>not</b> try to verify
4// the correctness of the response it gets, because it's inteded to be used
5// with different mock engines and all we care about is if we're able to
6// run the commands or not.
7
8#include "config.h"
9
10#include <memcached/protocol_binary.h>
11
12#include <getopt.h>
13#include <cstdlib>
14#include <cstdio>
15#include <string>
16#include <string.h>
17#include <list>
18#include <cassert>
19#include <sys/types.h>
20#include <unistd.h>
21#include <fcntl.h>
22#include <poll.h>
23#include <cerrno>
24#include <platform/platform.h>
25
26using namespace std;
27
28
29class Connection {
30public:
31    Connection(const string &_host, const string &_port,
32               char *sndbuf, size_t sndbufsz, bool loop) :
33        sendBuffer(sndbuf), sendBufferSize(sndbufsz), loopSendBuffer(loop),
34        sendBufferOffset(0), devZeroSize(8192), devZero(new char[devZeroSize]),
35        sock(INVALID_SOCKET), host(_host), port(_port)
36    {
37
38    }
39
40    ~Connection() {
41        if (sock != INVALID_SOCKET) {
42            closesocket(sock);
43        }
44        delete []devZero;
45    }
46
47
48    void main(void) {
49        while (true) {
50            if (sock == INVALID_SOCKET) {
51                connect();
52            }
53
54            while (sock != INVALID_SOCKET) {
55                struct pollfd fds[1];
56                fds[0].fd = sock;
57                if (toSend == 0) {
58                    fds[0].events = POLLIN;
59                } else {
60                    fds[0].events = POLLIN | POLLOUT;
61                }
62                cb_assert(poll(fds, 1, -1) != -1);
63
64                if (fds[0].revents & POLLIN) {
65                    drainInput();
66                } else if (fds[0].revents & POLLOUT) {
67                    doSendData();
68                } else {
69                    // what is this?
70                    abort();
71                }
72            }
73        }
74    }
75
76protected:
77    /**
78     * Try to connect to the server
79     * @return false if we failed to connect to the server
80     */
81    bool connect(void)
82    {
83        struct addrinfo *ai = NULL;
84        struct addrinfo hints;
85
86        memset(&hints, 0, sizeof(hints));
87        hints.ai_family = AF_UNSPEC;
88        hints.ai_protocol = IPPROTO_TCP;
89        hints.ai_socktype = SOCK_STREAM;
90
91        if (getaddrinfo(host.c_str(), port.c_str(), &hints, &ai) != 0) {
92            return false;
93        }
94
95        struct addrinfo *e = ai;
96        do {
97            if ((sock = socket(e->ai_family, e->ai_socktype,
98                               e->ai_protocol)) != -1) {
99                if (::connect(sock, ai->ai_addr, ai->ai_addrlen) == -1) {
100                    close(sock);
101                    sock = -1;
102                }
103            }
104            if (sock == -1) {
105                e = e->ai_next;
106            } else {
107                break;
108            }
109        } while (e) ;
110
111        fcntl(sock, F_SETFL, fcntl(sock, F_GETFL) | O_NONBLOCK);
112
113        freeaddrinfo(ai);
114
115        sendBufferOffset = 0;
116        toSend = sendBufferSize - sendBufferOffset;
117        return true;
118    }
119
120    void drainInput(void) {
121        ssize_t nr;
122
123        while ((nr = recv(sock, devZero, devZeroSize, 0)) > 0) {
124            // Discard the data :)
125        }
126
127        if ((nr == -1 && errno != EWOULDBLOCK) || nr == 0) {
128            close(sock);
129            sock = -1;
130        }
131    }
132
133    void doSendData(void) {
134        ssize_t nw;
135        cb_assert(toSend > 0);
136
137        while ((nw = send(sock, sendBuffer+sendBufferOffset, toSend, 0)) > 0) {
138            sendBufferOffset += nw;
139            toSend = sendBufferSize - sendBufferOffset;
140            if (toSend == 0) {
141                if (loopSendBuffer) {
142                    sendBufferOffset = 0;
143                    toSend = sendBufferSize - sendBufferOffset;
144                } else {
145                    return;
146                }
147            }
148        }
149
150        if (nw == -1 && errno != EWOULDBLOCK) {
151            close(sock);
152            sock = -1;
153        }
154    }
155
156    char *sendBuffer;
157    size_t toSend;
158    size_t sendBufferSize;
159    bool loopSendBuffer;
160    size_t sendBufferOffset;
161
162    size_t devZeroSize;
163    char* devZero;
164    int sock;
165    const string host;
166    const string port;
167};
168
169
170
171extern "C" {
172    static void connection_main(void *arg) {
173        Connection *c = reinterpret_cast<Connection*>(arg);
174        c->main();
175    }
176}
177
178static void bangit(list<Connection*> conn)
179{
180    // @todo rewrite to use libevent and async io ;)
181    list<cb_thread_t> tids;
182    list<Connection*>::iterator iter;
183    for (iter = conn.begin(); iter != conn.end(); ++iter) {
184        cb_thread_t tid;
185        void *arg = reinterpret_cast<void*>(*iter);
186        cb_assert(cb_create_thread(&tid, connection_main, arg, 0) == 0);
187        tids.push_back(tid);
188    }
189
190    list<pthread_t>::iterator ii;
191    for (ii = tids.begin(); ii != tids.end(); ++ii) {
192        cb_assert(cb_join_thread(*ii) == 0);
193    }
194}
195
196static Connection *createTapConnection(const char *host, const char *port)
197{
198    protocol_binary_request_tap_connect *req;
199    char *message = new char[sizeof(req->bytes)];
200    size_t messagesize = sizeof(req->bytes);
201    req = reinterpret_cast<protocol_binary_request_tap_connect *>(message);
202    memset(req->bytes, 0, sizeof(req->bytes));
203
204    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
205    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_TAP_CONNECT;
206    req->message.header.request.keylen = htons(0);
207    req->message.header.request.extlen = 4;
208    req->message.header.request.bodylen = htonl(4);
209    req->message.body.flags = ntohl(TAP_CONNECT_FLAG_DUMP);
210
211    return new Connection(host, port, message, messagesize, false);
212}
213
214static size_t insertRandomGet(char *ptr, size_t buffsz)
215{
216    protocol_binary_request_get *req;
217
218    if (buffsz < (sizeof(req->bytes) + 1)) {
219        return 0;
220    }
221
222    req = reinterpret_cast<protocol_binary_request_get*>(ptr);
223    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
224    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_GET;
225    req->message.header.request.keylen = htons(1);
226    req->message.header.request.extlen = 0;
227    req->message.header.request.bodylen = htonl(1);
228    ptr += sizeof(req->bytes);
229    *ptr = 'a';
230    return sizeof(req->bytes) + 1;
231}
232
233static size_t insertRandomAdd(char *ptr, size_t buffsz)
234{
235    protocol_binary_request_add *req;
236
237    if (buffsz < (sizeof(req->bytes) + 1)) {
238        return 0;
239    }
240
241    req = reinterpret_cast<protocol_binary_request_add*>(ptr);
242    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
243    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_ADD;
244    req->message.header.request.keylen = htons(1);
245    req->message.header.request.extlen = 8;
246    req->message.header.request.bodylen = htonl(9);
247    req->message.body.flags = 0;
248    req->message.body.expiration = 0;
249    ptr += sizeof(req->bytes);
250    *ptr = 'a';
251    return sizeof(req->bytes) + 1;
252}
253
254static size_t insertRandomSet(char *ptr, size_t buffsz)
255{
256    protocol_binary_request_set *req;
257
258    if (buffsz < (sizeof(req->bytes) + 1)) {
259        return 0;
260    }
261
262    req = reinterpret_cast<protocol_binary_request_set*>(ptr);
263    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
264    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_SET;
265    req->message.header.request.keylen = htons(1);
266    req->message.header.request.extlen = 8;
267    req->message.header.request.bodylen = htonl(9);
268    req->message.body.flags = 0;
269    req->message.body.expiration = 0;
270    ptr += sizeof(req->bytes);
271    *ptr = 'a';
272    return sizeof(req->bytes) + 1;
273}
274
275static size_t insertRandomReplace(char *ptr, size_t buffsz)
276{
277    protocol_binary_request_replace *req;
278
279    if (buffsz < (sizeof(req->bytes) + 1)) {
280        return 0;
281    }
282
283    req = reinterpret_cast<protocol_binary_request_replace*>(ptr);
284    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
285    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_REPLACE;
286    req->message.header.request.keylen = htons(1);
287    req->message.header.request.extlen = 8;
288    req->message.header.request.bodylen = htonl(9);
289    req->message.body.flags = 0;
290    req->message.body.expiration = 0;
291    ptr += sizeof(req->bytes);
292    *ptr = 'a';
293    return sizeof(req->bytes) + 1;
294}
295
296static size_t insertRandomDelete(char *ptr, size_t buffsz)
297{
298    protocol_binary_request_delete *req;
299
300    if (buffsz < (sizeof(req->bytes) + 1)) {
301        return 0;
302    }
303
304    req = reinterpret_cast<protocol_binary_request_delete*>(ptr);
305    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
306    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_DELETE;
307    req->message.header.request.keylen = htons(1);
308    req->message.header.request.extlen = 0;
309    req->message.header.request.bodylen = htonl(1);
310    ptr += sizeof(req->bytes);
311    *ptr = 'a';
312    return sizeof(req->bytes) + 1;
313}
314
315static size_t insertRandomFlush(char *ptr, size_t buffsz)
316{
317    protocol_binary_request_flush *req;
318
319    if (buffsz < sizeof(req->bytes)) {
320        return 0;
321    }
322
323    req = reinterpret_cast<protocol_binary_request_flush*>(ptr);
324    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
325    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_FLUSH;
326    req->message.header.request.extlen = 4;
327    req->message.header.request.bodylen = htonl(4);
328    return sizeof(req->bytes);
329}
330
331static size_t insertRandomNoop(char *ptr, size_t buffsz)
332{
333    protocol_binary_request_noop *req;
334
335    if (buffsz < sizeof(req->bytes)) {
336        return 0;
337    }
338
339    req = reinterpret_cast<protocol_binary_request_noop*>(ptr);
340    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
341    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_NOOP;
342    return sizeof(req->bytes);
343}
344
345static size_t insertRandomVersion(char *ptr, size_t buffsz)
346{
347    protocol_binary_request_version *req;
348
349    if (buffsz < sizeof(req->bytes)) {
350        return 0;
351    }
352
353    req = reinterpret_cast<protocol_binary_request_version*>(ptr);
354    req->message.header.request.magic = PROTOCOL_BINARY_REQ;
355    req->message.header.request.opcode = PROTOCOL_BINARY_CMD_VERSION;
356    return sizeof(req->bytes);
357}
358
359static Connection *createGetSetConnection(const char *host, const char *port)
360{
361    const size_t buffersize = 1024*1024;
362    char *message = new char[buffersize];
363    memset(message, 0, buffersize);
364    size_t msgsize = 0;
365    char *curr = message;
366    while (msgsize < buffersize) {
367        long id = random() % 10;
368        size_t size;
369        switch (id) {
370        case 0:
371            size = insertRandomGet(curr, buffersize - msgsize);
372            break;
373        case 1:
374            size = insertRandomAdd(curr, buffersize - msgsize);
375            break;
376        case 2:
377            size = insertRandomReplace(curr, buffersize - msgsize);
378            break;
379        case 3:
380            size = insertRandomDelete(curr, buffersize - msgsize);
381            break;
382        case 4:
383            size = insertRandomFlush(curr, buffersize - msgsize);
384            break;
385        case 5:
386            size = insertRandomNoop(curr, buffersize - msgsize);
387            break;
388        case 6:
389            size = insertRandomVersion(curr, buffersize - msgsize);
390            break;
391        default:
392            size = insertRandomSet(curr, buffersize - msgsize);
393        }
394        curr += size;
395        msgsize += size;
396        if (size == 0) {
397            break;
398        }
399    }
400
401    return new Connection(host, port, message, msgsize, true);
402}
403
404
405/**
406 * Program entry point. Connect to a memcached server and use the binary
407 * protocol to retrieve a given set of stats.
408 *
409 * @param argc argument count
410 * @param argv argument vector
411 * @return 0 if success, error code otherwise
412 */
413int main(int argc, char **argv)
414{
415    int cmd;
416    const char *port = "11211";
417    const char *host = NULL;
418    int connections = 10;
419    char *ptr;
420
421    /* Initialize the socket subsystem */
422    cb_initialize_sockets();
423
424    while ((cmd = getopt(argc, argv, "h:p:c:")) != EOF) {
425        switch (cmd) {
426        case 'h' :
427            host = optarg;
428            ptr = strchr(optarg, ':');
429            if (ptr != NULL) {
430                *ptr = '\0';
431                port = ptr + 1;
432            }
433            break;
434        case 'p' :
435            port = optarg;
436            break;
437        case 'c' :
438            connections = atoi(optarg);
439            break;
440        default:
441            fprintf(stderr,
442                    "Usage mcbasher [-h host[:port]] [-p port] [-c connections]*\n");
443            return 1;
444        }
445    }
446
447    if (host == NULL) {
448        host = "localhost";
449    }
450
451    list<Connection*> conns;
452    for (int ii = 0; ii < connections; ++ii) {
453        Connection *c;
454        if ((ii & 1) == 0) {
455            c = createTapConnection(host, port);
456        } else {
457            c = createGetSetConnection(host, port);
458        }
459
460        conns.push_back(c);
461    }
462
463    bangit(conns);
464
465    return 0;
466}
467