--- /dev/null
+/*
+ * Copyright (c) 2009, Rambler media
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ * * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ * * Redistributions in binary form must reproduce the above copyright
+ * notice, this list of conditions and the following disclaimer in the
+ * documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY Rambler media ''AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL Rambler BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+/*
+ * Ipmark is custom plugin for marking ip with some weight, it understand several commands:
+ * - add <ip[/mask]> value
+ * - delete <ip[/mask]>
+ * - check <ip>
+ *
+ * This plugin is a sample of custom filters system in rspamd
+ */
+
+#include "../../../config.h"
+#include "../../../cfg_file.h"
+#include "../../../radix.h"
+
+#define ADD_COMMAND "add"
+#define DELETE_COMMAND "delete"
+#define CHECK_COMMAND "check"
+
+
+enum ipmark_command {
+ COMMAND_ADD,
+ COMMAND_DELETE,
+ COMMAND_CHECK
+};
+
+/* Exported functions */
+void module_init (struct config_file *cfg);
+void* before_connect (void);
+gboolean parse_line (const char *line, size_t len, char **output, void *user_data);
+void after_connect (char **output, char **log_line, void *user_data);
+void module_fin (void);
+
+/* Internal variables */
+char *filename = NULL;
+radix_tree_t *radix = NULL;
+
+/* Implementation */
+
+char *
+get_module_opt (struct config_file *cfg, char *module_name, char *opt_name)
+{
+ GList *cur_opt;
+ struct module_opt *cur;
+
+ cur_opt = g_hash_table_lookup (cfg->modules_opts, module_name);
+ if (cur_opt == NULL) {
+ return NULL;
+ }
+
+ while (cur_opt) {
+ cur = cur_opt->data;
+ if (strcmp (cur->param, opt_name) == 0) {
+ return cur->value;
+ }
+ cur_opt = g_list_next (cur_opt);
+ }
+
+ return NULL;
+}
+
+static gboolean
+parse_ipmask (const char *begin, struct in_addr *ina, int *mask, int *value)
+{
+ const char *pos;
+ char ip_buf[sizeof ("255.255.255.255")], mask_buf[3] = { '\0', '\0', '\0' }, *p;
+ int state = 1, dots = 0;
+
+ bzero (ip_buf, sizeof (ip_buf));
+ bzero (mask_buf, sizeof (mask_buf));
+ pos = begin;
+
+ while (*pos && state < 5) {
+ switch (state) {
+ case 1:
+ /* Begin parse ip */
+ if (g_ascii_isspace (*p)) {
+ state = 3;
+ }
+ else if (p - ip_buf >= sizeof (ip_buf) || dots > 3) {
+ return FALSE;
+ }
+ if (g_ascii_isdigit (*pos)) {
+ *p ++ = *pos ++;
+ }
+ else if (*pos == '.') {
+ *p ++ = *pos ++;
+ dots ++;
+ }
+ else if (*pos == '/') {
+ pos ++;
+ p = mask_buf;
+ state = 2;
+ }
+ else {
+ /* Invalid character */
+ return FALSE;
+ }
+ break;
+ case 2:
+ /* Parse mask */
+ if (g_ascii_isspace (*p)) {
+ state = 3;
+ }
+ else if (p - mask_buf > 2) {
+ return FALSE;
+ }
+ if (g_ascii_isdigit (*pos)) {
+ *p ++ = *pos ++;
+ }
+ else {
+ return FALSE;
+ }
+ break;
+ case 3:
+ if (!g_ascii_isspace (*p)) {
+ state = 4;
+ }
+ else {
+ p ++;
+ }
+ break;
+ case 4:
+ *value = strtol (p, NULL, 10);
+ state = 99;
+ break;
+ }
+ }
+
+ if (!inet_aton (ip_buf, ina)) {
+ return FALSE;
+ }
+
+ if (mask_buf[0] != '\0') {
+ /* Also parse mask */
+ *mask = (mask_buf[0] - '0') * 10 + mask_buf[1] - '0';
+ if (*mask > 32) {
+ return FALSE;
+ }
+ }
+ else {
+ *mask = 32;
+ }
+
+ *mask = 0xFFFFFFFF << (32 - *mask);
+
+ return TRUE;
+}
+
+static void
+read_radix_file (void)
+{
+ FILE *f;
+ char buf[BUFSIZ];
+ struct in_addr ina;
+ int mask, value;
+
+ f = fopen (filename, "r");
+ if (f != NULL) {
+ while (fgets (buf, sizeof (buf), f)) {
+ if (parse_ipmask (buf, &ina, &mask, &value)) {
+ (void)radix32tree_add (radix, ntohl (ina.s_addr), mask, (uintptr_t)value);
+ }
+ }
+
+ fclose (f);
+ }
+}
+
+static gboolean
+write_cb_func (uint32_t key, uint32_t level, uintptr_t value, void *user_data)
+{
+ FILE *f = user_data;
+ struct in_addr ina;
+
+ ina.s_addr = htonl (value);
+
+ fprintf (f, "%s/%d %d\n", inet_ntoa (ina), level, (int)value);
+
+ return FALSE;
+}
+
+static void
+write_radix_file (void)
+{
+ FILE *f;
+
+ /* Traverse throught radix tree */
+ f = fopen (filename, "w");
+ if (f != NULL) {
+ radix32tree_traverse (radix, write_cb_func, f);
+ fclose (f);
+ }
+}
+
+void
+module_init (struct config_file *cfg)
+{
+ char *value;
+
+ if (cfg && (value = get_module_opt (cfg, "ipmark", "file")) != NULL) {
+ filename = g_strdup (value);
+ }
+
+ radix = radix_tree_create ();
+ if (filename) {
+ read_radix_file ();
+ }
+}
+
+void *
+before_connect (void)
+{
+ /* In fact we do not need any session data, so just return NULL */
+ return NULL;
+}
+
+void
+module_fin (void)
+{
+ if (filename) {
+ write_radix_file ();
+ g_free (filename);
+ filename = NULL;
+ }
+ if (radix) {
+ radix_tree_free (radix);
+ radix = NULL;
+ }
+
+}
+
+gboolean
+parse_line (const char *line, size_t len, char **output, void *user_data)
+{
+ char ip_buf[sizeof ("255.255.255.255")], mask_buf[3] = {'\0', '\0', '\0'};
+ const char *p;
+ char *c = ip_buf, *err_str;
+ struct in_addr ina;
+ int state = 0, next_state, dots;
+ int16_t value;
+ uint32_t mask;
+ enum ipmark_command cmd;
+
+ /* Parse input line */
+ p = line;
+ while (p - line < len && state < 100) {
+ switch (state) {
+ case 0:
+ /* Expect command */
+ if (g_ascii_strncasecmp (line, ADD_COMMAND, sizeof (ADD_COMMAND) - 1) == 0) {
+ state = 99;
+ next_state = 1;
+ cmd = COMMAND_ADD;
+ p += sizeof (ADD_COMMAND);
+ }
+ else if (g_ascii_strncasecmp (line, DELETE_COMMAND, sizeof (DELETE_COMMAND) - 1) == 0) {
+ state = 99;
+ next_state = 1;
+ cmd = COMMAND_DELETE;
+ p += sizeof (DELETE_COMMAND);
+ }
+ else if (g_ascii_strncasecmp (line, CHECK_COMMAND, sizeof (CHECK_COMMAND) - 1) == 0) {
+ state = 99;
+ next_state = 1;
+ cmd = COMMAND_CHECK;
+ p += sizeof (CHECK_COMMAND);
+ }
+ else {
+ state = 100;
+ }
+ break;
+ case 1:
+ /* Expect ip or ipmask */
+ if (c - ip_buf >= sizeof (ip_buf) || dots > 3) {
+ state = 100;
+ }
+ if (g_ascii_isdigit (*p)) {
+ *c ++ = *p ++;
+ }
+ else if (*p == '.') {
+ *c ++ = *p ++;
+ dots ++;
+ }
+ else if (*p == '/') {
+ p ++;
+ c = mask_buf;
+ state = 2;
+ }
+ else if (g_ascii_isspace (*p)) {
+ if (cmd == COMMAND_ADD) {
+ next_state = 3;
+ }
+ else {
+ next_state = 100;
+ }
+ state = 99;
+ }
+ else {
+ /* Invalid character */
+ state = 100;
+ }
+ break;
+ case 2:
+ /* Parse mask */
+ if (c - mask_buf > 2) {
+ state = 100;
+ }
+ if (g_ascii_isdigit (*p)) {
+ *c ++ = *p ++;
+ }
+ else if (g_ascii_isspace (*p)) {
+ if (cmd == COMMAND_ADD) {
+ next_state = 3;
+ }
+ else {
+ next_state = 100;
+ }
+ state = 99;
+ }
+ else {
+ state = 100;
+ }
+ break;
+ case 3:
+ errno = 0;
+ value = strtol (p, &err_str, 10);
+ if (errno != 0) {
+ state = 100;
+ }
+ else {
+ state = 101;
+ }
+ break;
+ case 99:
+ /* Skip spaces */
+ if (g_ascii_isspace (*p)) {
+ p ++;
+ }
+ else {
+ state = next_state;
+ }
+ break;
+ }
+ }
+
+ if (state == 100 || !inet_aton (ip_buf, &ina)) {
+ /* Error occured */
+ *output = g_strdup ("ERR: invalid command");
+ return FALSE;
+ }
+
+ /* Process mask */
+ if (mask_buf[0] == '\0') {
+ /* Assume /32 mask */
+ mask = 0xFFFFFFFF;
+ }
+ else {
+ mask = (mask_buf[0] - '0') * 10 + mask_buf[1] - '0';
+ if (mask > 32) {
+ mask = 32;
+ }
+
+ mask = 0xFFFFFFFF << (32 - mask);
+ }
+
+ /* Process command */
+ switch (cmd) {
+ case COMMAND_ADD:
+ state = radix32tree_add (radix, ntohl (ina.s_addr), mask, (uintptr_t)value);
+ if (state == 0) {
+ *output = g_strdup_printf ("OK: new value %d", (int)value);
+ }
+ else if (state == -1) {
+ *output = g_strdup ("ERR: cannot insert value");
+ }
+ else {
+ *output = g_strdup_printf ("OK: new value %d", state);
+ }
+ break;
+ case COMMAND_DELETE:
+ if (radix32tree_delete (radix, ntohl (ina.s_addr), mask) == 0) {
+ *output = g_strdup ("OK: address deleted");
+ }
+ else {
+ *output = g_strdup ("ERR: address not found");
+ }
+ break;
+ case COMMAND_CHECK:
+ if ((value = radix32tree_find (radix, ntohl (ina.s_addr))) != RADIX_NO_VALUE) {
+ *output = g_strdup_printf ("OK: %d", (int)value);
+ }
+ else {
+ *output = g_strdup ("ERR: address not found");
+ }
+ break;
+ }
+
+ return TRUE;
+}
+
+void after_connect (char **output, char **log_line, void *user_data)
+{
+ /* Placeholder */
+ return;
+}
+
return tree;
}
-
-int
-radix32tree_insert (radix_tree_t * tree, uint32_t key, uint32_t mask, unsigned char value)
+enum radix_insert_type {
+ RADIX_INSERT,
+ RADIX_ADD,
+ RADIX_REPLACE
+};
+
+static uintptr_t
+radix32tree_insert_common (radix_tree_t * tree, uint32_t key, uint32_t mask, uintptr_t value, enum radix_insert_type type)
{
uint32_t bit;
radix_node_t *node, *next;
while (bit & mask) {
if (key & bit) {
next = node->right;
-
}
else {
next = node->left;
if (next) {
if (node->value != RADIX_NO_VALUE) {
- return 1;
+ /* Value was found, switch on insert type */
+ switch (type) {
+ case RADIX_INSERT:
+ return 1;
+ case RADIX_ADD:
+ node->value += value;
+ return value;
+ case RADIX_REPLACE:
+ node->value = value;
+ return 1;
+ }
}
node->value = value;
+ node->key = key;
return 0;
}
/* Inserting value in trie creating all path components */
}
node->value = value;
+ node->key = key;
return 0;
}
+int
+radix32tree_insert (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value)
+{
+ return (int)radix32tree_insert_common (tree, key, mask, value, RADIX_INSERT);
+}
+
+uintptr_t
+radix32tree_add (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value)
+{
+ return radix32tree_insert_common (tree, key, mask, value, RADIX_ADD);
+}
+
+int
+radix32tree_replace (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value)
+{
+ return (int)radix32tree_insert_common (tree, key, mask, value, RADIX_REPLACE);
+}
+
+/*
+ * per recursion step:
+ * ptr + ptr + ptr + int = 4 words
+ * result = 1 word
+ * 5 words total in stack
+ */
+static gboolean
+radix_recurse_nodes (radix_node_t *node, radix_tree_traverse_func func, void *user_data, int level)
+{
+ if (node->left) {
+ if (radix_recurse_nodes (node->left, func, user_data, level + 1)) {
+ return TRUE;
+ }
+ }
+
+ if (node->value != RADIX_NO_VALUE) {
+ if (func (node->key, level, node->value, user_data)) {
+ return TRUE;
+ }
+ }
+
+ if (node->right) {
+ if (radix_recurse_nodes (node->right, func, user_data, level + 1)) {
+ return TRUE;
+ }
+ }
+
+ return FALSE;
+}
+
+void
+radix32tree_traverse (radix_tree_t *tree, radix_tree_traverse_func func, void *user_data)
+{
+ radix_recurse_nodes (tree->root, func, user_data, 0);
+}
+
int
radix32tree_delete (radix_tree_t * tree, uint32_t key, uint32_t mask)
}
-unsigned char
+uintptr_t
radix32tree_find (radix_tree_t * tree, uint32_t key)
{
uint32_t bit;
#include "config.h"
#include "mem_pool.h"
-#define RADIX_NO_VALUE (unsigned char)-1
+#define RADIX_NO_VALUE (uintptr_t)-1
typedef struct radix_node_s radix_node_t;
radix_node_t *right;
radix_node_t *left;
radix_node_t *parent;
- unsigned char value;
+ uintptr_t value;
+ uint32_t key;
};
memory_pool_t *pool;
} radix_tree_t;
+typedef gboolean (*radix_tree_traverse_func)(uint32_t key, uint32_t mask, uintptr_t value, void *user_data);
+/**
+ * Create new radix tree
+ */
radix_tree_t *radix_tree_create ();
-int radix32tree_insert (radix_tree_t *tree, uint32_t key, uint32_t mask, unsigned char value);
+
+/**
+ * Insert value to radix tree
+ * returns: 1 if value already exists
+ * 0 if operation was successfull
+ * -1 if there was some error
+ */
+int radix32tree_insert (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value);
+
+/**
+ * Add value to radix tree or insert it if value does not exists
+ * returns: value if value already exists and was added
+ * 0 if value was inserted
+ * -1 if there was some error
+ */
+uintptr_t radix32tree_add (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value);
+
+/**
+ * Replace value in radix tree or insert it if value does not exists
+ * returns: 1 if value already exists and was replaced
+ * 0 if value was inserted
+ * -1 if there was some error
+ */
+int radix32tree_replace (radix_tree_t *tree, uint32_t key, uint32_t mask, uintptr_t value);
+
+/**
+ * Delete value from radix tree
+ * returns: 1 if value does not exist
+ * 0 if value was deleted
+ * -1 if there was some error
+ */
int radix32tree_delete (radix_tree_t *tree, uint32_t key, uint32_t mask);
-unsigned char radix32tree_find (radix_tree_t *tree, uint32_t key);
+
+/**
+ * Find value in radix tree
+ * returns: value if value was found
+ * RADIX_NO_VALUE if value was not found
+ */
+uintptr_t radix32tree_find (radix_tree_t *tree, uint32_t key);
+
+/**
+ * Traverse via the whole tree calling specified callback
+ */
+void radix32tree_traverse (radix_tree_t *tree, radix_tree_traverse_func func, void *user_data);
+
+/**
+ * Frees radix tree
+ */
void radix_tree_free (radix_tree_t *tree);
#endif