1 /* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2 #include <stdio.h>
3 #include <stdlib.h>
4 #include <string.h>
5 
6 #include "protocol_extension.h"
7 
8 static const char *get_name(const void *cmd_cookie);
9 static bool accept_command(const void *cmd_cookie, void *cookie,
10                            int argc, token_t *argv, size_t *ndata,
11                            char **ptr);
12 static ENGINE_ERROR_CODE execute_command(const void *cmd_cookie, const void *cookie,
13                                          int argc, token_t *argv,
14                                          ENGINE_ERROR_CODE (*response_handler)(const void *cookie,
15                                                                   int nbytes,
16                                                                   const char *dta));
17 static void abort_command(const void *cmd_cookie, const void *cookie);
18 
19 static EXTENSION_ASCII_PROTOCOL_DESCRIPTOR scrub_descriptor = {
20     .get_name = get_name,
21     .accept = accept_command,
22     .execute = execute_command,
23     .abort = abort_command,
24     .cookie = &scrub_descriptor
25 };
26 
27 GET_SERVER_API server_api;
28 
get_name(const void * cmd_cookie)29 static const char *get_name(const void *cmd_cookie) {
30     return "scrub";
31 }
32 
accept_command(const void * cmd_cookie,void * cookie,int argc,token_t * argv,size_t * ndata,char ** ptr)33 static bool accept_command(const void *cmd_cookie, void *cookie,
34                            int argc, token_t *argv, size_t *ndata,
35                            char **ptr) {
36     return strcmp(argv[0].value, "scrub") == 0;
37 }
38 
my_response_handler(const void * key,uint16_t keylen,const void * ext,uint8_t extlen,const void * body,uint32_t bodylen,uint8_t datatype,uint16_t status,uint64_t cas,const void * cookie)39 static bool my_response_handler(const void *key, uint16_t keylen,
40                                 const void *ext, uint8_t extlen,
41                                 const void *body, uint32_t bodylen,
42                                 uint8_t datatype, uint16_t status,
43                                 uint64_t cas, const void *cookie)
44 {
45     uint16_t *rval = (uint16_t*)cookie;
46     *rval = status;
47     return true;
48 }
49 
execute_command(const void * cmd_cookie,const void * cookie,int argc,token_t * argv,ENGINE_ERROR_CODE (* response_handler)(const void * cookie,int nbytes,const char * dta))50 static ENGINE_ERROR_CODE execute_command(const void *cmd_cookie, const void *cookie,
51                                          int argc, token_t *argv,
52                                          ENGINE_ERROR_CODE (*response_handler)(const void *cookie,
53                                                                                int nbytes,
54                                                                                const char *dta))
55 {
56     protocol_binary_request_header request = {
57         .request.magic = (uint8_t)PROTOCOL_BINARY_REQ,
58         .request.opcode = PROTOCOL_BINARY_CMD_SCRUB
59     };
60 
61     uint16_t status = 0;
62     SERVER_HANDLE_V1 *server = server_api();
63     ENGINE_HANDLE_V1 *v1 = (ENGINE_HANDLE_V1*)server->engine;
64     if (v1 == NULL) {
65         return response_handler(cookie, 29, "SERVER_ERROR internal error\r\n");
66     }
67 
68     v1->unknown_command(server->engine, &status, &request, my_response_handler);
69 
70     if (status == PROTOCOL_BINARY_RESPONSE_SUCCESS) {
71         return response_handler(cookie, 4, "OK\r\n");
72     } else if (status == PROTOCOL_BINARY_RESPONSE_EBUSY) {
73         return response_handler(cookie, 6, "BUSY\r\n");
74     } else {
75         return response_handler(cookie, 7, "ERROR\r\n");
76     }
77 }
78 
abort_command(const void * cmd_cookie,const void * cookie)79 static void abort_command(const void *cmd_cookie, const void *cookie)
80 {
81     /* EMPTY */
82 }
83 
84 MEMCACHED_PUBLIC_API
memcached_extensions_initialize(const char * config,GET_SERVER_API get_server_api)85 EXTENSION_ERROR_CODE memcached_extensions_initialize(const char *config,
86                                                      GET_SERVER_API get_server_api) {
87     server_api = get_server_api;
88     SERVER_HANDLE_V1 *server = get_server_api();
89     if (server == NULL) {
90         return EXTENSION_FATAL;
91     }
92     if (!server->extension->register_extension(EXTENSION_ASCII_PROTOCOL,
93                                                &scrub_descriptor)) {
94         return EXTENSION_FATAL;
95     }
96 
97     return EXTENSION_SUCCESS;
98 }
99