Add support for binary protocol extensions

5 files changed, 175 insertions, 48 deletions

diff --git a/daemon/memcached.c b/daemon/memcached.c
index 7cc67e2..f4a328d 100644
--- a/daemon/memcached.c
+++ b/daemon/memcached.c
@@ -2583,6 +2583,37 @@ static void ship_tap_log(conn *c) {
     }
 }
 
+
+static ENGINE_ERROR_CODE default_unknown_command(EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *descriptor,
+                                                 ENGINE_HANDLE* handle,
+                                                 const void* cookie,
+                                                 protocol_binary_request_header *request,
+                                                 ADD_RESPONSE response)
+{
+    return settings.engine.v1->unknown_command(handle, cookie, request, response);
+}
+
+struct request_lookup {
+    EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *descriptor;
+    BINARY_COMMAND_CALLBACK callback;
+};
+
+static struct request_lookup request_handlers[0x100];
+
+static void initialize_binary_lookup_map(void) {
+    for (int ii = 0; ii < 0x100; ++ii) {
+        request_handlers[ii].descriptor = NULL;
+        request_handlers[ii].callback = default_unknown_command;
+    }
+}
+
+static void setup_binary_lookup_cmd(EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *descriptor,
+                                    uint8_t cmd,
+                                    BINARY_COMMAND_CALLBACK new_handler) {
+    request_handlers[cmd].descriptor = descriptor;
+    request_handlers[cmd].callback = new_handler;
+}
+
 static void process_bin_unknown_packet(conn *c) {
     void *packet = c->rcurr - (c->binary_header.request.bodylen +
                                sizeof(c->binary_header));
@@ -2592,24 +2623,31 @@ static void process_bin_unknown_packet(conn *c) {
     c->ewouldblock = false;
 
     if (ret == ENGINE_SUCCESS) {
-        ret = settings.engine.v1->unknown_command(settings.engine.v0, c, packet,
-                                                  binary_response_handler);
+        struct request_lookup *rq = request_handlers + c->binary_header.request.opcode;
+        ret = rq->callback(rq->descriptor, settings.engine.v0, c, packet,
+                           binary_response_handler);
     }
 
-    if (ret == ENGINE_SUCCESS) {
+    switch (ret) {
+    case ENGINE_SUCCESS:
         if (c->dynamic_buffer.buffer != NULL) {
             write_and_free(c, c->dynamic_buffer.buffer, c->dynamic_buffer.offset);
             c->dynamic_buffer.buffer = NULL;
         } else {
             conn_set_state(c, conn_new_cmd);
         }
-    } else if (ret == ENGINE_ENOTSUP) {
-        write_bin_packet(c, PROTOCOL_BINARY_RESPONSE_UNKNOWN_COMMAND, 0);
-    } else if (ret == ENGINE_EWOULDBLOCK) {
+        break;
+    case ENGINE_EWOULDBLOCK:
         c->ewouldblock = true;
-    } else {
-        /* FATAL ERROR, shut down connection */
+        break;
+    case ENGINE_DISCONNECT:
         conn_set_state(c, conn_closing);
+        break;
+    default:
+        // Release the dynamic buffer.. it may be partial..
+        free(c->dynamic_buffer.buffer);
+        c->dynamic_buffer.buffer = NULL;
+        write_bin_packet(c, engine_error_2_protocol_error(ret), 0);
     }
 }
 
@@ -2835,8 +2873,6 @@ static void process_bin_packet(conn *c) {
     }
 }
 
-
-
 typedef void (*RESPONSE_HANDLER)(conn*);
 /**
  * A map between the response packets op-code and the function to handle
@@ -3834,6 +3870,12 @@ static void process_stat_settings(ADD_STAT add_stats, void *c) {
          ptr = ptr->next) {
         APPEND_STAT("ascii_extension", "%s", ptr->get_name(ptr->cookie));
     }
+
+    for (EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *ptr = settings.extensions.binary;
+         ptr != NULL;
+         ptr = ptr->next) {
+        APPEND_STAT("binary_extension", "%s", ptr->get_name());
+    }
 }
 
 static char *process_stat(conn *c, token_t *tokens, const size_t ntokens) {
@@ -6530,6 +6572,28 @@ static bool register_extension(extension_type_t type, void *extension)
         }
         return true;
 
+    case EXTENSION_BINARY_PROTOCOL:
+        if (settings.extensions.binary != NULL) {
+            EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *last;
+            for (last = settings.extensions.binary; last->next != NULL;
+                 last = last->next) {
+                if (last == extension) {
+                    return false;
+                }
+            }
+            if (last == extension) {
+                return false;
+            }
+            last->next = extension;
+            last->next->next = NULL;
+        } else {
+            settings.extensions.binary = extension;
+            settings.extensions.binary->next = NULL;
+        }
+
+        ((EXTENSION_BINARY_PROTOCOL_DESCRIPTOR*)extension)->setup(setup_binary_lookup_cmd);
+        return true;
+
     default:
         return false;
     }
@@ -6592,6 +6656,13 @@ static void unregister_extension(extension_type_t type, void *extension)
         }
         break;
 
+
+    case EXTENSION_BINARY_PROTOCOL:
+        settings.extensions.logger->log(EXTENSION_LOG_WARNING, NULL,
+                                        "You can't unregister a binary command handler!");
+        abort();
+        break;
+
     default:
         ;
     }
@@ -6613,6 +6684,9 @@ static void* get_extension(extension_type_t type)
     case EXTENSION_ASCII_PROTOCOL:
         return settings.extensions.ascii;
 
+    case EXTENSION_BINARY_PROTOCOL:
+        return settings.extensions.binary;
+
     default:
         return NULL;
     }
@@ -6866,6 +6940,8 @@ int main (int argc, char **argv) {
     /* init settings */
     settings_init();
 
+    initialize_binary_lookup_map();
+
     if (memcached_initialize_stderr_logger(get_server_api) != EXTENSION_SUCCESS) {
         fprintf(stderr, "Failed to initialize log system\n");
         return EX_OSERR;
@@ -7345,11 +7421,11 @@ int main (int argc, char **argv) {
         exit(EXIT_FAILURE);
     }
 
-    if(!init_engine(engine_handle,engine_config,settings.extensions.logger)) {
+    if (!init_engine(engine_handle,engine_config,settings.extensions.logger)) {
         return false;
     }
 
-    if(settings.verbose > 0) {
+    if (settings.verbose > 0) {
         log_engine_details(engine_handle,settings.extensions.logger);
     }
     settings.engine.v1 = (ENGINE_HANDLE_V1 *) engine_handle;
diff --git a/daemon/memcached.h b/daemon/memcached.h
index c8a0983..32697e0 100644
--- a/daemon/memcached.h
+++ b/daemon/memcached.h
@@ -209,6 +209,7 @@ struct settings {
         EXTENSION_DAEMON_DESCRIPTOR *daemons;
         EXTENSION_LOGGER_DESCRIPTOR *logger;
         EXTENSION_ASCII_PROTOCOL_DESCRIPTOR *ascii;
+        EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *binary;
     } extensions;
 };
 
diff --git a/include/memcached/engine.h b/include/memcached/engine.h
index fa34410..6acea49 100644
--- a/include/memcached/engine.h
+++ b/include/memcached/engine.h
@@ -49,41 +49,6 @@ extern "C" {
 
 #define ENGINE_INTERFACE_VERSION 1
 
-    /**
-     * Callback for any function producing stats.
-     *
-     * @param key the stat's key
-     * @param klen length of the key
-     * @param val the stat's value in an ascii form (e.g. text form of a number)
-     * @param vlen length of the value
-     * @param cookie magic callback cookie
-     */
-    typedef void (*ADD_STAT)(const char *key, const uint16_t klen,
-                             const char *val, const uint32_t vlen,
-                             const void *cookie);
-
-    /**
-     * Callback for adding a response backet
-     * @param key The key to put in the response
-     * @param keylen The length of the key
-     * @param ext The data to put in the extended field in the response
-     * @param extlen The number of bytes in the ext field
-     * @param body The data body
-     * @param bodylen The number of bytes in the body
-     * @param datatype This is currently not used and should be set to 0
-     * @param status The status code of the return packet (see in protocol_binary
-     *               for the legal values)
-     * @param cas The cas to put in the return packet
-     * @param cookie The cookie provided by the frontend
-     * @return true if return message was successfully created, false if an
-     *              error occured that prevented the message from being sent
-     */
-    typedef bool (*ADD_RESPONSE)(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);
-
 
     /**
      * Abstract interface to an engine.
diff --git a/include/memcached/engine_common.h b/include/memcached/engine_common.h
index e918be4..dea0e9f 100644
--- a/include/memcached/engine_common.h
+++ b/include/memcached/engine_common.h
@@ -1,6 +1,9 @@
+/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
 #ifndef ENGINE_COMMON_H
 #define ENGINE_COMMON_H
 
+#include <stdint.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -9,6 +12,42 @@ extern "C" {
         uint64_t interface; /**< The version number on the engine structure */
     } ENGINE_HANDLE;
 
+    /**
+     * Callback for any function producing stats.
+     *
+     * @param key the stat's key
+     * @param klen length of the key
+     * @param val the stat's value in an ascii form (e.g. text form of a number)
+     * @param vlen length of the value
+     * @param cookie magic callback cookie
+     */
+    typedef void (*ADD_STAT)(const char *key, const uint16_t klen,
+                             const char *val, const uint32_t vlen,
+                             const void *cookie);
+
+    /**
+     * Callback for adding a response backet
+     * @param key The key to put in the response
+     * @param keylen The length of the key
+     * @param ext The data to put in the extended field in the response
+     * @param extlen The number of bytes in the ext field
+     * @param body The data body
+     * @param bodylen The number of bytes in the body
+     * @param datatype This is currently not used and should be set to 0
+     * @param status The status code of the return packet (see in protocol_binary
+     *               for the legal values)
+     * @param cas The cas to put in the return packet
+     * @param cookie The cookie provided by the frontend
+     * @return true if return message was successfully created, false if an
+     *              error occured that prevented the message from being sent
+     */
+    typedef bool (*ADD_RESPONSE)(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);
+
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/include/memcached/extension.h b/include/memcached/extension.h
index d978d75..6ec0040 100644
--- a/include/memcached/extension.h
+++ b/include/memcached/extension.h
@@ -3,6 +3,10 @@
 #define MEMCACHED_EXTENSION_H
 
 #include <stdbool.h>
+#include <stdint.h>
+#include <memcached/engine_common.h>
+#include <memcached/protocol_binary.h>
+#include <memcached/types.h>
 #include <memcached/server_api.h>
 
 #ifdef __cplusplus
@@ -43,7 +47,11 @@ extern "C" {
         /**
          * Command extension for the ASCII protocol
          */
-        EXTENSION_ASCII_PROTOCOL
+        EXTENSION_ASCII_PROTOCOL,
+        /**
+         * Command extension for the binary protocol
+         */
+        EXTENSION_BINARY_PROTOCOL
     } extension_type_t;
 
     /**
@@ -203,6 +211,44 @@ extern "C" {
         struct extension_ascii_protocol_descriptor *next;
     } EXTENSION_ASCII_PROTOCOL_DESCRIPTOR;
 
+
+    typedef struct extension_binary_protocol_descriptor EXTENSION_BINARY_PROTOCOL_DESCRIPTOR;
+
+    typedef ENGINE_ERROR_CODE (*BINARY_COMMAND_CALLBACK)(EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *descriptor,
+                                                         ENGINE_HANDLE* handle,
+                                                         const void* cookie,
+                                                         protocol_binary_request_header *request,
+                                                         ADD_RESPONSE response);
+
+    /**
+     * ASCII protocol extensions must provide the following descriptor to
+     * extend the capabilities of the ascii protocol. The memcached core
+     * will probe each command in the order they are registered, so you should
+     * register the most likely command to be used first (or you could register
+     * only one descriptor and do a better dispatch routine inside your own
+     * implementation of accept / execute).
+     */
+    struct extension_binary_protocol_descriptor {
+        /**
+         * Get the name of the descriptor. The memory area returned by this
+         * function has to be valid until the descriptor is unregistered.
+         */
+        const char* (*get_name)(void);
+
+        void (*setup)(void (*add)(EXTENSION_BINARY_PROTOCOL_DESCRIPTOR *descriptor,
+                                  uint8_t cmd,
+                                  BINARY_COMMAND_CALLBACK new_handler));
+
+        /**
+         * Deamon descriptors are stored in a linked list in the memcached
+         * core by using this pointer. Please do not modify this pointer
+         * by yourself until you have unregistered the descriptor.
+         * The <b>only</b> time it is safe for an extension to walk this
+         * list is during initialization of the modules.
+         */
+        struct extension_binary_protocol_descriptor *next;
+    };
+
     /**
      * The signature for the "memcached_extensions_initialize" function
      * exported from the loadable module.