Implement a set of tree_* functions for storing arbitrary pointers in splay

trees with uint64_t keys. Also add x{m,c}alloc and xstrdup helpers.

ok gilles@

4 files changed, 277 insertions, 4 deletions

diff --git a/usr.sbin/smtpd/smtpd.h b/usr.sbin/smtpd/smtpd.h
index b2b72a90f4c..b98cf72f6f9 100644
--- a/usr.sbin/smtpd/smtpd.h
+++ b/usr.sbin/smtpd/smtpd.h
@@ -1,4 +1,4 @@
-/*	$OpenBSD: smtpd.h,v 1.313 2012/07/29 17:21:43 gilles Exp $	*/
+/*	$OpenBSD: smtpd.h,v 1.314 2012/08/07 21:47:57 eric Exp $	*/
 
 /*
  * Copyright (c) 2008 Gilles Chehade <gilles@openbsd.org>
@@ -1159,6 +1159,23 @@ size_t	stat_increment(int);
 size_t	stat_decrement(int);
 
 
+/* tree.c */
+SPLAY_HEAD(tree, treeentry);
+#define tree_init(t) SPLAY_INIT((t))
+#define tree_empty(t) SPLAY_EMPTY((t))
+int tree_check(struct tree *, uint64_t);
+void *tree_set(struct tree *, uint64_t, void *);
+void tree_xset(struct tree *, uint64_t, void *);
+void *tree_get(struct tree *, uint64_t);
+void *tree_xget(struct tree *, uint64_t);
+void *tree_pop(struct tree *, uint64_t);
+void *tree_xpop(struct tree *, uint64_t);
+int tree_poproot(struct tree *, uint64_t *, void **);
+int tree_root(struct tree *, uint64_t *, void **);
+int tree_iter(struct tree *, void **, uint64_t *, void **);
+void tree_merge(struct tree *, struct tree *);
+
+
 /* user.c */
 struct user_backend *user_backend_lookup(enum user_type);
 
@@ -1198,3 +1215,6 @@ int mvpurge(char *, char *);
 const char *parse_smtp_response(char *, size_t, char **, int *);
 int text_to_netaddr(struct netaddr *, char *);
 int text_to_relayhost(struct relayhost *, char *);
+void *xmalloc(size_t, const char *);
+void *xcalloc(size_t, size_t, const char *);
+char *xstrdup(const char *, const char *);
diff --git a/usr.sbin/smtpd/smtpd/Makefile b/usr.sbin/smtpd/smtpd/Makefile
index ebaae6d81d6..3d9467405e9 100644
--- a/usr.sbin/smtpd/smtpd/Makefile
+++ b/usr.sbin/smtpd/smtpd/Makefile
@@ -1,4 +1,4 @@
-#	$OpenBSD: Makefile,v 1.45 2012/07/14 09:05:58 eric Exp $
+#	$OpenBSD: Makefile,v 1.46 2012/08/07 21:47:58 eric Exp $
 
 .PATH:		${.CURDIR}/.. ${.CURDIR}/../../../lib/libc/asr
 
@@ -10,7 +10,7 @@ SRCS=		aliases.c auth.c bounce.c config.c control.c		\
 		mfa.c mfa_session.c mta.c mta_session.c	parse.y		\
 		queue.c queue_backend.c ruleset.c			\
 		scheduler.c scheduler_backend.c smtp.c smtp_session.c	\
-		smtpd.c	ssl.c ssl_privsep.c stats.c user.c util.c
+		smtpd.c	ssl.c ssl_privsep.c stats.c tree.c user.c util.c
 
 # backends
 SRCS+=		auth_bsd.c
