* Add first custom filter for making marks for ip addresses and networks

* Some additions to radix tree library:
  - allow tree traverse
  - add new insert methods (add and replace)
  - store keys in radix nodes (thought we can calculate key by bits, but I think that storing key is not too expensive)
  - values in a tree are now uintptr_t

CMakeLists.txt

@@ -7,7 +7,7 @@ PROJECT(rspamd C)
 
 SET(RSPAMD_VERSION_MAJOR 0)
 SET(RSPAMD_VERSION_MINOR 2)
-SET(RSPAMD_VERSION_PATCH 8)
+SET(RSPAMD_VERSION_PATCH 9)
 
 SET(RSPAMD_VERSION         "${RSPAMD_VERSION_MAJOR}.${RSPAMD_VERSION_MINOR}.${RSPAMD_VERSION_PATCH}")
 SET(RSPAMD_MASTER_SITE_URL "http://cebka.pp.ru/hg/rspamd")
@@ -412,6 +412,7 @@ ENDIF(ENABLE_LUA MATCHES "ON")
 
 ADD_SUBDIRECTORY(src/json)
 ADD_SUBDIRECTORY(src/evdns)
+ADD_SUBDIRECTORY(src/plugins/custom)
 
 SET(TOKENIZERSSRC  src/tokenizers/tokenizers.c
 				src/tokenizers/osb.c)

src/plugins/custom/CMakeLists.txt

@@ -0,0 +1 @@
+ADD_SUBDIRECTORY(ipmark)

src/plugins/custom/ipmark/CMakeLists.txt

@@ -0,0 +1,10 @@
+# IPmark plugin makefile
+SET(IPMARKSRC		  ipmark.c
+					  ../../../radix.c
+					  ../../../mem_pool.c
+)
+
+ADD_LIBRARY(rspamd_ipmark SHARED ${IPMARKSRC})
+TARGET_LINK_LIBRARIES(rspamd_ipmark ${GLIB2_LIBRARIES})
+
+INSTALL(TARGETS rspamd_ipmark DESTINATION lib)

src/plugins/custom/ipmark/ipmark.c

@@ -0,0 +1,429 @@
+/*
+ * 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;
+}
+

src/radix.c

@@ -55,9 +55,14 @@ radix_tree_create ()
 	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;
@@ -70,7 +75,6 @@ radix32tree_insert (radix_tree_t * tree, uint32_t key, uint32_t mask, unsigned c
 	while (bit & mask) {
 		if (key & bit) {
 			next = node->right;
-
 		}
 		else {
 			next = node->left;
@@ -86,10 +90,21 @@ radix32tree_insert (radix_tree_t * tree, uint32_t key, uint32_t mask, unsigned c
 
 	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 */
@@ -117,10 +132,65 @@ radix32tree_insert (radix_tree_t * tree, uint32_t key, uint32_t mask, unsigned c
 	}
 
 	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)
@@ -186,7 +256,7 @@ 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;

src/radix.h

@@ -4,7 +4,7 @@
 #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;
 
@@ -12,7 +12,8 @@ struct radix_node_s {
     radix_node_t *right;
     radix_node_t *left;
     radix_node_t *parent;
-    unsigned char value;
+    uintptr_t value;
+	uint32_t key;
 };
 
 
@@ -22,11 +23,60 @@ typedef struct {
 	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