diff --git a/usr.sbin/smtpd/tree.c b/usr.sbin/smtpd/tree.c
new file mode 100644
index 00000000000..ce147e33c02
--- /dev/null
+++ b/usr.sbin/smtpd/tree.c
@@ -0,0 +1,219 @@
+/*	$OpenBSD: tree.c,v 1.1 2012/08/07 21:47:58 eric Exp $	*/
+
+/*
+ * Copyright (c) 2012 Eric Faurot <eric@openbsd.org>
+ *
+ * Permission to use, copy, modify, and distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ */
+
+#include <sys/types.h>
+#include <sys/tree.h>
+
+#include <sys/socket.h>	/* for smtpd.h */
+#include <sys/queue.h>	/* for smtpd.h */
+#include <stdio.h>	/* for smtpd.h */
+#include <imsg.h>	/* for smtpd.h */
+
+#include <err.h>
+#include <inttypes.h>
+#include <stdlib.h>
+
+#include "smtpd.h"
+
+struct treeentry {
+	SPLAY_ENTRY(treeentry)	 entry;
+	uint64_t		 id;
+	void			*data;
+};
+
+static int treeentry_cmp(struct treeentry *, struct treeentry *);
+
+SPLAY_PROTOTYPE(tree, treeentry, entry, treeentry_cmp);
+
+int
+tree_check(struct tree *t, uint64_t id)
+{
+	struct treeentry	key;
+
+	key.id = id;
+	return (SPLAY_FIND(tree, t, &key) != NULL);
+}
+
+void *
+tree_set(struct tree *t, uint64_t id, void *data)
+{
+	struct treeentry	*entry, key;
+	char			*old;
+
+	key.id = id;
+	if ((entry = SPLAY_FIND(tree, t, &key)) == NULL) {
+		entry = xmalloc(sizeof *entry, "tree_set");
+		entry->id = id;
+		SPLAY_INSERT(tree, t, entry);
+		old = NULL;
+	} else
+		old = entry->data;
+
+	entry->data = data;
+
+	return (old);
+}
+
+void
+tree_xset(struct tree *t, uint64_t id, void *data)
+{
+	struct treeentry	*entry;
+
+	entry = xmalloc(sizeof *entry, "tree_xset");
+	entry->id = id;
+	entry->data = data;
+	if (SPLAY_INSERT(tree, t, entry))
+		errx(1, "tree_xset(%p, 0x%016"PRIx64 ")", t, id);
+}
+
+void *
+tree_get(struct tree *t, uint64_t id)
+{
+	struct treeentry	key, *entry;
+
+	key.id = id;
+	if ((entry = SPLAY_FIND(tree, t, &key)) == NULL)
+		return (NULL);
+
+	return (entry->data);
+}
+
+void *
+tree_xget(struct tree *t, uint64_t id)
+{
+	struct treeentry	key, *entry;
+
+	key.id = id;
+	if ((entry = SPLAY_FIND(tree, t, &key)) == NULL)
+		errx(1, "tree_get(%p, 0x%016"PRIx64 ")", t, id);
+
+	return (entry->data);
+}
+
+void *
+tree_pop(struct tree *t, uint64_t id)
+{
+	struct treeentry	key, *entry;
+	void			*data;
+
+	key.id = id;
+	if ((entry = SPLAY_FIND(tree, t, &key)) == NULL)
+		return (NULL);
+
+	data = entry->data;
+	SPLAY_REMOVE(tree, t, entry);
+	free(entry);
+
+	return (data);
+}
+
+void *
+tree_xpop(struct tree *t, uint64_t id)
+{
+	struct treeentry	key, *entry;
+	void			*data;
+
+	key.id = id;
+	if ((entry = SPLAY_FIND(tree, t, &key)) == NULL)
+		errx(1, "tree_xpop(%p, 0x%016" PRIx64 ")", t, id);
+
+	data = entry->data;
+	SPLAY_REMOVE(tree, t, entry);
+	free(entry);
+
+	return (data);
+}
+
+int
+tree_poproot(struct tree *t, uint64_t *id, void **data)
+{
+	struct treeentry	*entry;
+
+	entry = SPLAY_ROOT(t);
+	if (entry == NULL)
+		return (0);
+	if (id)
+		*id = entry->id;
+	if (data)
+		*data = entry->data;
+	SPLAY_REMOVE(tree, t, entry);
+	free(entry);
+	return (1);
+}
+
+int
+tree_root(struct tree *t, uint64_t *id, void **data)
+{
+	struct treeentry	*entry;
+
+	entry = SPLAY_ROOT(t);
+	if (entry == NULL)
+		return (0);
+	if (id)
+		*id = entry->id;
+	if (data)
+		*data = entry->data;
+	return (1);
+}
+
+int
+tree_iter(struct tree *t, void **hdl, uint64_t *id, void **data)
+{
+	struct treeentry *curr = *hdl;
+
+	if (curr == NULL)
+		curr = SPLAY_MIN(tree, t);
+	else
+		curr = SPLAY_NEXT(tree, t, curr);
+
+	if (curr) {
+		*hdl = curr;
+		if (id)
+			*id = curr->id;
+		if (data)
+			*data = curr->data;
+		return (1);
+	}
+
+	return (0);
+}
+
+void
+tree_merge(struct tree *dst, struct tree *src)
+{
+	struct treeentry	*entry;
+
+	while (!SPLAY_EMPTY(src)) {
+		entry = SPLAY_ROOT(src);
+		SPLAY_REMOVE(tree, src, entry);
+		if (SPLAY_INSERT(tree, dst, entry))
+			errx(1, "tree_merge: duplicate");
+	}
+}
+
+static int
+treeentry_cmp(struct treeentry *a, struct treeentry *b)
+{
+	if (a->id < b->id)
+		return (-1);
+	if (a->id > b->id)
+		return (1);
+	return (0);
+}
+
+SPLAY_GENERATE(tree, treeentry, entry, treeentry_cmp);
diff --git a/usr.sbin/smtpd/util.c b/usr.sbin/smtpd/util.c
index 53b0dd8d884..7000a366c1d 100644
--- a/usr.sbin/smtpd/util.c
+++ b/usr.sbin/smtpd/util.c
@@ -1,9 +1,10 @@
-/*	$OpenBSD: util.c,v 1.67 2012/07/29 17:21:43 gilles Exp $	*/
+/*	$OpenBSD: util.c,v 1.68 2012/08/07 21:47:58 eric Exp $	*/
 
 /*
  * Copyright (c) 2000,2001 Markus Friedl.  All rights reserved.
  * Copyright (c) 2008 Gilles Chehade <gilles@openbsd.org>
  * Copyright (c) 2009 Jacek Masiulaniec <jacekm@dobremiasto.net>
+ * Copyright (c) 2012 Eric Faurot <eric@openbsd.org>
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -85,6 +86,39 @@ const char *log_sockaddr(struct sockaddr *);
 
 static int temp_inet_net_pton_ipv6(const char *, void *, size_t);
 
+void *
+xmalloc(size_t size, const char *where)
+{
+	void	*r;
+
+	if ((r = malloc(size)) == NULL)
+		errx(1, "%s: malloc(%zu)", where, size);
+
+	return (r);
+}
+
+void *
+xcalloc(size_t nmemb, size_t size, const char *where)
+{
+	void	*r;
+
+	if ((r = calloc(nmemb, size)) == NULL)
+		errx(1, "%s: calloc(%zu, %zu)", where, nmemb, size);
+
+	return (r);
+}
+
+char *
+xstrdup(const char *str, const char *where)
+{
+	char	*r;
+
+	if ((r = strdup(str)) == NULL)
+		errx(1, "%s: strdup(%p)", where, str);
+
+	return (r);
+}
+
 int
 bsnprintf(char *str, size_t size, const char *format, ...)
 {