--- a/src/auth/auth-client-connection.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/auth/auth-client-connection.c	Sun Aug 15 06:40:30 2004 +0300
@@ -25,22 +25,19 @@
 			     const void *data,
 			     struct auth_client_connection *conn)
 {
+	struct const_iovec iov[2];
 	ssize_t ret;
 
-	ret = o_stream_send(conn->output, reply, sizeof(*reply));
-	if ((size_t)ret == sizeof(*reply)) {
-		if (reply->data_size == 0) {
-			/* all sent */
-			auth_client_connection_unref(conn);
-			return;
-		}
+	iov[0].iov_base = reply;
+	iov[0].iov_len = sizeof(*reply);
+	iov[1].iov_base = data;
+	iov[2].iov_len = reply->data_size;
 
-		ret = o_stream_send(conn->output, data, reply->data_size);
-		if ((size_t)ret == reply->data_size) {
-			/* all sent */
-			auth_client_connection_unref(conn);
-			return;
-		}
+	ret = o_stream_sendv(conn->output, iov, 2);
+	if (ret == (ssize_t)(iov[0].iov_len + iov[1].iov_len)) {
+		/* all sent */
+		auth_client_connection_unref(conn);
+		return;
 	}
 
 	if (ret >= 0) {
@@ -188,6 +185,7 @@
 	static unsigned int connect_uid_counter = 0;
 	struct auth_client_connection *conn;
 	struct auth_client_handshake_reply handshake_reply;
+	struct const_iovec iov[2];
 
 	pool_t pool;
 
@@ -214,10 +212,12 @@
 	handshake_reply = *master->handshake_reply;
 	handshake_reply.connect_uid = conn->connect_uid;
 
-	if (o_stream_send(conn->output, &handshake_reply,
-			  sizeof(handshake_reply)) < 0 ||
-	    o_stream_send(conn->output, master->handshake_reply + 1,
-			  handshake_reply.data_size) < 0) {
+	iov[0].iov_base = &handshake_reply;
+	iov[0].iov_len = sizeof(handshake_reply);
+	iov[1].iov_base = master->handshake_reply + 1;
+	iov[2].iov_len = handshake_reply.data_size;
+
+	if (o_stream_sendv(conn->output, iov, 2) < 0) {
 		auth_client_connection_destroy(conn);
 		conn = NULL;
 	}

--- a/src/auth/auth-master-connection.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/auth/auth-master-connection.c	Sun Aug 15 06:40:30 2004 +0300
@@ -31,6 +31,7 @@
 	unsigned int tag;
 };
 
+static void master_output(void *context);
 static void auth_master_connection_close(struct auth_master_connection *conn);
 static int auth_master_connection_unref(struct auth_master_connection *conn);
 
@@ -94,23 +95,20 @@
 	ssize_t ret;
 
 	reply->tag = tag;
-	for (;;) {
-		ret = o_stream_send(conn->output, reply, reply_size);
-		if (ret < 0) {
-			/* master died, kill ourself too */
-			auth_master_connection_close(conn);
-			break;
-		}
 
-		if ((size_t)ret == reply_size)
-			break;
+	ret = o_stream_send(conn->output, reply, reply_size);
+	if (ret < 0) {
+		/* master died, kill ourself too */
+		auth_master_connection_close(conn);
+		return;
+	}
+	i_assert((size_t)ret == reply_size);
 
-		/* buffer full, we have to block */
-		i_warning("Master transmit buffer full, blocking..");
-		if (o_stream_flush(conn->output) < 0) {
-			/* transmit error, probably master died */
-			auth_master_connection_close(conn);
-			break;
+	if (o_stream_get_buffer_used_size(conn->output) >= MAX_OUTBUF_SIZE) {
+		/* buffer full, stop accepting more input */
+		if (conn->io != NULL) {
+			io_remove(conn->io);
+			conn->io = NULL;
 		}
 	}
 }
@@ -191,6 +189,23 @@
 	}
 }
 
+static void master_output(void *context)
+{
+	struct auth_master_connection *conn = context;
+	int ret;
+
+	if ((ret = o_stream_flush(conn->output)) < 0) {
+		/* transmit error, probably master died */
+		auth_master_connection_close(conn);
+		return;
+	}
+
+	if (o_stream_get_buffer_used_size(conn->output) <= MAX_OUTBUF_SIZE/2) {
+		/* allow input again */
+		conn->io = io_add(conn->fd, IO_READ, master_input, conn);
+	}
+}
+
 static void master_get_handshake_reply(struct auth_master_connection *master)
 {
 	struct mech_module_list *list;
@@ -242,7 +257,8 @@
 		io_remove(conn->io);
 
 	conn->output = o_stream_create_file(fd, default_pool,
-					    MAX_OUTBUF_SIZE, FALSE);
+					    (size_t)-1, FALSE);
+	o_stream_set_flush_callback(conn->output, master_output, conn);
 	conn->io = io_add(fd, IO_READ, master_input, conn);
 
 	conn->fd = fd;
@@ -269,12 +285,13 @@
 {
 	struct auth_master_handshake_reply reply;
 
-	/* just a note to master that we're ok. if we die before,
-	   master should shutdown itself. */
+	/* just a note to master that we're ok. if we die before, it means
+	   we're broken and a simple restart most likely won't help. */
 	if (conn->output != NULL) {
 		memset(&reply, 0, sizeof(reply));
 		reply.server_pid = conn->pid;
-		o_stream_send(conn->output, &reply, sizeof(reply));
+		if (o_stream_send(conn->output, &reply, sizeof(reply)) < 0)
+                        auth_master_connection_close(conn);
 	}
 }
 
@@ -290,8 +307,10 @@
 	o_stream_close(conn->output);
 	conn->output = NULL;
 
-	io_remove(conn->io);
-	conn->io = NULL;
+	if (conn->io != NULL) {
+		io_remove(conn->io);
+		conn->io = NULL;
+	}
 }
 
 void auth_master_connection_destroy(struct auth_master_connection *conn)

--- a/src/auth/password-scheme.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/auth/password-scheme.c	Sun Aug 15 06:40:30 2004 +0300
@@ -64,7 +64,7 @@
 			/* stop at next '$' */
 			p = strchr(p+1, '$');
 			if (p != NULL)
-				*password = t_strdup_until(*password, p);
+				*password = t_strdup_until(*password + 3, p);
 			return "MD5";
 		}
 	}

--- a/src/imap-login/client-authenticate.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap-login/client-authenticate.c	Sun Aug 15 06:40:30 2004 +0300
@@ -53,7 +53,6 @@
 	client_send_tagline(client, msg != NULL ?
 			    t_strconcat("NO ", msg, NULL) :
 			    "NO Authentication failed.");
-	o_stream_flush(client->output);
 
 	/* get back to normal client input */
 	if (client->common.io != NULL)
@@ -86,6 +85,9 @@
 				  const unsigned char *data, size_t size)
 {
 	buffer_t *buf;
+	const void *buf_data;
+	size_t buf_size;
+	ssize_t ret;
 
 	t_push();
 
@@ -95,9 +97,11 @@
 	base64_encode(data, size, buf);
 	buffer_append(buf, "\r\n", 2);
 
-	o_stream_send(client->output, buffer_get_data(buf, NULL),
-		      buffer_get_used_size(buf));
-	o_stream_flush(client->output);
+	buf_data = buffer_get_data(buf, &buf_size);
+	if ((ret = o_stream_send(client->output, buf_data, buf_size) < 0))
+		client_destroy(client, "Disconnected");
+	else if ((size_t)ret != buf_size)
+		client_destroy(client, "Transmit buffer full");
 
 	t_pop();
 }
@@ -315,6 +319,8 @@
 	info.remote_ip = client->common.ip;
 
 	client_ref(client);
+	o_stream_uncork(client->output);
+
 	client->common.auth_request =
 		auth_client_request_new(auth_client, NULL, &info,
 					authenticate_callback, client, &error);

--- a/src/imap-login/client.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap-login/client.c	Sun Aug 15 06:40:30 2004 +0300
@@ -15,10 +15,13 @@
 #include "auth-client.h"
 #include "ssl-proxy.h"
 
-/* max. size of one parameter in line */
-#define MAX_INBUF_SIZE 512
+/* max. size of one parameter in line, or max reply length in SASL
+   authentication */
+#define MAX_INBUF_SIZE 4096
 
-#define MAX_OUTBUF_SIZE 1024
+/* max. size of output buffer. if it gets full, the client is disconnected.
+   SASL authentication gives the largest output. */
+#define MAX_OUTBUF_SIZE 4096
 
 /* maximum length for IMAP command line. */
 #define MAX_IMAP_LINE 8192
@@ -100,10 +103,51 @@
 	return TRUE;
 }
 
-static int cmd_starttls(struct imap_client *client)
+static void client_start_tls(struct imap_client *client)
 {
 	int fd_ssl;
 
+	fd_ssl = ssl_proxy_new(client->common.fd, &client->common.ip,
+			       &client->common.proxy);
+	if (fd_ssl == -1) {
+		client_send_line(client, "* BYE TLS initialization failed.");
+		client_destroy(client, "TLS initialization failed.");
+		return;
+	}
+
+	client->tls = TRUE;
+	client->secured = TRUE;
+	client_set_title(client);
+
+	client->common.fd = fd_ssl;
+	i_stream_unref(client->input);
+	o_stream_unref(client->output);
+	imap_parser_destroy(client->parser);
+
+	/* CRLF is lost from buffer when streams are reopened. */
+	client->skip_line = FALSE;
+
+	client_open_streams(client, fd_ssl);
+	client->common.io = io_add(client->common.fd, IO_READ,
+				   client_input, client);
+}
+
+static void client_output_starttls(void *context)
+{
+	struct imap_client *client = context;
+	int ret;
+
+	if ((ret = o_stream_flush(client->output)) < 0) {
+		client_destroy(client, "Disconnected");
+		return;
+	}
+
+	if (ret > 0)
+		client_start_tls(client);
+}
+
+static int cmd_starttls(struct imap_client *client)
+{
 	if (client->tls) {
 		client_send_tagline(client, "BAD TLS is already active.");
 		return TRUE;
@@ -114,40 +158,21 @@
 		return TRUE;
 	}
 
-	client_send_tagline(client, "OK Begin TLS negotiation now.");
-	o_stream_flush(client->output);
-
-	/* must be removed before ssl_proxy_new(), since it may
-	   io_add() the same fd. */
+	/* remove input handler, SSL proxy gives us a new fd. we also have to
+	   remove it in case we have to wait for buffer to be flushed */
 	if (client->common.io != NULL) {
 		io_remove(client->common.io);
 		client->common.io = NULL;
 	}
 
-	fd_ssl = ssl_proxy_new(client->common.fd, &client->common.ip,
-			       &client->common.proxy);
-	if (fd_ssl != -1) {
-		client->tls = TRUE;
-		client->secured = TRUE;
-                client_set_title(client);
-
-		/* we skipped it already, so don't ignore next command */
-		client->skip_line = FALSE;
-
-		client->common.fd = fd_ssl;
-
-		i_stream_unref(client->input);
-		o_stream_unref(client->output);
-		imap_parser_destroy(client->parser);
-
-		client_open_streams(client, fd_ssl);
-		client->common.io = io_add(client->common.fd, IO_READ,
-					   client_input, client);
+	client_send_tagline(client, "OK Begin TLS negotiation now.");
+	if (o_stream_get_buffer_used_size(client->output) != 0) {
+		/* the buffer has to be flushed */
+		o_stream_set_flush_callback(client->output,
+					    client_output_starttls, client);
 	} else {
-		client_send_line(client, "* BYE TLS initialization failed.");
-		client_destroy(client, "TLS initialization failed.");
+		client_start_tls(client);
 	}
-
 	return TRUE;
 }
 
@@ -292,22 +317,21 @@
 	if (!client_read(client))
 		return;
 
+	client_ref(client);
+
 	if (!auth_client_is_connected(auth_client)) {
 		/* we're not yet connected to auth process -
 		   don't allow any commands */
 		client_send_line(client,
 			"* OK Waiting for authentication process to respond..");
 		client->input_blocked = TRUE;
-		return;
+	} else {
+		o_stream_cork(client->output);
+		while (client_handle_input(client)) ;
+		o_stream_uncork(client->output);
 	}
 
-	client_ref(client);
-
-	o_stream_cork(client->output);
-	while (client_handle_input(client)) ;
-
-	if (client_unref(client))
-		o_stream_flush(client->output);
+	client_unref(client);
 }
 
 static void client_destroy_oldest(void)
@@ -452,8 +476,18 @@
 
 void client_send_line(struct imap_client *client, const char *line)
 {
-	o_stream_send_str(client->output, line);
-	o_stream_send(client->output, "\r\n", 2);
+	struct const_iovec iov[2];
+	ssize_t ret;
+
+	iov[0].iov_base = line;
+	iov[0].iov_len = strlen(line);
+	iov[1].iov_base = "\r\n";
+	iov[1].iov_len = 2;
+
+	if ((ret = o_stream_sendv(client->output, iov, 2)) < 0)
+		client_destroy(client, "Disconnected");
+	else if ((size_t)ret != iov[0].iov_len + iov[1].iov_len)
+		client_destroy(client, "Transmit buffer full");
 }
 
 void client_send_tagline(struct imap_client *client, const char *line)

--- a/src/imap/Makefile.am	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/Makefile.am	Sun Aug 15 06:40:30 2004 +0300
@@ -64,7 +64,7 @@
 	commands-util.c \
 	imap-expunge.c \
 	imap-fetch.c \
-	imap-fetch-body-section.c \
+	imap-fetch-body.c \
 	imap-messageset.c \
 	imap-search.c \
 	imap-sort.c \

--- a/src/imap/client.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/client.c	Sun Aug 15 06:40:30 2004 +0300
@@ -1,4 +1,4 @@
-/* Copyright (C) 2002 Timo Sirainen */
+/* Copyright (C) 2002-2004 Timo Sirainen */
 
 #include "common.h"
 #include "ioloop.h"
@@ -10,60 +10,36 @@
 
 #include <stdlib.h>
 
-/* If we can't send a buffer in a minute, disconnect the client */
-#define CLIENT_OUTPUT_TIMEOUT (60*1000)
-
-/* If we don't soon receive expected data from client while processing
-   a command, disconnect the client */
-#define CLIENT_CMDINPUT_TIMEOUT CLIENT_OUTPUT_TIMEOUT
-
-/* Disconnect client when it sends too many bad commands in a row */
-#define CLIENT_MAX_BAD_COMMANDS 20
-
 extern struct mail_storage_callbacks mail_storage_callbacks;
 
 static struct client *my_client; /* we don't need more than one currently */
 static struct timeout *to_idle;
 
-static void client_output_timeout(void *context)
-{
-	struct client *client = context;
-
-	i_stream_close(client->input);
-	o_stream_close(client->output);
-}
-
-static void client_input_timeout(void *context)
-{
-	struct client *client = context;
-
-	client_disconnect_with_error(client,
-		"Disconnected for inactivity while waiting for command data.");
-}
+static void client_input(void *context);
+static void client_output(void *context);
 
 struct client *client_create(int hin, int hout, struct namespace *namespaces)
 {
 	struct client *client;
 
+	/* always use nonblocking I/O */
+	net_set_nonblock(hin, TRUE);
+	net_set_nonblock(hout, TRUE);
+
 	client = i_new(struct client, 1);
 	client->input = i_stream_create_file(hin, default_pool,
 					     imap_max_line_length, FALSE);
-	client->output = o_stream_create_file(hout, default_pool, 4096, FALSE);
+	client->output = o_stream_create_file(hout, default_pool,
+					      (size_t)-1, FALSE);
 
-	/* set timeout for reading expected data (eg. APPEND). This is
-	   different from the actual idle time. */
-	i_stream_set_blocking(client->input, CLIENT_CMDINPUT_TIMEOUT,
-			      client_input_timeout, client);
+	o_stream_set_flush_callback(client->output, client_output, client);
 
-	/* set timeout for sending data */
-	o_stream_set_blocking(client->output, CLIENT_OUTPUT_TIMEOUT,
-			      client_output_timeout, client);
-
-	client->io = io_add(hin, IO_READ, _client_input, client);
+	client->io = io_add(hin, IO_READ, client_input, client);
 	client->parser = imap_parser_create(client->input, client->output,
 					    imap_max_line_length);
         client->last_input = ioloop_time;
 
+	client->cmd_pool = pool_alloconly_create("command pool", 8192);
 	client->keywords.pool = pool_alloconly_create("mailbox_keywords", 512);
 	client->namespaces = namespaces;
 
@@ -83,14 +59,13 @@
 
 void client_destroy(struct client *client)
 {
-	o_stream_flush(client->output);
-
 	if (client->mailbox != NULL)
 		mailbox_close(client->mailbox);
 	namespace_deinit(client->namespaces);
 
 	imap_parser_destroy(client->parser);
-	io_remove(client->io);
+	if (client->io != NULL)
+		io_remove(client->io);
 
 	if (client->idle_to != NULL)
 		timeout_remove(client->idle_to);
@@ -99,6 +74,7 @@
 	o_stream_unref(client->output);
 
 	pool_unref(client->keywords.pool);
+	pool_unref(client->cmd_pool);
 	i_free(client);
 
 	/* quit the program */
@@ -108,7 +84,7 @@
 
 void client_disconnect(struct client *client)
 {
-	o_stream_flush(client->output);
+	(void)o_stream_flush(client->output);
 
 	i_stream_close(client->input);
 	o_stream_close(client->output);
@@ -168,7 +144,6 @@
 				    cmd, ": ", msg, NULL);
 	}
 
-	client->cmd_error = TRUE;
 	client_send_tagline(client, error);
 
 	if (++client->bad_counter >= CLIENT_MAX_BAD_COMMANDS) {
@@ -234,12 +209,22 @@
 
 void _client_reset_command(struct client *client)
 {
+	/* reset input idle time because command output might have taken a
+	   long time and we don't want to disconnect client immediately then */
+	client->last_input = ioloop_time;
+
+	client->command_pending = FALSE;
+	if (client->io == NULL) {
+		client->io = io_add(i_stream_get_fd(client->input),
+				    IO_READ, client_input, client);
+	}
+
 	client->cmd_tag = NULL;
 	client->cmd_name = NULL;
 	client->cmd_func = NULL;
-	client->cmd_error = FALSE;
 	client->cmd_uid = FALSE;
 
+	p_clear(client->cmd_pool);
         imap_parser_reset(client->parser);
 }
 
@@ -268,8 +253,7 @@
 {
         if (client->cmd_func != NULL) {
 		/* command is being executed - continue it */
-		client->input_skip_line = TRUE;
-		if (client->cmd_func(client) || client->cmd_error) {
+		if (client->cmd_func(client)) {
 			/* command execution was finished */
 			_client_reset_command(client);
                         client->bad_counter = 0;
@@ -293,12 +277,14 @@
                 client->cmd_tag = imap_parser_read_word(client->parser);
 		if (client->cmd_tag == NULL)
 			return FALSE; /* need more data */
+		client->cmd_tag = p_strdup(client->cmd_pool, client->cmd_tag);
 	}
 
 	if (client->cmd_name == NULL) {
-                client->cmd_name = imap_parser_read_word(client->parser);
+		client->cmd_name = imap_parser_read_word(client->parser);
 		if (client->cmd_name == NULL)
 			return FALSE; /* need more data */
+		client->cmd_name = p_strdup(client->cmd_pool, client->cmd_name);
 	}
 
 	if (client->cmd_name == '\0') {
@@ -315,7 +301,7 @@
 		_client_reset_command(client);
 	} else {
 		client->input_skip_line = TRUE;
-		if (client->cmd_func(client) || client->cmd_error) {
+		if (client->cmd_func(client)) {
 			/* command execution was finished */
 			_client_reset_command(client);
                         client->bad_counter = 0;
@@ -328,10 +314,16 @@
 	return TRUE;
 }
 
-void _client_input(void *context)
+static void client_input(void *context)
 {
 	struct client *client = context;
 
+	if (client->command_pending) {
+		/* already processing one command. wait. */
+		io_remove(client->io);
+		client->io = NULL;
+	}
+
 	client->last_input = ioloop_time;
 
 	switch (i_stream_read(client->input)) {
@@ -353,20 +345,56 @@
 	o_stream_cork(client->output);
 	while (client_handle_input(client))
 		;
-	o_stream_flush(client->output);
+	o_stream_uncork(client->output);
 
 	if (client->output->closed)
 		client_destroy(client);
 }
 
+static void client_output(void *context)
+{
+	struct client *client = context;
+	int ret;
+
+	if ((ret = o_stream_flush(client->output)) < 0) {
+		client_destroy(client);
+		return;
+	}
+
+	client->last_output = ioloop_time;
+
+	if (client->command_pending) {
+		o_stream_cork(client->output);
+		if (client->cmd_func(client)) {
+			/* command execution was finished */
+			_client_reset_command(client);
+                        client->bad_counter = 0;
+		}
+		o_stream_uncork(client->output);
+	}
+}
+
 static void idle_timeout(void *context __attr_unused__)
 {
+	time_t idle_time;
+
 	if (my_client == NULL)
 		return;
 
-	if (ioloop_time - my_client->last_input >= CLIENT_IDLE_TIMEOUT) {
-		client_send_line(my_client,
-				 "* BYE Disconnected for inactivity.");
+	idle_time = ioloop_time -
+		I_MAX(my_client->last_input, my_client->last_output);
+
+	if (my_client->command_pending &&
+	    o_stream_get_buffer_used_size(my_client->output) > 0 &&
+	    idle_time >= CLIENT_OUTPUT_TIMEOUT) {
+		/* client isn't reading our output */
+		client_destroy(my_client);
+	} else if (idle_time >= CLIENT_IDLE_TIMEOUT) {
+		/* client isn't sending us anything */
+		if (!my_client->command_pending) {
+			client_send_line(my_client,
+					 "* BYE Disconnected for inactivity.");
+		}
 		client_destroy(my_client);
 	}
 }

--- a/src/imap/client.h	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/client.h	Sun Aug 15 06:40:30 2004 +0300
@@ -27,18 +27,20 @@
 	unsigned int select_counter; /* increased when mailbox is changed */
 	uint32_t messages_count, recent_count;
 
-	time_t last_input;
+	time_t last_input, last_output;
 	unsigned int bad_counter;
 
 	struct imap_parser *parser;
-	const char *cmd_tag; /* tag of command (allocated from parser pool), */
-	const char *cmd_name; /* command name (allocated from parser pool) */
+	pool_t cmd_pool;
+	const char *cmd_tag;
+	const char *cmd_name;
 	command_func_t *cmd_func;
+	void *cmd_context;
 
 	struct timeout *idle_to;
 	unsigned int idle_expunge;
 
-	unsigned int cmd_error:1;
+	unsigned int command_pending:1;
 	unsigned int cmd_uid:1; /* used UID command */
 	unsigned int rawlog:1;
 	unsigned int input_skip_line:1; /* skip all the data until we've
@@ -73,7 +75,6 @@
 void clients_init(void);
 void clients_deinit(void);
 
-void _client_input(void *context);
 void _client_reset_command(struct client *client);
 
 #endif

--- a/src/imap/cmd-append.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/cmd-append.c	Sun Aug 15 06:40:30 2004 +0300
@@ -121,19 +121,16 @@
 
 			/* need more data */
 			ret = i_stream_read(client->input);
-			if (ret == -2) {
-				client_send_command_error(client,
-							  "Too long argument.");
-				break;
-			}
 			if (ret < 0) {
-				/* disconnected */
-				client->cmd_error = TRUE;
+				if (ret == -2) {
+					client_send_command_error(client,
+						"Too long argument.");
+				}
 				break;
 			}
 		}
 
-		if (client->cmd_error)
+		if (ret < 0)
 			break;
 
 		if (args->type == IMAP_ARG_EOL) {

--- a/src/imap/cmd-fetch.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/cmd-fetch.c	Sun Aug 15 06:40:30 2004 +0300
@@ -1,4 +1,4 @@
-/* Copyright (C) 2002 Timo Sirainen */
+/* Copyright (C) 2002-2004 Timo Sirainen */
 
 #include "common.h"
 #include "commands.h"
@@ -6,360 +6,64 @@
 #include "imap-search.h"
 #include "mail-search.h"
 
-/* Parse next digits in string into integer. Returns FALSE if the integer
-   becomes too big and wraps. */
-static int read_uoff_t(char **p, uoff_t *value)
+static int
+fetch_parse_args(struct client *client, struct imap_fetch_context *ctx,
+		 struct imap_arg *arg)
 {
-	uoff_t prev;
+	const char *str;
+
+	if (arg->type == IMAP_ARG_ATOM) {
+		str = str_ucase(IMAP_ARG_STR(arg));
 
-	*value = 0;
-	while (**p >= '0' && **p <= '9') {
-		prev = *value;
-		*value = *value * 10 + (**p - '0');
+		/* handle macros first */
+		if (strcmp(str, "ALL") == 0) {
+			if (!imap_fetch_init_handler(ctx, "FLAGS") ||
+			    !imap_fetch_init_handler(ctx, "INTERNALDATE") ||
+			    !imap_fetch_init_handler(ctx, "RFC822.SIZE") ||
+			    !imap_fetch_init_handler(ctx, "ENVELOPE"))
+				return FALSE;
+		} else if (strcmp(str, "FAST") == 0) {
+			if (!imap_fetch_init_handler(ctx, "FLAGS") ||
+			    !imap_fetch_init_handler(ctx, "INTERNALDATE") ||
+			    !imap_fetch_init_handler(ctx, "RFC822.SIZE"))
+				return FALSE;
+		} else if (strcmp(str, "FULL") == 0) {
+			if (!imap_fetch_init_handler(ctx, "FLAGS") ||
+			    !imap_fetch_init_handler(ctx, "INTERNALDATE") ||
+			    !imap_fetch_init_handler(ctx, "RFC822.SIZE") ||
+			    !imap_fetch_init_handler(ctx, "ENVELOPE") ||
+			    !imap_fetch_init_handler(ctx, "BODY"))
+				return FALSE;
+		} else {
+			if (!imap_fetch_init_handler(ctx, str))
+				return FALSE;
+		}
+	} else {
+		arg = IMAP_ARG_LIST(arg)->args;
+		while (arg->type == IMAP_ARG_ATOM) {
+			str = str_ucase(IMAP_ARG_STR(arg));
+			if (!imap_fetch_init_handler(ctx, str))
+				return FALSE;
+			arg++;
+		}
+		if (arg->type != IMAP_ARG_EOL) {
+			client_send_command_error(client,
+				"FETCH list contains non-atoms.");
+			return FALSE;
+		}
+	}
 
-		if (*value < prev)
+	if (client->cmd_uid) {
+		if (!imap_fetch_init_handler(ctx, "UID"))
 			return FALSE;
-
-		(*p)++;
 	}
 
 	return TRUE;
 }
 
-static int check_header_section(const char *section)
-{
-	/* HEADER, HEADER.FIELDS (list), HEADER.FIELDS.NOT (list) */
-	if (*section == '\0')
-		return TRUE;
-
-	if (strncmp(section, ".FIELDS", 7) != 0)
-		return FALSE;
-
-	section += 7;
-	if (strncmp(section, ".NOT", 4) == 0)
-		section += 4;
-
-	while (*section == ' ') section++;
-	if (*section++ != '(')
-		return FALSE;
-
-	while (*section != '\0' && *section != ')') {
-		if (*section == '(')
-			return FALSE;
-		section++;
-	}
-
-	if (*section++ != ')')
-		return FALSE;
-
-	if (*section != '\0')
-		return FALSE;
-	return TRUE;
-}
-
-static int check_section(struct client *client, const char *section,
-			 enum mail_fetch_field *fetch_data)
-{
-	if (*section == '\0') {
-		*fetch_data |= MAIL_FETCH_STREAM_HEADER |
-			MAIL_FETCH_STREAM_BODY;
-		return TRUE;
-	}
-
-	if (strcmp(section, "TEXT") == 0) {
-		*fetch_data |= MAIL_FETCH_STREAM_BODY;
-		return TRUE;
-	}
-
-	if (strncmp(section, "HEADER", 6) == 0) {
-		/* exact header matches could be cached */
-		if (strncmp(section, "HEADER.FIELDS ", 14) != 0)
-			*fetch_data |= MAIL_FETCH_STREAM_HEADER;
-
-		if (check_header_section(section+6))
-			return TRUE;
-	} else if (*section >= '0' && *section <= '9') {
-		*fetch_data |= MAIL_FETCH_STREAM_BODY |
-			MAIL_FETCH_MESSAGE_PARTS;
-
-		while ((*section >= '0' && *section <= '9') ||
-		       *section == '.') section++;
-
-		if (*section == '\0')
-			return TRUE;
-		if (strcmp(section, "MIME") == 0 ||
-		    strcmp(section, "TEXT") == 0)
-			return TRUE;
-
-		if (strncmp(section, "HEADER", 6) == 0 &&
-		    check_header_section(section+6))
-			return TRUE;
-	}
-
-	client_send_tagline(client, t_strconcat(
-		"BAD Invalid BODY[] section: ", section, NULL));
-	return FALSE;
-}
-
-/* BODY[] and BODY.PEEK[] items. item points to next character after '[' */
-static int parse_body_section(struct client *client, const char *item, int peek,
-			      enum mail_fetch_field *fetch_data,
-			      struct imap_fetch_body_data ***bodies)
-{
-	/* @UNSAFE */
-	struct imap_fetch_body_data *body;
-	uoff_t num;
-	char *p;
-
-	body = t_new(struct imap_fetch_body_data, 1);
-	body->peek = peek;
-
-	p = t_strdup_noconst(item);
-
-	/* read section */
-	body->section = p;
-	for (; *p != ']'; p++) {
-		if (*p == '\0') {
-			client_send_tagline(client, t_strconcat(
-				"BAD Missing ']' with ", item, NULL));
-			return FALSE;
-		}
-	}
-	*p++ = '\0';
-
-	if (!check_section(client, body->section, fetch_data))
-		return FALSE;
-
-	/* <start.end> */
-	body->skip = 0;
-	body->max_size = (uoff_t)-1;
-	if (*p != '<' && *p != '\0') {
-		client_send_tagline(client, t_strconcat(
-			"BAD Unexpected character after ']' with ",
-			item, NULL));
-	} else if (*p == '<') {
-		/* read start */
-		p++;
-
-		body->skip_set = TRUE;
-		if (!read_uoff_t(&p, &num) || num > OFF_T_MAX) {
-			/* wrapped */
-			client_send_tagline(client, t_strconcat(
-				"BAD Too big partial start with ", item, NULL));
-			return FALSE;
-		}
-		body->skip = num;
-
-		if (*p == '.') {
-			/* read end */
-			p++;
-			if (!read_uoff_t(&p, &num) || num > OFF_T_MAX) {
-				/* wrapped */
-				client_send_tagline(client, t_strconcat(
-					"BAD Too big partial end with ",
-					item, NULL));
-				return FALSE;
-			}
-
-                        body->max_size = num;
-		}
-
-		if (*p != '>') {
-			client_send_tagline(client, t_strconcat(
-				"BAD Invalid partial ", item, NULL));
-			return FALSE;
-		}
-	}
-
-	**bodies = body;
-	*bodies = &body->next;
-	return TRUE;
-}
-
-static int parse_arg(struct client *client, struct imap_arg *arg,
-		     enum mail_fetch_field *fetch_data,
-		     enum imap_fetch_field *imap_data,
-		     struct imap_fetch_body_data ***bodies)
+static void cmd_fetch_finish(struct client *client, int failed)
 {
-	char *item;
-
-	if (arg->type != IMAP_ARG_ATOM) {
-		client_send_command_error(client,
-					  "FETCH list contains non-atoms.");
-		return FALSE;
-	}
-
-	item = str_ucase(IMAP_ARG_STR(arg));
-
-	switch (*item) {
-	case 'A':
-		if (strcmp(item, "ALL") == 0) {
-			*fetch_data |= MAIL_FETCH_FLAGS |
-				MAIL_FETCH_RECEIVED_DATE |
-				MAIL_FETCH_SIZE |
-				MAIL_FETCH_IMAP_ENVELOPE;
-		} else
-			item = NULL;
-		break;
-	case 'B':
-		/* all start with BODY so skip it */
-		if (strncmp(item, "BODY", 4) != 0) {
-			item = NULL;
-			break;
-		}
-		item += 4;
-
-		if (*item == '\0') {
-			/* BODY */
-			*fetch_data |= MAIL_FETCH_IMAP_BODY;
-		} else if (*item == '[') {
-			/* BODY[...] */
-			if (!parse_body_section(client, item+1, FALSE,
-						fetch_data, bodies))
-				return FALSE;
-		} else if (strncmp(item, ".PEEK[", 6) == 0) {
-			/* BODY.PEEK[..] */
-			if (!parse_body_section(client, item+6, TRUE,
-						fetch_data, bodies))
-				return FALSE;
-		} else if (strcmp(item, "STRUCTURE") == 0) {
-			/* BODYSTRUCTURE */
-			*fetch_data |= MAIL_FETCH_IMAP_BODYSTRUCTURE;
-		} else
-			item = NULL;
-		break;
-	case 'E':
-		if (strcmp(item, "ENVELOPE") == 0)
-			*fetch_data |= MAIL_FETCH_IMAP_ENVELOPE;
-		else
-			item = NULL;
-		break;
-	case 'F':
-		if (strcmp(item, "FLAGS") == 0)
-			*fetch_data |= MAIL_FETCH_FLAGS;
-		else if (strcmp(item, "FAST") == 0) {
-			*fetch_data |= MAIL_FETCH_FLAGS |
-				MAIL_FETCH_RECEIVED_DATE |
-				MAIL_FETCH_SIZE;
-		} else if (strcmp(item, "FULL") == 0) {
-			*fetch_data |= MAIL_FETCH_FLAGS |
-				MAIL_FETCH_RECEIVED_DATE |
-				MAIL_FETCH_SIZE |
-				MAIL_FETCH_IMAP_ENVELOPE |
-				MAIL_FETCH_IMAP_BODY;
-		} else
-			item = NULL;
-		break;
-	case 'I':
-		if (strcmp(item, "INTERNALDATE") == 0)
-			*fetch_data |= MAIL_FETCH_RECEIVED_DATE;
-		else
-			item = NULL;
-		break;
-	case 'R':
-		/* all start with RFC822 so skip it */
-		if (strncmp(item, "RFC822", 6) != 0) {
-			item = NULL;
-			break;
-		}
-		item += 6;
-
-		if (*item == '\0') {
-			/* RFC822 */
-			*fetch_data |= MAIL_FETCH_STREAM_HEADER |
-				MAIL_FETCH_STREAM_BODY;
-			*imap_data |= IMAP_FETCH_RFC822;
-			break;
-		}
-
-		/* only items beginning with "RFC822." left */
-		if (*item != '.') {
-			item = NULL;
-			break;
-		}
-		item++;
-
-		if (strcmp(item, "HEADER") == 0) {
-			*fetch_data |= MAIL_FETCH_STREAM_HEADER;
-			*imap_data |= IMAP_FETCH_RFC822_HEADER;
-		} else if (strcmp(item, "TEXT") == 0) {
-			*fetch_data |= MAIL_FETCH_STREAM_BODY;
-			*imap_data |= IMAP_FETCH_RFC822_TEXT;
-		} else if (strcmp(item, "SIZE") == 0)
-			*fetch_data |= MAIL_FETCH_SIZE;
-		else
-			item = NULL;
-		break;
-	case 'U':
-		if (strcmp(item, "UID") == 0)
-			*imap_data |= IMAP_FETCH_UID;
-		else
-			item = NULL;
-		break;
-	default:
-		item = NULL;
-		break;
-	}
-
-	if (item == NULL) {
-		/* unknown item */
-		client_send_tagline(client, t_strconcat(
-			"BAD Invalid item ", IMAP_ARG_STR(arg), NULL));
-		return FALSE;
-	}
-
-	return TRUE;
-}
-
-int cmd_fetch(struct client *client)
-{
-	struct imap_arg *args, *listargs;
-	enum mail_fetch_field fetch_data;
-	enum imap_fetch_field imap_data;
-	struct imap_fetch_body_data *bodies, **bodies_p;
-	struct mail_search_arg *search_arg;
-	const char *messageset;
-	int ret;
-
-	if (!client_read_args(client, 2, 0, &args))
-		return FALSE;
-
-	if (!client_verify_open_mailbox(client))
-		return TRUE;
-
-	messageset = imap_arg_string(&args[0]);
-	if (messageset == NULL ||
-	    (args[1].type != IMAP_ARG_LIST && args[1].type != IMAP_ARG_ATOM)) {
-		client_send_command_error(client, "Invalid arguments.");
-		return TRUE;
-	}
-
-	/* parse items argument */
-	fetch_data = 0; imap_data = 0; bodies = NULL; bodies_p = &bodies;
-	if (args[1].type == IMAP_ARG_ATOM) {
-		if (!parse_arg(client, &args[1], &fetch_data,
-			       &imap_data, &bodies_p))
-			return TRUE;
-	} else {
-		listargs = IMAP_ARG_LIST(&args[1])->args;
-		while (listargs->type != IMAP_ARG_EOL) {
-			if (!parse_arg(client, listargs, &fetch_data,
-				       &imap_data, &bodies_p))
-				return TRUE;
-
-			listargs++;
-		}
-	}
-
-	if (client->cmd_uid)
-		imap_data |= IMAP_FETCH_UID;
-
-	search_arg = imap_search_get_arg(client, messageset, client->cmd_uid);
-	if (search_arg == NULL)
-		return TRUE;
-
-	ret = imap_fetch(client, fetch_data, imap_data, bodies, search_arg);
-	if (ret == 0) {
+	if (!failed) {
 		if ((client_workarounds &
 		     WORKAROUND_OE6_FETCH_NO_NEWMAIL) == 0) {
 			if (client->cmd_uid)
@@ -389,6 +93,67 @@
 			client_send_storage_error(client, storage);
 		}
 	}
+}
 
+static int cmd_fetch_continue(struct client *client)
+{
+        struct imap_fetch_context *ctx = client->cmd_context;
+	int ret;
+
+	if ((ret = imap_fetch(ctx)) == 0) {
+		/* unfinished */
+		return FALSE;
+	}
+	if (imap_fetch_deinit(ctx) < 0)
+		ret = -1;
+	cmd_fetch_finish(client, ret < 0);
 	return TRUE;
 }
+
+int cmd_fetch(struct client *client)
+{
+        struct imap_fetch_context *ctx;
+	struct imap_arg *args;
+	struct mail_search_arg *search_arg;
+	const char *messageset;
+	int ret;
+
+	if (!client_read_args(client, 2, 0, &args))
+		return FALSE;
+
+	if (!client_verify_open_mailbox(client))
+		return TRUE;
+
+	messageset = imap_arg_string(&args[0]);
+	if (messageset == NULL ||
+	    (args[1].type != IMAP_ARG_LIST && args[1].type != IMAP_ARG_ATOM)) {
+		client_send_command_error(client, "Invalid arguments.");
+		return TRUE;
+	}
+
+	search_arg = imap_search_get_arg(client, messageset, client->cmd_uid);
+	if (search_arg == NULL)
+		return TRUE;
+
+	ctx = imap_fetch_init(client);
+	if (ctx == NULL)
+		return TRUE;
+
+	if (!fetch_parse_args(client, ctx, &args[1])) {
+		imap_fetch_deinit(ctx);
+		return TRUE;
+	}
+
+	imap_fetch_begin(ctx, search_arg);
+	if ((ret = imap_fetch(ctx)) == 0) {
+		/* unfinished */
+		client->command_pending = TRUE;
+		client->cmd_func = cmd_fetch_continue;
+		client->cmd_context = ctx;
+		return FALSE;
+	}
+	if (imap_fetch_deinit(ctx) < 0)
+		ret = -1;
+	cmd_fetch_finish(client, ret < 0);
+	return TRUE;
+}

--- a/src/imap/cmd-idle.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/cmd-idle.c	Sun Aug 15 06:40:30 2004 +0300
@@ -26,8 +26,7 @@
 	}
 
 	io_remove(client->io);
-	client->io = io_add(i_stream_get_fd(client->input),
-			    IO_READ, _client_input, client);
+	client->io = NULL;
 
 	if (client->mailbox != NULL)
 		mailbox_notify_changes(client->mailbox, 0, NULL, NULL);
@@ -38,7 +37,7 @@
 	else
 		client_send_tagline(client, "BAD Expected DONE.");
 
-	o_stream_flush(client->output);
+	o_stream_uncork(client->output);
 
 	_client_reset_command(client);
 	client->bad_counter = 0;

--- a/src/imap/cmd-search.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/cmd-search.c	Sun Aug 15 06:40:30 2004 +0300
@@ -58,7 +58,6 @@
 	struct mail_search_arg *sargs;
 	struct imap_arg *args;
 	int args_count;
-	pool_t pool;
 	const char *error, *charset;
 
 	args_count = imap_parser_read_args(client->parser, 0, 0, &args);
@@ -91,9 +90,7 @@
 		charset = NULL;
 	}
 
-	pool = pool_alloconly_create("mail_search_args", 2048);
-
-	sargs = imap_search_args_build(pool, client->mailbox, args, &error);
+	sargs = imap_search_args_build(client->cmd_pool, client->mailbox, args, &error);
 	if (sargs == NULL) {
 		/* error in search arguments */
 		client_send_tagline(client, t_strconcat("NO ", error, NULL));
@@ -108,6 +105,5 @@
 					  mailbox_get_storage(client->mailbox));
 	}
 
-	pool_unref(pool);
 	return TRUE;
 }

--- a/src/imap/common.h	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/common.h	Sun Aug 15 06:40:30 2004 +0300
@@ -7,6 +7,15 @@
 /* Disconnect client after idling this many seconds */
 #define CLIENT_IDLE_TIMEOUT (60*30)
 
+/* If we can't send anything to client for this long, disconnect the client */
+#define CLIENT_OUTPUT_TIMEOUT (5*60)
+
+/* Stop buffering more data into output stream after this many bytes */
+#define CLIENT_OUTPUT_OPTIMAL_SIZE 2048
+
+/* Disconnect client when it sends too many bad commands in a row */
+#define CLIENT_MAX_BAD_COMMANDS 20
+
 /* RFC-2683 recommends at least 8000 bytes. Some clients however don't
    break large message sets to multiple commands, so we're pretty liberal
    by default. */

--- a/src/imap/imap-fetch-body-section.c	Sun Aug 15 05:54:47 2004 +0300
+++ /dev/null	Thu Jan 01 00:00:00 1970 +0000
@@ -1,554 +0,0 @@
-/* Copyright (C) 2002 Timo Sirainen */
-
-#include "common.h"
-#include "buffer.h"
-#include "istream.h"
-#include "ostream.h"
-#include "message-parser.h"
-#include "message-send.h"
-#include "mail-storage.h"
-#include "imap-fetch.h"
-
-#include <ctype.h>
-#include <unistd.h>
-
-/* For FETCH[HEADER.FIELDS*] we need to modify the header data before sending
-   it. We can either save it in memory and then send it, or we can parse it
-   twice, first calculating the size and then send it. This value specifies
-   the maximum amount of memory we allow to allocate before using
-   double-parsing. */
-#define MAX_HEADER_BUFFER_SIZE (32*1024)
-
-struct fetch_header_field_context {
-        struct imap_fetch_context *fetch_ctx;
-	struct mail *mail;
-
-	buffer_t *dest;
-	struct ostream *output;
-	uoff_t dest_size;
-
-	uoff_t skip, max_size;
-	const char *const *fields;
-	int (*match_func) (const char *const *, const char *, size_t);
-
-	unsigned int fix_nuls:1;
-};
-
-struct partial_cache {
-	unsigned int select_counter;
-	unsigned int uid;
-
-	uoff_t physical_start;
-	int cr_skipped;
-	struct message_size pos;
-};
-
-static struct partial_cache partial = { 0, 0, 0, 0, { 0, 0, 0 } };
-
-static int seek_partial(unsigned int select_counter, unsigned int uid,
-			struct partial_cache *partial, struct istream *stream,
-			uoff_t physical_start, uoff_t virtual_skip)
-{
-	int cr_skipped;
-
-	if (select_counter == partial->select_counter && uid == partial->uid &&
-	    physical_start == partial->physical_start &&
-	    virtual_skip >= partial->pos.virtual_size) {
-		/* we can use the cache */
-		virtual_skip -= partial->pos.virtual_size;
-	} else {
-		partial->select_counter = select_counter;
-		partial->uid = uid;
-		partial->physical_start = physical_start;
-		partial->cr_skipped = FALSE;
-		memset(&partial->pos, 0, sizeof(partial->pos));
-	}
-
-	i_stream_seek(stream, partial->physical_start +
-		      partial->pos.physical_size);
-	message_skip_virtual(stream, virtual_skip, &partial->pos,
-			     partial->cr_skipped, &cr_skipped);
-
-	partial->cr_skipped = FALSE;
-	return cr_skipped;
-}
-
-static uoff_t get_send_size(const struct imap_fetch_body_data *body,
-			    uoff_t max_size)
-{
-	uoff_t size;
-
-	if (body->skip >= max_size)
-		return 0;
-
-	size = max_size - body->skip;
-	return size <= body->max_size ? size : body->max_size;
-}
-
-static int fetch_data(struct imap_fetch_context *ctx,
-		      const struct imap_fetch_body_data *body,
-		      struct mail *mail, struct istream *input,
-		      uoff_t physical_start, const struct message_size *size)
-{
-	const char *str;
-	uoff_t send_size;
-	off_t ret;
-	int skip_cr, last_cr;
-
-	send_size = get_send_size(body, size->virtual_size);
-
-	str = t_strdup_printf("%s {%"PRIuUOFF_T"}\r\n", ctx->prefix, send_size);
-	if (o_stream_send_str(ctx->output, str) < 0)
-		return FALSE;
-
-	skip_cr = seek_partial(ctx->select_counter, mail->uid,
-			       &partial, input, physical_start, body->skip);
-
-	ret = message_send(ctx->output, input, size, skip_cr, send_size,
-			   &last_cr, !mail->has_no_nuls);
-	if (ret > 0) {
-		partial.cr_skipped = last_cr != 0;
-		partial.pos.physical_size =
-			input->v_offset - partial.physical_start;
-		partial.pos.virtual_size += ret;
-	}
-
-	if (ret != (off_t)send_size) {
-		/* Input stream gave less data then we expected. Two choices
-		   here: either we fill the missing data with spaces or we
-		   disconnect the client.
-
-		   We shouldn't really ever get here. One reason is if mail
-		   was deleted from NFS server while we were reading it.
-		   Another is some temporary disk error.
-
-		   If we filled the missing data the client could cache it,
-		   and if it was just a temporary error the message would be
-		   permanently left corrupted in client's local cache. So, we
-		   disconnect the client and hope that next try works. */
-		o_stream_close(ctx->output);
-		return FALSE;
-	}
-
-	return TRUE;
-}
-
-/* fetch BODY[] or BODY[TEXT] */
-static int fetch_body(struct imap_fetch_context *ctx,
-		      const struct imap_fetch_body_data *body,
-		      struct mail *mail, int fetch_header)
-{
-	struct message_size hdr_size, body_size;
-	struct istream *stream;
-
-	stream = mail->get_stream(mail, &hdr_size, &body_size);
-	if (stream == NULL)
-		return FALSE;
-
-	if (fetch_header)
-		message_size_add(&body_size, &hdr_size);
-
-	return fetch_data(ctx, body, mail, stream,
-			  fetch_header ? 0 : hdr_size.physical_size,
-			  &body_size);
-}
-
-static int header_match(const char *const *fields,
-			const char *name, size_t size)
-{
-	const char *name_start, *name_end, *field;
-
-	if (size == 0)
-		return FALSE;
-
-	name_start = name;
-	name_end = name + size;
-
-	for (; *fields != NULL; fields++) {
-		field = *fields;
-		if (*field == '\0')
-			continue;
-
-		for (name = name_start; name != name_end; name++) {
-			/* field has been uppercased long time ago while
-			   parsing FETCH command */
-			if (i_toupper(*name) != *field)
-				break;
-
-			field++;
-			if (*field == '\0') {
-				if (name+1 == name_end)
-					return TRUE;
-				break;
-			}
-		}
-	}
-
-	return FALSE;
-}
-
-static int header_match_not(const char *const *fields,
-			    const char *name, size_t size)
-{
-	return !header_match(fields, name, size);
-}
-
-static int header_match_mime(const char *const *fields __attr_unused__,
-			     const char *name, size_t size)
-{
-	if (strncasecmp(name, "Content-", 8) == 0)
-		return TRUE;
-
-	if (size == 12 && strcasecmp(name, "Mime-Version") == 0)
-		return TRUE;
-
-	return FALSE;
-}
-
-static int fetch_header_append(struct fetch_header_field_context *ctx,
-			       const void *data, size_t size)
-{
-	const unsigned char *str = data;
-	size_t i;
-
-	if (ctx->skip > 0) {
-		if (ctx->skip >= size) {
-			ctx->skip -= size;
-			return TRUE;
-		}
-
-		str += ctx->skip;
-		size -= ctx->skip;
-		ctx->skip = 0;
-	}
-
-	if (ctx->dest_size + size > ctx->max_size) {
-		i_assert(ctx->dest_size <= ctx->max_size);
-		size = ctx->max_size - ctx->dest_size;
-	}
-
-	ctx->dest_size += size;
-
-	if (ctx->fix_nuls && (ctx->dest != NULL || ctx->output != NULL)) {
-		for (i = 0; i < size; ) {
-			if (str[i] != 0) {
-				i++;
-				continue;
-			}
-
-			/* NUL found, change it to #128 */
-			if (ctx->dest != NULL) {
-				buffer_append(ctx->dest, str, i);
-				buffer_append(ctx->dest, "\x80", 1);
-			} else {
-				if (o_stream_send(ctx->output, str, i) < 0 ||
-				    o_stream_send(ctx->output, "\x80", 1) < 0)
-					return FALSE;
-			}
-
-			str += i+1;
-			size -= i+1;
-			i = 0;
-		}
-	}
-
-	if (ctx->dest != NULL)
-		buffer_append(ctx->dest, str, size);
-	else {
-		if (o_stream_send(ctx->output, str, size) < 0)
-			return FALSE;
-	}
-
-	return ctx->dest_size < ctx->max_size;
-}
-
-static int fetch_header_fields(struct istream *input, const char *section,
-			       struct fetch_header_field_context *ctx)
-{
-	struct message_header_parser_ctx *hdr_ctx;
-	struct message_header_line *hdr;
-
-	if (strncmp(section, "HEADER.FIELDS ", 14) == 0) {
-		ctx->fields = imap_fetch_get_body_fields(section + 14);
-		ctx->match_func = header_match;
-
-		if (ctx->fetch_ctx->body_fetch_from_cache) {
-			input = ctx->mail->
-				get_headers(ctx->mail,
-					    ctx->fetch_ctx->headers_ctx);
-			if (input == NULL)
-				return FALSE;
-		}
-	} else if (strncmp(section, "HEADER.FIELDS.NOT ", 18) == 0) {
-		ctx->fields = imap_fetch_get_body_fields(section + 18);
-		ctx->match_func = header_match_not;
-	} else if (strcmp(section, "MIME") == 0) {
-		/* Mime-Version + Content-* fields */
-		ctx->match_func = header_match_mime;
-	} else {
-		i_warning("BUG: Accepted invalid section from user: '%s'",
-			  section);
-		return FALSE;
-	}
-
-	ctx->dest_size = 0;
-
-	hdr_ctx = message_parse_header_init(input, NULL, FALSE);
-	while ((hdr = message_parse_header_next(hdr_ctx)) != NULL) {
-		/* see if we want this field. */
-		if (!ctx->match_func(ctx->fields, hdr->name, hdr->name_len))
-			continue;
-
-		if (!hdr->continued) {
-			if (!fetch_header_append(ctx, hdr->name, hdr->name_len))
-				break;
-			if (!fetch_header_append(ctx, hdr->middle,
-						 hdr->middle_len))
-				break;
-		}
-		if (!fetch_header_append(ctx, hdr->value, hdr->value_len))
-			break;
-		if (!hdr->no_newline) {
-			if (!fetch_header_append(ctx, "\r\n", 2))
-				break;
-		}
-	}
-
-	/* FIXME: We'll just always add the end of headers mark now.
-	   mail-storage should rather include it in the header stream..
-	   however, not much of a problem since all non-broken mails have it.
-
-	   Also, Netscape 4.x seems to require this or it won't show the
-	   mail.. So if we do make this as RFC says, we'll need to add
-	   netscape-workaround. */
-	(void)fetch_header_append(ctx, "\r\n", 2);
-
-	message_parse_header_deinit(hdr_ctx);
-
-	i_assert(ctx->dest_size <= ctx->max_size);
-	i_assert(ctx->dest == NULL ||
-		 buffer_get_used_size(ctx->dest) == ctx->dest_size);
-	return TRUE;
-}
-
-/* fetch wanted headers from given data */
-static int fetch_header_from(struct imap_fetch_context *ctx,
-			     struct istream *input,
-			     const struct message_size *size, struct mail *mail,
-			     const struct imap_fetch_body_data *body,
-			     const char *header_section)
-{
-	struct fetch_header_field_context hdr_ctx;
-	const char *str;
-	const void *data;
-	size_t data_size;
-	uoff_t start_offset;
-	int failed;
-
-	/* HEADER, MIME, HEADER.FIELDS (list), HEADER.FIELDS.NOT (list) */
-
-	if (strcmp(header_section, "HEADER") == 0) {
-		/* all headers */
-		return fetch_data(ctx, body, mail, input, 0, size);
-	}
-
-	/* partial headers - copy the wanted fields into memory, inserting
-	   missing CRs on the way. If the header is too large, calculate 
-	   the size first and then send the data directly to output stream. */
-
-	memset(&hdr_ctx, 0, sizeof(hdr_ctx));
-	hdr_ctx.mail = mail;
-	hdr_ctx.fetch_ctx = ctx;
-	hdr_ctx.skip = body->skip;
-	hdr_ctx.max_size = body->max_size;
-	hdr_ctx.fix_nuls = !mail->has_no_nuls;
-
-	failed = FALSE;
-	start_offset = input == NULL ? 0 : input->v_offset;
-
-	t_push();
-
-	/* first pass, we need at least the size */
-	if (!ctx->body_fetch_from_cache &&
-	    size->virtual_size > MAX_HEADER_BUFFER_SIZE &&
-	    body->max_size > MAX_HEADER_BUFFER_SIZE) {
-		if (!fetch_header_fields(input, header_section, &hdr_ctx))
-			failed = TRUE;
-
-		i_assert(hdr_ctx.dest_size <= size->virtual_size);
-	} else {
-		hdr_ctx.dest =
-			buffer_create_dynamic(pool_datastack_create(),
-					      I_MIN(size->virtual_size, 8192),
-					      (size_t)-1);
-		if (!fetch_header_fields(input, header_section, &hdr_ctx))
-			failed = TRUE;
-	}
-
-	if (!failed) {
-		str = t_strdup_printf("%s {%"PRIuUOFF_T"}\r\n",
-				      ctx->prefix, hdr_ctx.dest_size);
-		if (o_stream_send_str(ctx->output, str) < 0)
-			failed = TRUE;
-	}
-
-	if (!failed) {
-		if (hdr_ctx.dest == NULL) {
-			/* second pass, write the data to output stream */
-			uoff_t first_size = hdr_ctx.dest_size;
-
-			hdr_ctx.output = ctx->output;
-			i_stream_seek(input, start_offset);
-
-			if (!failed &&
-			    !fetch_header_fields(input, header_section,
-						 &hdr_ctx))
-				failed = TRUE;
-
-			i_assert(first_size == hdr_ctx.dest_size);
-		} else {
-			data = buffer_get_data(hdr_ctx.dest, &data_size);
-			if (o_stream_send(ctx->output, data, data_size) < 0)
-				failed = TRUE;
-		}
-	}
-
-	t_pop();
-	return !failed;
-}
-
-static int fetch_header(struct imap_fetch_context *ctx, struct mail *mail,
-			const struct imap_fetch_body_data *body)
-{
-	struct istream *stream;
-	struct message_size hdr_size;
-
-	if (ctx->body_fetch_from_cache) {
-		memset(&hdr_size, 0, sizeof(hdr_size));
-		stream = NULL;
-	} else {
-		stream = mail->get_stream(mail, &hdr_size, NULL);
-		if (stream == NULL)
-			return FALSE;
-	}
-
-	return fetch_header_from(ctx, stream, &hdr_size,
-				 mail, body, body->section);
-}
-
-/* Find message_part for section (eg. 1.3.4) */
-static int part_find(struct mail *mail, const struct imap_fetch_body_data *body,
-		     const struct message_part **part_r, const char **section)
-{
-	const struct message_part *part;
-	const char *path;
-	unsigned int num;
-
-	part = mail->get_parts(mail);
-	if (part == NULL)
-		return FALSE;
-
-	path = body->section;
-	while (*path >= '0' && *path <= '9' && part != NULL) {
-		/* get part number */
-		num = 0;
-		while (*path != '\0' && *path != '.') {
-			if (*path < '0' || *path > '9')
-				return FALSE;
-			num = num*10 + (*path - '0');
-			path++;
-		}
-
-		if (*path == '.')
-			path++;
-
-		if (part->flags & MESSAGE_PART_FLAG_MULTIPART) {
-			/* find the part */
-			part = part->children;
-			for (; num > 1 && part != NULL; num--)
-				part = part->next;
-		} else {
-			/* only 1 allowed with non-multipart messages */
-			if (num != 1)
-				part = NULL;
-		}
-
-		if (part != NULL &&
-		    (part->flags & MESSAGE_PART_FLAG_MESSAGE_RFC822) &&
-		    ((*path >= '0' && *path <= '9') ||
-		     strncmp(path, "HEADER", 6) == 0)) {
-			/* if remainder of path is a number or "HEADER",
-			   skip the message/rfc822 part */
-			part = part->children;
-		}
-	}
-
-	*part_r = part;
-	*section = path;
-	return TRUE;
-}
-
-static int fetch_part(struct imap_fetch_context *ctx, struct mail *mail,
-		      const struct imap_fetch_body_data *body)
-{
-	struct istream *stream;
-	const struct message_part *part;
-	const char *section;
-
-	if (!part_find(mail, body, &part, &section))
-		return FALSE;
-
-	if (part == NULL) {
-		/* part doesn't exist */
-		return o_stream_send_str(ctx->output, ctx->prefix) > 0 &&
-			o_stream_send_str(ctx->output, " NIL") > 0;
-	}
-
-	stream = mail->get_stream(mail, NULL, NULL);
-	if (stream == NULL)
-		return FALSE;
-
-	if (*section == '\0' || strcmp(section, "TEXT") == 0) {
-		return fetch_data(ctx, body, mail, stream,
-				  part->physical_pos +
-				  part->header_size.physical_size,
-				  &part->body_size);
-	}
-
-	if (strncmp(section, "HEADER", 6) == 0 ||
-	    strcmp(section, "MIME") == 0) {
-		i_stream_seek(stream, part->physical_pos);
-		return fetch_header_from(ctx, stream, &part->header_size,
-					 mail, body, section);
-	}
-
-	i_warning("BUG: Accepted invalid section from user: '%s'",
-		  body->section);
-	return FALSE;
-}
-
-int imap_fetch_body_section(struct imap_fetch_context *ctx,
-			    const struct imap_fetch_body_data *body,
-			    struct mail *mail)
-{
-	ctx->prefix = !body->skip_set ?
-		t_strdup_printf(" BODY[%s]", body->section) :
-		t_strdup_printf(" BODY[%s]<%"PRIuUOFF_T">",
-				body->section, body->skip);
-	if (ctx->first) {
-		ctx->prefix++; ctx->first = FALSE;
-	}
-
-	if (*body->section == '\0')
-		return fetch_body(ctx, body, mail, TRUE);
-	if (strcmp(body->section, "TEXT") == 0)
-		return fetch_body(ctx, body, mail, FALSE);
-	if (strncmp(body->section, "HEADER", 6) == 0)
-		return fetch_header(ctx, mail, body);
-	if (*body->section >= '0' && *body->section <= '9')
-		return fetch_part(ctx, mail, body);
-
-	i_warning("BUG: Accepted invalid section from user: '%s'",
-		  body->section);
-	return FALSE;
-}

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/src/imap/imap-fetch-body.c	Sun Aug 15 06:40:30 2004 +0300
@@ -0,0 +1,838 @@
+/* Copyright (C) 2002-2004 Timo Sirainen */
+
+#include "common.h"
+#include "buffer.h"
+#include "str.h"
+#include "istream.h"
+#include "ostream.h"
+#include "istream-header-filter.h"
+#include "message-parser.h"
+#include "message-send.h"
+#include "mail-storage.h"
+#include "imap-fetch.h"
+
+#include <stdlib.h>
+#include <ctype.h>
+#include <unistd.h>
+
+struct imap_fetch_body_data {
+	struct imap_fetch_body_data *next;
+
+        struct mailbox_header_lookup_ctx *header_ctx;
+	const char *section; /* NOTE: always uppercased */
+	uoff_t skip, max_size; /* if you don't want max_size,
+	                          set it to (uoff_t)-1 */
+	unsigned int skip_set:1;
+	unsigned int peek:1;
+};
+
+struct partial_cache {
+	unsigned int select_counter;
+	unsigned int uid;
+
+	uoff_t physical_start;
+	int cr_skipped;
+	struct message_size pos;
+};
+
+static struct partial_cache partial = { 0, 0, 0, 0, { 0, 0, 0 } };
+
+static const char *const *
+imap_fetch_get_body_fields(const char *fields, size_t *count_r)
+{
+	const char **arr, **p;
+	size_t count;
+
+	i_assert(*fields == '(');
+
+	arr = t_strsplit_spaces(t_strcut(fields+1, ')'), " ");
+	for (count = 0, p = arr; *p != NULL; p++)
+		count++;
+
+	qsort(arr, count, sizeof(*arr), strcasecmp_p);
+
+	*count_r = count;
+	return arr;
+}
+
+static int seek_partial(unsigned int select_counter, unsigned int uid,
+			struct partial_cache *partial, struct istream *stream,
+			uoff_t virtual_skip)
+{
+	int cr_skipped;
+
+	if (select_counter == partial->select_counter && uid == partial->uid &&
+	    stream->v_offset == partial->physical_start &&
+	    virtual_skip >= partial->pos.virtual_size) {
+		/* we can use the cache */
+		virtual_skip -= partial->pos.virtual_size;
+	} else {
+		partial->select_counter = select_counter;
+		partial->uid = uid;
+		partial->physical_start = stream->v_offset;
+		partial->cr_skipped = FALSE;
+		memset(&partial->pos, 0, sizeof(partial->pos));
+	}
+
+	i_stream_seek(stream, partial->physical_start +
+		      partial->pos.physical_size);
+	message_skip_virtual(stream, virtual_skip, &partial->pos,
+			     partial->cr_skipped, &cr_skipped);
+
+	partial->cr_skipped = FALSE;
+	return cr_skipped;
+}
+
+static uoff_t get_send_size(const struct imap_fetch_body_data *body,
+			    uoff_t max_size)
+{
+	uoff_t size;
+
+	if (body->skip >= max_size)
+		return 0;
+
+	size = max_size - body->skip;
+	return size <= body->max_size ? size : body->max_size;
+}
+
+static string_t *get_prefix(struct imap_fetch_context *ctx,
+			    const struct imap_fetch_body_data *body,
+			    uoff_t size)
+{
+	string_t *str;
+
+	str = t_str_new(128);
+	if (ctx->first)
+		ctx->first = FALSE;
+	else
+		str_append_c(str, ' ');
+
+	str_printfa(str, "BODY[%s]", body->section);
+	if (body->skip_set)
+		str_printfa(str, "<%"PRIuUOFF_T">", body->skip);
+
+	if (size != (uoff_t)-1)
+		str_printfa(str, " {%"PRIuUOFF_T"}\r\n", size);
+	else
+		str_append(str, " NIL");
+	return str;
+}
+
+static off_t imap_fetch_send(struct ostream *output, struct istream *input,
+			     int cr_skipped, uoff_t virtual_size, int *last_cr)
+{
+	const unsigned char *msg;
+	size_t i, size;
+	uoff_t vsize_left, sent;
+	off_t ret;
+	unsigned char add;
+	int blocks = FALSE;
+
+	/* go through the message data and insert CRs where needed.  */
+	sent = 0; vsize_left = virtual_size;
+	while (vsize_left > 0 && !blocks &&
+	       i_stream_read_data(input, &msg, &size, 0) > 0) {
+		add = '\0';
+		for (i = 0; i < size && vsize_left > 0; i++) {
+			vsize_left--;
+
+			if (msg[i] == '\n') {
+				if ((i > 0 && msg[i-1] != '\r') ||
+				    (i == 0 && !cr_skipped)) {
+					/* missing CR */
+					add = '\r';
+					break;
+				}
+			} else if (msg[i] == '\0') {
+				add = 128;
+				break;
+			}
+		}
+
+		if ((ret = o_stream_send(output, msg, i)) < 0)
+			return -1;
+		if (ret < i) {
+			add = '\0';
+			blocks = TRUE;
+		}
+		i_stream_skip(input, ret);
+		sent += ret;
+
+		cr_skipped = ret > 0 && msg[ret-1] == '\r';
+		if (add != '\0') {
+			if ((ret = o_stream_send(output, &add, 1)) < 0)
+				return -1;
+			if (ret == 0)
+				blocks = TRUE;
+			else {
+				sent++;
+				cr_skipped = add == '\r';
+				if (add == 128)
+					i_stream_skip(input, 1);
+			}
+		}
+	}
+
+	if ((uoff_t)sent != virtual_size && !blocks) {
+		/* Input stream gave less data then we expected. Two choices
+		   here: either we fill the missing data with spaces or we
+		   disconnect the client.
+
+		   We shouldn't really ever get here. One reason is if mail
+		   was deleted from NFS server while we were reading it.
+		   Another is some temporary disk error.
+
+		   If we filled the missing data the client could cache it,
+		   and if it was just a temporary error the message would be
+		   permanently left corrupted in client's local cache. So, we
+		   disconnect the client and hope that next try works. */
+		o_stream_close(output);
+		return -1;
+	}
+
+	*last_cr = cr_skipped;
+	return sent;
+}
+
+static int fetch_stream_send(struct imap_fetch_context *ctx)
+{
+	off_t ret;
+
+	o_stream_set_max_buffer_size(ctx->client->output, 0);
+	ret = imap_fetch_send(ctx->client->output, ctx->cur_input,
+			      ctx->skip_cr, ctx->cur_size - ctx->cur_offset,
+			      &ctx->skip_cr);
+	o_stream_set_max_buffer_size(ctx->client->output, (size_t)-1);
+
+	if (ret < 0)
+		return -1;
+
+	ctx->cur_offset += ret;
+	if (ctx->update_partial) {
+		partial.cr_skipped = ctx->skip_cr != 0;
+		partial.pos.physical_size =
+			ctx->cur_input->v_offset - partial.physical_start;
+		partial.pos.virtual_size += ret;
+	}
+
+	return ctx->cur_offset == ctx->cur_size;
+}
+
+static int fetch_stream_send_direct(struct imap_fetch_context *ctx)
+{
+	off_t ret;
+
+	o_stream_set_max_buffer_size(ctx->client->output, 0);
+	ret = o_stream_send_istream(ctx->client->output, ctx->cur_input);
+	o_stream_set_max_buffer_size(ctx->client->output, (size_t)-1);
+
+	if (ret < 0)
+		return -1;
+
+	ctx->cur_offset += ret;
+	return ctx->cur_offset == ctx->cur_size;
+}
+
+static int fetch_stream(struct imap_fetch_context *ctx,
+			const struct message_size *size)
+{
+	struct istream *input;
+
+	if (size->physical_size == size->virtual_size &&
+	    ctx->cur_mail->has_no_nuls) {
+		/* no need to kludge with CRs, we can use sendfile() */
+		input = i_stream_create_limit(default_pool, ctx->cur_input,
+					      ctx->cur_input->v_offset,
+					      ctx->cur_size);
+		i_stream_unref(ctx->cur_input);
+		ctx->cur_input = input;
+
+		ctx->cont_handler = fetch_stream_send_direct;
+	} else {
+                ctx->cont_handler = fetch_stream_send;
+	}
+
+	return ctx->cont_handler(ctx);
+}
+
+static int fetch_data(struct imap_fetch_context *ctx,
+		      const struct imap_fetch_body_data *body,
+		      const struct message_size *size)
+{
+	string_t *str;
+
+	ctx->cur_size = get_send_size(body, size->virtual_size);
+
+	str = get_prefix(ctx, body, ctx->cur_size);
+	if (o_stream_send(ctx->client->output,
+			  str_data(str), str_len(str)) < 0)
+		return -1;
+
+	if (!ctx->update_partial) {
+		message_skip_virtual(ctx->cur_input, body->skip, NULL, FALSE,
+				     &ctx->skip_cr);
+	} else {
+		ctx->skip_cr =
+			seek_partial(ctx->select_counter, ctx->cur_mail->uid,
+				     &partial, ctx->cur_input, body->skip);
+	}
+
+	return fetch_stream(ctx, size);
+}
+
+static int fetch_body(struct imap_fetch_context *ctx, struct mail *mail,
+		      void *context)
+{
+	const struct imap_fetch_body_data *body = context;
+	const struct message_size *fetch_size;
+	struct message_size hdr_size, body_size;
+
+	ctx->cur_input = mail->get_stream(mail, &hdr_size, &body_size);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = TRUE;
+
+	switch (body->section[0]) {
+	case '\0':
+		/* BODY[] - fetch everything */
+		message_size_add(&body_size, &hdr_size);
+                fetch_size = &body_size;
+		break;
+	case 'H':
+		/* BODY[HEADER] - fetch only header */
+                fetch_size = &hdr_size;
+		break;
+	case 'T':
+		/* BODY[TEXT] - skip header */
+		i_stream_skip(ctx->cur_input, hdr_size.physical_size);
+                fetch_size = &body_size;
+		break;
+	default:
+		i_unreached();
+	}
+
+	return fetch_data(ctx, body, fetch_size);
+}
+
+static void header_filter_mime(struct message_header_line *hdr,
+			       int *matched, void *context __attr_unused__)
+{
+	if (hdr == NULL)
+		return;
+
+	*matched = strncasecmp(hdr->name, "Content-", 8) == 0 ||
+		strcasecmp(hdr->name, "Mime-Version") == 0;
+}
+
+static int fetch_header_partial_from(struct imap_fetch_context *ctx,
+				     const struct imap_fetch_body_data *body,
+				     const char *header_section)
+{
+	const char *const *fields;
+	struct message_size msg_size;
+	struct istream *input;
+	size_t size, fields_count;
+
+	/* MIME, HEADER.FIELDS (list), HEADER.FIELDS.NOT (list) */
+
+	if (strncmp(header_section, "HEADER.FIELDS ", 14) == 0) {
+		fields = imap_fetch_get_body_fields(header_section + 14,
+						    &fields_count);
+		input = i_stream_create_header_filter(ctx->client->cmd_pool,
+						      ctx->cur_input, FALSE,
+						      fields, fields_count,
+						      NULL, NULL);
+	} else if (strncmp(header_section, "HEADER.FIELDS.NOT ", 18) == 0) {
+		fields = imap_fetch_get_body_fields(header_section + 18,
+						    &fields_count);
+		input = i_stream_create_header_filter(ctx->client->cmd_pool,
+						      ctx->cur_input, TRUE,
+						      fields, fields_count,
+						      NULL, NULL);
+	} else if (strcmp(header_section, "MIME") == 0) {
+		/* Mime-Version + Content-* fields */
+		input = i_stream_create_header_filter(ctx->client->cmd_pool,
+						      ctx->cur_input, FALSE,
+						      NULL, 0,
+						      header_filter_mime, NULL);
+	} else {
+		i_error("BUG: Accepted invalid section from user: '%s'",
+			header_section);
+		return -1;
+	}
+
+	i_stream_unref(ctx->cur_input);
+	ctx->cur_input = input;
+	ctx->update_partial = FALSE;
+
+	/* FIXME: We'll just always add the end of headers mark now.
+	   mail-storage should rather include it in the header stream..
+	   however, not much of a problem since all non-broken mails have it.
+
+	   Also, Netscape 4.x seems to require this or it won't show the
+	   mail.. So if we do make this as RFC says, we'll need to add
+	   netscape-workaround. */
+
+	// FIXME: we rely on the current behavior of header filter..
+	(void)i_stream_get_data(ctx->cur_input, &size);
+	memset(&msg_size, 0, sizeof(msg_size));
+	msg_size.physical_size = msg_size.virtual_size = size;
+	return fetch_data(ctx, body, &msg_size);
+}
+
+static int fetch_body_header_partial(struct imap_fetch_context *ctx,
+				     struct mail *mail, void *context)
+{
+	const struct imap_fetch_body_data *body = context;
+
+	ctx->cur_input = mail->get_stream(mail, NULL, NULL);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = FALSE;
+
+	return fetch_header_partial_from(ctx, body, body->section);
+}
+
+static int fetch_body_header_fields(struct imap_fetch_context *ctx,
+				    struct mail *mail, void *context)
+{
+	const struct imap_fetch_body_data *body = context;
+	struct message_size size;
+
+	ctx->cur_input = mail->get_headers(mail, body->header_ctx);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = FALSE;
+
+	message_get_body_size(ctx->cur_input, &size, NULL);
+	i_stream_seek(ctx->cur_input, 0);
+
+	return fetch_data(ctx, body, &size);
+}
+
+/* Find message_part for section (eg. 1.3.4) */
+static int part_find(struct mail *mail, const struct imap_fetch_body_data *body,
+		     const struct message_part **part_r, const char **section)
+{
+	const struct message_part *part;
+	const char *path;
+	unsigned int num;
+
+	part = mail->get_parts(mail);
+	if (part == NULL)
+		return -1;
+
+	path = body->section;
+	while (*path >= '0' && *path <= '9' && part != NULL) {
+		/* get part number */
+		num = 0;
+		while (*path != '\0' && *path != '.') {
+			if (*path < '0' || *path > '9')
+				return FALSE;
+			num = num*10 + (*path - '0');
+			path++;
+		}
+
+		if (*path == '.')
+			path++;
+
+		if (part->flags & MESSAGE_PART_FLAG_MULTIPART) {
+			/* find the part */
+			part = part->children;
+			for (; num > 1 && part != NULL; num--)
+				part = part->next;
+		} else {
+			/* only 1 allowed with non-multipart messages */
+			if (num != 1)
+				part = NULL;
+		}
+
+		if (part != NULL &&
+		    (part->flags & MESSAGE_PART_FLAG_MESSAGE_RFC822) &&
+		    ((*path >= '0' && *path <= '9') ||
+		     strncmp(path, "HEADER", 6) == 0)) {
+			/* if remainder of path is a number or "HEADER",
+			   skip the message/rfc822 part */
+			part = part->children;
+		}
+	}
+
+	*part_r = part;
+	*section = path;
+	return 0;
+}
+
+static int fetch_body_mime(struct imap_fetch_context *ctx, struct mail *mail,
+			   void *context)
+{
+	const struct imap_fetch_body_data *body = context;
+	const struct message_part *part;
+	const char *section;
+
+	if (part_find(mail, body, &part, &section) < 0)
+		return -1;
+
+	if (part == NULL) {
+		/* part doesn't exist */
+		string_t *str = get_prefix(ctx, body, (uoff_t)-1);
+		if (o_stream_send(ctx->client->output,
+				  str_data(str), str_len(str)) < 0)
+			return -1;
+		return 1;
+	}
+
+	ctx->cur_input = mail->get_stream(mail, NULL, NULL);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = TRUE;
+
+	if (*section == '\0' || strcmp(section, "TEXT") == 0) {
+		i_stream_seek(ctx->cur_input, part->physical_pos +
+			      part->header_size.physical_size);
+		return fetch_data(ctx, body, &part->body_size);
+	}
+
+	if (strcmp(section, "HEADER") == 0) {
+		/* all headers */
+		return fetch_data(ctx, body, &part->header_size);
+	}
+
+	if (strncmp(section, "HEADER", 6) == 0 ||
+	    strcmp(section, "MIME") == 0) {
+		i_stream_seek(ctx->cur_input, part->physical_pos);
+		return fetch_header_partial_from(ctx, body, section);
+	}
+
+	i_error("BUG: Accepted invalid section from user: '%s'", body->section);
+	return 1;
+}
+
+static int fetch_body_header_fields_check(const char *section)
+{
+	if (*section++ != '(')
+		return FALSE;
+
+	while (*section != '\0' && *section != ')') {
+		if (*section == '(')
+			return FALSE;
+		section++;
+	}
+
+	if (*section++ != ')')
+		return FALSE;
+
+	if (*section != '\0')
+		return FALSE;
+	return TRUE;
+}
+
+static int fetch_body_header_fields_init(struct imap_fetch_context *ctx,
+					 struct imap_fetch_body_data *body,
+					 const char *section)
+{
+	const char *const *headers, *const *arr;
+	size_t count;
+
+	if (!fetch_body_header_fields_check(section))
+		return FALSE;
+
+	if ((ctx->fetch_data & (MAIL_FETCH_STREAM_HEADER |
+				MAIL_FETCH_STREAM_BODY)) != 0) {
+		/* we'll need to open the file anyway, don't try to get the
+		   headers from cache. */
+		imap_fetch_add_handler(ctx, fetch_body_header_partial, body);
+		return TRUE;
+	}
+
+	t_push();
+        headers = imap_fetch_get_body_fields(section, &count);
+
+	for (arr = headers; *arr != NULL; arr++) {
+		char *hdr = p_strdup(ctx->client->cmd_pool, *arr);
+		buffer_append(ctx->all_headers_buf, &hdr, sizeof(hdr));
+	}
+
+	body->header_ctx = mailbox_header_lookup_init(ctx->box, headers);
+	imap_fetch_add_handler(ctx, fetch_body_header_fields, body);
+	t_pop();
+	return TRUE;
+}
+
+static int fetch_body_section_name_init(struct imap_fetch_context *ctx,
+					struct imap_fetch_body_data *body)
+{
+	const char *section = body->section;
+
+	if (*section == '\0') {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_HEADER |
+			MAIL_FETCH_STREAM_BODY;
+		imap_fetch_add_handler(ctx, fetch_body, body);
+		return TRUE;
+	}
+
+	if (strcmp(section, "TEXT") == 0) {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_BODY;
+		imap_fetch_add_handler(ctx, fetch_body, body);
+		return TRUE;
+	}
+
+	if (strncmp(section, "HEADER", 6) == 0) {
+		/* exact header matches could be cached */
+		if (section[6] == '\0') {
+			ctx->fetch_data |= MAIL_FETCH_STREAM_HEADER;
+			imap_fetch_add_handler(ctx, fetch_body, body);
+			return TRUE;
+		}
+
+		if (strncmp(section, "HEADER.FIELDS ", 14) == 0 &&
+		    fetch_body_header_fields_init(ctx, body, section+14))
+			return TRUE;
+
+		if (strncmp(section, "HEADER.FIELDS.NOT ", 18) == 0 &&
+		    fetch_body_header_fields_check(section+18)) {
+			imap_fetch_add_handler(ctx, fetch_body_header_partial,
+					       body);
+			return TRUE;
+		}
+	} else if (*section >= '0' && *section <= '9') {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_BODY |
+			MAIL_FETCH_MESSAGE_PARTS;
+
+		while ((*section >= '0' && *section <= '9') ||
+		       *section == '.') section++;
+
+		if (*section == '\0' ||
+		    strcmp(section, "MIME") == 0 ||
+		    strcmp(section, "TEXT") == 0 ||
+		    strcmp(section, "HEADER") == 0 ||
+		    (strncmp(section, "HEADER.FIELDS ", 14) == 0 &&
+		     fetch_body_header_fields_check(section+14)) ||
+		    (strncmp(section, "HEADER.FIELDS.NOT ", 18) == 0 &&
+		     fetch_body_header_fields_check(section+18))) {
+			imap_fetch_add_handler(ctx, fetch_body_mime, body);
+			return TRUE;
+		}
+	}
+
+	client_send_command_error(ctx->client,
+		"Invalid BODY[..] parameter: Unknown or broken section");
+	return FALSE;
+}
+
+/* Parse next digits in string into integer. Returns FALSE if the integer
+   becomes too big and wraps. */
+static int read_uoff_t(const char **p, uoff_t *value)
+{
+	uoff_t prev;
+
+	*value = 0;
+	while (**p >= '0' && **p <= '9') {
+		prev = *value;
+		*value = *value * 10 + (**p - '0');
+
+		if (*value < prev)
+			return FALSE;
+
+		(*p)++;
+	}
+
+	return TRUE;
+}
+
+int fetch_body_section_init(struct imap_fetch_context *ctx, const char *arg)
+{
+	struct imap_fetch_body_data *body;
+	const char *p = arg + 4;
+
+	body = p_new(ctx->client->cmd_pool, struct imap_fetch_body_data, 1);
+	body->max_size = (uoff_t)-1;
+
+	if (strncmp(p, ".PEEK", 5) == 0) {
+		body->peek = TRUE;
+		p += 5;
+	} else {
+		ctx->flags_update_seen = TRUE;
+	}
+
+	if (*p != '[') {
+		client_send_command_error(ctx->client,
+			"Invalid BODY[..] parameter: Missing '['");
+		return FALSE;
+	}
+
+	body->section = p+1;
+	p = strchr(body->section, ']');
+	if (p == NULL) {
+		client_send_command_error(ctx->client,
+			"Invalid BODY[..] parameter: Missing ']'");
+		return FALSE;
+	}
+	body->section = p_strdup_until(ctx->client->cmd_pool, body->section, p);
+
+	if (*++p == '<') {
+		/* <start.end> */
+		p++;
+		body->skip_set = TRUE;
+
+		if (!read_uoff_t(&p, &body->skip) || body->skip > OFF_T_MAX) {
+			/* wrapped */
+			client_send_command_error(ctx->client,
+				"Invalid BODY[..] parameter: "
+				"Too big partial start");
+			return FALSE;
+		}
+
+		if (*p == '.') {
+			p++;
+			if (!read_uoff_t(&p, &body->max_size) ||
+			    body->max_size > OFF_T_MAX) {
+				/* wrapped */
+				client_send_command_error(ctx->client,
+					"Invalid BODY[..] parameter: "
+					"Too big partial end");
+				return FALSE;
+			}
+		}
+
+		if (*p != '>') {
+			client_send_command_error(ctx->client,
+				"Invalid BODY[..] parameter: Missing '>'");
+			return FALSE;
+		}
+	}
+
+	return fetch_body_section_name_init(ctx, body);
+}
+
+static int fetch_rfc822_size(struct imap_fetch_context *ctx, struct mail *mail,
+			     void *context __attr_unused__)
+{
+	uoff_t size;
+
+	size = mail->get_size(mail);
+	if (size == (uoff_t)-1)
+		return -1;
+
+	str_printfa(ctx->cur_str, "RFC822.SIZE %"PRIuUOFF_T" ", size);
+	return 1;
+}
+
+static int fetch_rfc822(struct imap_fetch_context *ctx, struct mail *mail,
+			void *context __attr_unused__)
+{
+	struct message_size hdr_size, body_size;
+	const char *str;
+
+	ctx->cur_input = mail->get_stream(mail, &hdr_size, &body_size);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = FALSE;
+
+	message_size_add(&body_size, &hdr_size);
+
+	if (ctx->cur_offset == 0) {
+		str = t_strdup_printf(" RFC822 {%"PRIuUOFF_T"}\r\n",
+				      body_size.virtual_size);
+		if (ctx->first) {
+			str++; ctx->first = FALSE;
+		}
+		if (o_stream_send_str(ctx->client->output, str) < 0)
+			return -1;
+	}
+
+        ctx->cur_size = body_size.virtual_size;
+	return fetch_stream(ctx, &body_size);
+}
+
+static int fetch_rfc822_header(struct imap_fetch_context *ctx,
+			       struct mail *mail, void *context __attr_unused__)
+{
+	struct message_size hdr_size;
+	const char *str;
+
+	ctx->cur_input = mail->get_stream(mail, &hdr_size, NULL);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = FALSE;
+
+	str = t_strdup_printf(" RFC822.HEADER {%"PRIuUOFF_T"}\r\n",
+			      hdr_size.virtual_size);
+	if (ctx->first) {
+		str++; ctx->first = FALSE;
+	}
+	if (o_stream_send_str(ctx->client->output, str) < 0)
+		return -1;
+
+        ctx->cur_size = hdr_size.virtual_size;
+	return fetch_stream(ctx, &hdr_size);
+}
+
+static int fetch_rfc822_text(struct imap_fetch_context *ctx, struct mail *mail,
+			     void *context __attr_unused__)
+{
+	struct message_size hdr_size, body_size;
+	const char *str;
+
+	ctx->cur_input = mail->get_stream(mail, &hdr_size, &body_size);
+	if (ctx->cur_input == NULL)
+		return -1;
+
+	i_stream_ref(ctx->cur_input);
+	ctx->update_partial = FALSE;
+
+	str = t_strdup_printf(" RFC822.TEXT {%"PRIuUOFF_T"}\r\n",
+			      body_size.virtual_size);
+	if (ctx->first) {
+		str++; ctx->first = FALSE;
+	}
+	if (o_stream_send_str(ctx->client->output, str) < 0)
+		return -1;
+
+	i_stream_seek(ctx->cur_input, hdr_size.physical_size);
+        ctx->cur_size = body_size.virtual_size;
+	return fetch_stream(ctx, &body_size);
+}
+
+int fetch_rfc822_init(struct imap_fetch_context *ctx, const char *arg)
+{
+	if (arg[6] == '\0') {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_HEADER |
+			MAIL_FETCH_STREAM_BODY;
+		ctx->flags_update_seen = TRUE;
+		imap_fetch_add_handler(ctx, fetch_rfc822, NULL);
+		return TRUE;
+	}
+
+	if (strcmp(arg+6, ".SIZE") == 0) {
+		ctx->fetch_data |= MAIL_FETCH_SIZE;
+		imap_fetch_add_handler(ctx, fetch_rfc822_size, NULL);
+		return TRUE;
+	}
+	if (strcmp(arg+6, ".HEADER") == 0) {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_HEADER;
+		imap_fetch_add_handler(ctx, fetch_rfc822_header, NULL);
+		return TRUE;
+	}
+	if (strcmp(arg+6, ".TEXT") == 0) {
+		ctx->fetch_data |= MAIL_FETCH_STREAM_BODY;
+		ctx->flags_update_seen = TRUE;
+		imap_fetch_add_handler(ctx, fetch_rfc822_text, NULL);
+		return TRUE;
+	}
+
+	client_send_command_error(ctx->client, t_strconcat(
+		"Unknown parameter ", arg, NULL));
+	return FALSE;
+}

--- a/src/imap/imap-fetch.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/imap-fetch.c	Sun Aug 15 06:40:30 2004 +0300
@@ -1,4 +1,4 @@
-/* Copyright (C) 2002 Timo Sirainen */
+/* Copyright (C) 2002-2004 Timo Sirainen */
 
 #include "common.h"
 #include "buffer.h"
@@ -12,419 +12,445 @@
 #include "imap-fetch.h"
 #include "imap-util.h"
 
-#include <unistd.h>
+#include <stdlib.h>
 
-const char *const *imap_fetch_get_body_fields(const char *fields)
+const struct imap_fetch_handler default_handlers[7];
+static buffer_t *fetch_handlers = NULL;
+
+static int imap_fetch_handler_cmp(const void *p1, const void *p2)
 {
-	const char **field_list, **field, **dest;
+        const struct imap_fetch_handler *h1 = p1, *h2 = p2;
+
+	return strcmp(h1->name, h2->name);
+}
+
+void imap_fetch_handlers_register(const struct imap_fetch_handler *handlers,
+				  size_t count)
+{
+	void *data;
+	size_t size;
 
-	while (*fields == ' ')
-		fields++;
-	if (*fields == '(')
-		fields++;
+	if (fetch_handlers == NULL) {
+		fetch_handlers = buffer_create_dynamic(default_pool,
+						       128, (size_t)-1);
+	}
+	buffer_append(fetch_handlers, handlers, sizeof(*handlers) * count);
+
+	data = buffer_get_modifyable_data(fetch_handlers, &size);
+	qsort(data, size / sizeof(*handlers), sizeof(*handlers),
+	      imap_fetch_handler_cmp);
+}
 
-	field_list = t_strsplit_spaces(t_strcut(fields, ')'), " ");
+static int imap_fetch_handler_bsearch(const void *name_p, const void *handler_p)
+{
+	const char *name = name_p;
+        const struct imap_fetch_handler *h = handler_p;
+	int i;
+
+	for (i = 0; h->name[i] != '\0'; i++) {
+		if (h->name[i] != name[i]) {
+			if (name[i] < 'A' || name[i] >= 'Z')
+				return -1;
+			return name[i] - h->name[i];
+		}
+	}
 
-	/* array ends at ")" element */
-	for (field = dest = field_list; *field != NULL; field++) {
-		*dest = *field;
-		dest++;
-	}
-	*dest = NULL;
+	return name[i] < 'A' || name[i] >= 'Z' ? 0 : -1;
+}
+
+int imap_fetch_init_handler(struct imap_fetch_context *ctx, const char *arg)
+{
+	const struct imap_fetch_handler *handler;
 
-	return field_list;
+	handler = bsearch(arg, fetch_handlers->data,
+			  fetch_handlers->used /
+			  sizeof(struct imap_fetch_handler),
+                          sizeof(struct imap_fetch_handler),
+			  imap_fetch_handler_bsearch);
+	if (handler == NULL)
+		i_panic("Called unknown handler: %s", arg);
+
+	return handler->init(ctx, arg);
 }
 
-static void fetch_uid(struct imap_fetch_context *ctx, struct mail *mail)
+struct imap_fetch_context *imap_fetch_init(struct client *client)
 {
-	str_printfa(ctx->str, "UID %u ", mail->uid);
+	struct imap_fetch_context *ctx;
+
+	if (fetch_handlers == NULL) {
+		imap_fetch_handlers_register(default_handlers,
+					     sizeof(default_handlers) /
+					     sizeof(default_handlers[0]));
+	}
+
+	ctx = p_new(client->cmd_pool, struct imap_fetch_context, 1);
+	ctx->client = client;
+	ctx->box = client->mailbox;
+
+	ctx->cur_str = str_new(default_pool, 8192);
+	ctx->seen_flag.flags = MAIL_SEEN;
+	ctx->all_headers_buf =
+		buffer_create_dynamic(client->cmd_pool, 128, (size_t)-1);
+	ctx->handlers =
+		buffer_create_dynamic(client->cmd_pool, 128, (size_t)-1);
+	return ctx;
 }
 
-static int fetch_flags(struct imap_fetch_context *ctx, struct mail *mail,
-		       const struct mail_full_flags *flags)
+void imap_fetch_add_handler(struct imap_fetch_context *ctx,
+			    imap_fetch_handler_t *handler, void *context)
 {
-	struct mail_full_flags full_flags;
+	const struct imap_fetch_context_handler *handlers;
+	struct imap_fetch_context_handler h;
+	size_t i, size;
+
+	if (context == NULL) {
+		/* don't allow duplicate handlers */
+		handlers = buffer_get_data(ctx->handlers, &size);
+		size /= sizeof(*handlers);
+
+		for (i = 0; i < size; i++) {
+			if (handlers[i].handler == handler &&
+			    handlers[i].context == NULL)
+				return;
+		}
+	}
 
-	if (flags == NULL) {
-		flags = mail->get_flags(mail);
-		if (flags == NULL)
-			return FALSE;
+	memset(&h, 0, sizeof(h));
+	h.handler = handler;
+	h.context = context;
+
+	buffer_append(ctx->handlers, &h, sizeof(h));
+}
+
+void imap_fetch_begin(struct imap_fetch_context *ctx,
+		      struct mail_search_arg *search_arg)
+{
+	const void *null = NULL;
+	const void *data;
+
+	if (ctx->flags_update_seen) {
+		if (mailbox_is_readonly(ctx->box))
+			ctx->flags_update_seen = FALSE;
+		else if (!ctx->flags_have_handler) {
+			ctx->flags_show_only_seen_changes = TRUE;
+			(void)imap_fetch_init_handler(ctx, "FLAGS");
+		}
 	}
 
-	if (ctx->update_seen) {
-		/* \Seen change isn't shown by get_flags() yet */
-		full_flags = *flags;
-		full_flags.flags |= MAIL_SEEN;
-		flags = &full_flags;
+	if (buffer_get_used_size(ctx->all_headers_buf) != 0 &&
+	    ((ctx->fetch_data & (MAIL_FETCH_STREAM_HEADER |
+				 MAIL_FETCH_STREAM_BODY)) == 0)) {
+		buffer_append(ctx->all_headers_buf, &null, sizeof(null));
+
+		data = buffer_get_data(ctx->all_headers_buf, NULL);
+		ctx->all_headers_ctx =
+			mailbox_header_lookup_init(ctx->box, data);
 	}
 
-	str_append(ctx->str, "FLAGS (");
-	imap_write_flags(ctx->str, flags);
-	str_append(ctx->str, ") ");
-	return TRUE;
+	ctx->trans = mailbox_transaction_begin(ctx->box, TRUE);
+	ctx->select_counter = ctx->client->select_counter;
+	ctx->search_ctx =
+		mailbox_search_init(ctx->trans, NULL, search_arg, NULL,
+				    ctx->fetch_data, ctx->all_headers_ctx);
 }
 
-static int fetch_internaldate(struct imap_fetch_context *ctx, struct mail *mail)
+int imap_fetch(struct imap_fetch_context *ctx)
 {
-	time_t time;
+	const struct imap_fetch_context_handler *handlers;
+	size_t size;
+	int ret;
+
+	if (ctx->cont_handler != NULL) {
+		if ((ret = ctx->cont_handler(ctx)) <= 0) {
+			if (ret < 0)
+				ctx->failed = TRUE;
+			return ret;
+		}
+
+		ctx->cont_handler = NULL;
+		ctx->cur_offset = 0;
+	}
+
+	handlers = buffer_get_data(ctx->handlers, &size);
+	size /= sizeof(*handlers);
+
+	for (;;) {
+		if (o_stream_get_buffer_used_size(ctx->client->output) >=
+		    CLIENT_OUTPUT_OPTIMAL_SIZE) {
+			ret = o_stream_flush(ctx->client->output);
+			if (ret <= 0)
+				return ret;
+		}
+
+		if (ctx->cur_mail == NULL) {
+			if (ctx->cur_input != NULL) {
+				i_stream_unref(ctx->cur_input);
+                                ctx->cur_input = NULL;
+			}
+
+			ctx->cur_mail = mailbox_search_next(ctx->search_ctx);
+			if (ctx->cur_mail == NULL)
+				break;
 
-	time = mail->get_received_date(mail);
-	if (time == (time_t)-1)
-		return FALSE;
+			str_printfa(ctx->cur_str, "* %u FETCH (",
+				    ctx->cur_mail->seq);
+			if (o_stream_send(ctx->client->output,
+					  str_data(ctx->cur_str),
+					  str_len(ctx->cur_str)) < 0) {
+				ctx->failed = TRUE;
+				return -1;
+			}
+
+			str_truncate(ctx->cur_str, 0);
+			str_append_c(ctx->cur_str, ' ');
+			ctx->first = TRUE;
+		}
+
+		for (; ctx->cur_handler < size; ctx->cur_handler++) {
+			t_push();
+			ret = handlers[ctx->cur_handler].
+				handler(ctx, ctx->cur_mail,
+					handlers[ctx->cur_handler].context);
+			t_pop();
 
-	str_printfa(ctx->str, "INTERNALDATE \"%s\" ", imap_to_datetime(time));
-	return TRUE;
+			if (ret <= 0) {
+				if (ret < 0)
+					ctx->failed = TRUE;
+				else
+					i_assert(ctx->cont_handler != NULL);
+				return ret;
+			}
+
+			ctx->cont_handler = NULL;
+			ctx->cur_offset = 0;
+		}
+
+		if (str_len(ctx->cur_str) > 1) {
+			if (o_stream_send(ctx->client->output,
+					  str_data(ctx->cur_str) + ctx->first,
+					  str_len(ctx->cur_str) - 1 -
+					  ctx->first) < 0) {
+				ctx->failed = TRUE;
+				return -1;
+			}
+			str_truncate(ctx->cur_str, 0);
+		}
+
+		if (o_stream_send(ctx->client->output, ")\r\n", 3) < 0) {
+			ctx->failed = TRUE;
+			return -1;
+		}
+
+		ctx->cur_mail = NULL;
+		ctx->cur_handler = 0;
+	}
+
+	return 1;
 }
 
-static int fetch_rfc822_size(struct imap_fetch_context *ctx, struct mail *mail)
+int imap_fetch_deinit(struct imap_fetch_context *ctx)
 {
-	uoff_t size;
+	str_free(ctx->cur_str);
+
+	if (ctx->cur_input != NULL) {
+		i_stream_unref(ctx->cur_input);
+		ctx->cur_input = NULL;
+	}
 
-	size = mail->get_size(mail);
-	if (size == (uoff_t)-1)
-		return FALSE;
+	if (ctx->search_ctx != NULL) {
+		if (mailbox_search_deinit(ctx->search_ctx) < 0)
+			ctx->failed = TRUE;
+	}
+	if (ctx->all_headers_ctx != NULL)
+		mailbox_header_lookup_deinit(ctx->all_headers_ctx);
 
-	str_printfa(ctx->str, "RFC822.SIZE %"PRIuUOFF_T" ", size);
-	return TRUE;
+	if (ctx->trans != NULL) {
+		if (ctx->failed)
+			mailbox_transaction_rollback(ctx->trans);
+		else {
+			if (mailbox_transaction_commit(ctx->trans) < 0)
+				ctx->failed = TRUE;
+		}
+	}
+	return ctx->failed ? -1 : 0;
 }
 
-static int fetch_body(struct imap_fetch_context *ctx, struct mail *mail)
+static int fetch_body(struct imap_fetch_context *ctx, struct mail *mail,
+		      void *context __attr_unused__)
 {
 	const char *body;
 
 	body = mail->get_special(mail, MAIL_FETCH_IMAP_BODY);
 	if (body == NULL)
-		return FALSE;
+		return -1;
 
-	if (ctx->first) {
-		if (o_stream_send_str(ctx->output, "BODY (") < 0)
-			return FALSE;
+	if (ctx->first)
 		ctx->first = FALSE;
-	} else {
-		if (o_stream_send_str(ctx->output, " BODY (") < 0)
-			return FALSE;
+	else {
+		if (o_stream_send(ctx->client->output, " ", 1) < 0)
+			return -1;
 	}
 
-	if (o_stream_send_str(ctx->output, body) < 0)
-		return FALSE;
+	if (o_stream_send(ctx->client->output, "BODY (", 6) < 0 ||
+	    o_stream_send_str(ctx->client->output, body) < 0 ||
+	    o_stream_send(ctx->client->output, ")", 1) < 0)
+		return -1;
+	return 1;
+}
 
-	if (o_stream_send(ctx->output, ")", 1) < 0)
-		return FALSE;
-	return TRUE;
+static int fetch_body_init(struct imap_fetch_context *ctx, const char *arg)
+{
+	if (arg[4] == '\0') {
+		ctx->fetch_data |= MAIL_FETCH_IMAP_BODY;
+		imap_fetch_add_handler(ctx, fetch_body, NULL);
+		return TRUE;
+	}
+	return fetch_body_section_init(ctx, arg);
 }
 
 static int fetch_bodystructure(struct imap_fetch_context *ctx,
-			       struct mail *mail)
+			       struct mail *mail, void *context __attr_unused__)
 {
 	const char *bodystructure;
 
 	bodystructure = mail->get_special(mail, MAIL_FETCH_IMAP_BODYSTRUCTURE);
 	if (bodystructure == NULL)
-		return FALSE;
+		return -1;
 
-	if (ctx->first) {
-		if (o_stream_send_str(ctx->output, "BODYSTRUCTURE (") < 0)
-			return FALSE;
+	if (ctx->first)
 		ctx->first = FALSE;
-	} else {
-		if (o_stream_send_str(ctx->output, " BODYSTRUCTURE (") < 0)
-			return FALSE;
+	else {
+		if (o_stream_send(ctx->client->output, " ", 1) < 0)
+			return -1;
 	}
 
-	if (o_stream_send_str(ctx->output, bodystructure) < 0)
-		return FALSE;
+	if (o_stream_send(ctx->client->output, "BODYSTRUCTURE (", 15) < 0 ||
+	    o_stream_send_str(ctx->client->output, bodystructure) < 0 ||
+	    o_stream_send(ctx->client->output, ")", 1) < 0)
+		return -1;
 
-	if (o_stream_send(ctx->output, ")", 1) < 0)
-		return FALSE;
+	return 1;
+}
+
+static int fetch_bodystructure_init(struct imap_fetch_context *ctx,
+				    const char *arg __attr_unused__)
+{
+	ctx->fetch_data |= MAIL_FETCH_IMAP_BODYSTRUCTURE;
+	imap_fetch_add_handler(ctx, fetch_bodystructure, NULL);
 	return TRUE;
 }
 
-static int fetch_envelope(struct imap_fetch_context *ctx, struct mail *mail)
+static int fetch_envelope(struct imap_fetch_context *ctx, struct mail *mail,
+			  void *context __attr_unused__)
 {
 	const char *envelope;
 
 	envelope = mail->get_special(mail, MAIL_FETCH_IMAP_ENVELOPE);
 	if (envelope == NULL)
-		return FALSE;
+		return -1;
 
-	if (ctx->first) {
-		if (o_stream_send_str(ctx->output, "ENVELOPE (") < 0)
-			return FALSE;
+	if (ctx->first)
 		ctx->first = FALSE;
-	} else {
-		if (o_stream_send_str(ctx->output, " ENVELOPE (") < 0)
-			return FALSE;
+	else {
+		if (o_stream_send(ctx->client->output, " ", 1) < 0)
+			return -1;
 	}
 
-	if (o_stream_send_str(ctx->output, envelope) < 0)
-		return FALSE;
+	if (o_stream_send(ctx->client->output, "ENVELOPE (", 10) < 0 ||
+	    o_stream_send_str(ctx->client->output, envelope) < 0 ||
+	    o_stream_send(ctx->client->output, ")", 1) < 0)
+		return -1;
+	return 1;
+}
 
-	if (o_stream_send(ctx->output, ")", 1) < 0)
-		return FALSE;
+static int fetch_envelope_init(struct imap_fetch_context *ctx,
+			       const char *arg __attr_unused__)
+{
+	ctx->fetch_data |= MAIL_FETCH_IMAP_ENVELOPE;
+	imap_fetch_add_handler(ctx, fetch_envelope, NULL);
 	return TRUE;
 }
 
-static int fetch_send_rfc822(struct imap_fetch_context *ctx, struct mail *mail)
-{
-	struct message_size hdr_size, body_size;
-	struct istream *stream;
-	const char *str;
-
-	stream = mail->get_stream(mail, &hdr_size, &body_size);
-	if (stream == NULL)
-		return FALSE;
-
-	message_size_add(&body_size, &hdr_size);
-
-	str = t_strdup_printf(" RFC822 {%"PRIuUOFF_T"}\r\n",
-			      body_size.virtual_size);
-	if (ctx->first) {
-		str++; ctx->first = FALSE;
-	}
-	if (o_stream_send_str(ctx->output, str) < 0)
-		return FALSE;
-
-	return message_send(ctx->output, stream, &body_size,
-			    0, body_size.virtual_size, NULL,
-			    !mail->has_no_nuls) >= 0;
-}
-
-static int fetch_send_rfc822_header(struct imap_fetch_context *ctx,
-				    struct mail *mail)
-{
-	struct message_size hdr_size;
-	struct istream *stream;
-	const char *str;
-
-	stream = mail->get_stream(mail, &hdr_size, NULL);
-	if (stream == NULL)
-		return FALSE;
-
-	str = t_strdup_printf(" RFC822.HEADER {%"PRIuUOFF_T"}\r\n",
-			      hdr_size.virtual_size);
-	if (ctx->first) {
-		str++; ctx->first = FALSE;
-	}
-	if (o_stream_send_str(ctx->output, str) < 0)
-		return FALSE;
-
-	return message_send(ctx->output, stream, &hdr_size,
-			    0, hdr_size.virtual_size, NULL,
-			    !mail->has_no_nuls) >= 0;
-}
-
-static int fetch_send_rfc822_text(struct imap_fetch_context *ctx,
-				  struct mail *mail)
-{
-	struct message_size hdr_size, body_size;
-	struct istream *stream;
-	const char *str;
-
-	stream = mail->get_stream(mail, &hdr_size, &body_size);
-	if (stream == NULL)
-		return FALSE;
-
-	str = t_strdup_printf(" RFC822.TEXT {%"PRIuUOFF_T"}\r\n",
-			      body_size.virtual_size);
-	if (ctx->first) {
-		str++; ctx->first = FALSE;
-	}
-	if (o_stream_send_str(ctx->output, str) < 0)
-		return FALSE;
-
-	i_stream_seek(stream, hdr_size.physical_size);
-	return message_send(ctx->output, stream, &body_size,
-			    0, body_size.virtual_size, NULL,
-			    !mail->has_no_nuls) >= 0;
-}
-
-static int fetch_mail(struct imap_fetch_context *ctx, struct mail *mail)
+static int fetch_flags(struct imap_fetch_context *ctx, struct mail *mail,
+		       void *context __attr_unused__)
 {
 	const struct mail_full_flags *flags;
-	struct imap_fetch_body_data *body;
-	size_t len, orig_len;
-	int failed, data_written, seen_updated = FALSE;
+	struct mail_full_flags full_flags;
+
+	flags = mail->get_flags(mail);
+	if (flags == NULL)
+		return -1;
 
-	if (!ctx->update_seen)
-		flags = NULL;
-	else {
-		flags = mail->get_flags(mail);
-		if (flags == NULL)
-			return FALSE;
+	if (ctx->flags_update_seen && (flags->flags & MAIL_SEEN) == 0) {
+		/* Add \Seen flag */
+		full_flags = *flags;
+		full_flags.flags |= MAIL_SEEN;
+		flags = &full_flags;
 
-		if ((flags->flags & MAIL_SEEN) == 0) {
-			if (mail->update_flags(mail, &ctx->seen_flag,
-					       MODIFY_ADD) < 0)
-				return FALSE;
-			seen_updated = TRUE;
-		}
+		if (mail->update_flags(mail, &ctx->seen_flag, MODIFY_ADD) < 0)
+			return -1;
+	} else if (ctx->flags_show_only_seen_changes) {
+		return 1;
 	}
 
-	t_push();
-
-	str_truncate(ctx->str, 0);
-	str_printfa(ctx->str, "* %u FETCH (", mail->seq);
-	orig_len = str_len(ctx->str);
-
-	failed = TRUE;
-	data_written = FALSE;
-	do {
-		/* write the data into temp string */
-		if (ctx->imap_data & IMAP_FETCH_UID)
-			fetch_uid(ctx, mail);
-		if ((ctx->fetch_data & MAIL_FETCH_FLAGS) || seen_updated)
-			if (!fetch_flags(ctx, mail, flags))
-				break;
-		if (ctx->fetch_data & MAIL_FETCH_RECEIVED_DATE)
-			if (!fetch_internaldate(ctx, mail))
-				break;
-		if (ctx->fetch_data & MAIL_FETCH_SIZE)
-			if (!fetch_rfc822_size(ctx, mail))
-				break;
-
-		/* send the data written into temp string */
-		len = str_len(ctx->str);
-		ctx->first = len == orig_len;
-
-		if (!ctx->first)
-			str_truncate(ctx->str, --len);
-		if (o_stream_send(ctx->output, str_data(ctx->str), len) < 0)
-			break;
-
-		data_written = TRUE;
+	str_append(ctx->cur_str, "FLAGS (");
+	imap_write_flags(ctx->cur_str, flags);
+	str_append(ctx->cur_str, ") ");
+	return 1;
+}
 
-		/* medium size data .. seems to be faster without
-		 putting through string */
-		if (ctx->fetch_data & MAIL_FETCH_IMAP_BODY)
-			if (!fetch_body(ctx, mail))
-				break;
-		if (ctx->fetch_data & MAIL_FETCH_IMAP_BODYSTRUCTURE)
-			if (!fetch_bodystructure(ctx, mail))
-				break;
-		if (ctx->fetch_data & MAIL_FETCH_IMAP_ENVELOPE)
-			if(!fetch_envelope(ctx, mail))
-				break;
+static int fetch_flags_init(struct imap_fetch_context *ctx,
+			    const char *arg __attr_unused__)
+{
+	ctx->flags_have_handler = TRUE;
+	ctx->fetch_data |= MAIL_FETCH_FLAGS;
+	imap_fetch_add_handler(ctx, fetch_flags, NULL);
+	return TRUE;
+}
 
-		/* large data */
-		if (ctx->imap_data & IMAP_FETCH_RFC822)
-			if (!fetch_send_rfc822(ctx, mail))
-				break;
-		if (ctx->imap_data & IMAP_FETCH_RFC822_HEADER)
-			if (!fetch_send_rfc822_header(ctx, mail))
-				break;
-		if (ctx->imap_data & IMAP_FETCH_RFC822_TEXT)
-			if (!fetch_send_rfc822_text(ctx, mail))
-				break;
-
-		failed = FALSE;
+static int fetch_internaldate(struct imap_fetch_context *ctx, struct mail *mail,
+			      void *context __attr_unused__)
+{
+	time_t time;
 
-		for (body = ctx->bodies; body != NULL; body = body->next) {
-			if (!imap_fetch_body_section(ctx, body, mail)) {
-				failed = TRUE;
-				break;
-			}
-		}
-	} while (0);
+	time = mail->get_received_date(mail);
+	if (time == (time_t)-1)
+		return -1;
 
-	if (data_written) {
-		if (o_stream_send(ctx->output, ")\r\n", 3) < 0)
-			failed = TRUE;
-	}
-
-	t_pop();
-	return !failed;
+	str_printfa(ctx->cur_str, "INTERNALDATE \"%s\" ",
+		    imap_to_datetime(time));
+	return 1;
 }
 
-int imap_fetch(struct client *client,
-	       enum mail_fetch_field fetch_data,
-	       enum imap_fetch_field imap_data,
-	       struct imap_fetch_body_data *bodies,
-	       struct mail_search_arg *search_args)
-{
-	struct mailbox *box = client->mailbox;
-	struct imap_fetch_context ctx;
-	struct mailbox_transaction_context *t;
-	struct mail *mail;
-	struct imap_fetch_body_data *body;
-	const char *null = NULL;
-	const char *const *arr;
-	buffer_t *buffer;
 
-	memset(&ctx, 0, sizeof(ctx));
-	ctx.fetch_data = fetch_data;
-	ctx.imap_data = imap_data;
-	ctx.bodies = bodies;
-	ctx.output = client->output;
-	ctx.select_counter = client->select_counter;
-	ctx.seen_flag.flags = MAIL_SEEN;
+static int fetch_internaldate_init(struct imap_fetch_context *ctx,
+				   const char *arg __attr_unused__)
+{
+	ctx->fetch_data |= MAIL_FETCH_RECEIVED_DATE;
+	imap_fetch_add_handler(ctx, fetch_internaldate, NULL);
+	return TRUE;
+}
 
-	if (!mailbox_is_readonly(box)) {
-		/* If we have any BODY[..] sections, \Seen flag is added for
-		   all messages. */
-		for (body = bodies; body != NULL; body = body->next) {
-			if (!body->peek) {
-				ctx.update_seen = TRUE;
-				break;
-			}
-		}
-
-		if (imap_data & (IMAP_FETCH_RFC822|IMAP_FETCH_RFC822_TEXT))
-			ctx.update_seen = TRUE;
-	}
+static int fetch_uid(struct imap_fetch_context *ctx, struct mail *mail,
+		     void *context __attr_unused__)
+{
+	str_printfa(ctx->cur_str, "UID %u ", mail->uid);
+	return 1;
+}
 
-	/* If we have only BODY[HEADER.FIELDS (...)] fetches, get them
-	   separately rather than parsing the full header so mail storage
-	   can try to cache them. */
-	ctx.body_fetch_from_cache = (imap_data & (IMAP_FETCH_RFC822 |
-						  IMAP_FETCH_RFC822_HEADER |
-						  IMAP_FETCH_RFC822_TEXT)) == 0;
-	if (ctx.body_fetch_from_cache) {
-		buffer = buffer_create_dynamic(pool_datastack_create(),
-					       64, (size_t)-1);
-		for (body = bodies; body != NULL; body = body->next) {
-			if (strncmp(body->section, "HEADER.FIELDS ", 14) != 0) {
-				ctx.body_fetch_from_cache = FALSE;
-				break;
-			}
-
-			arr = imap_fetch_get_body_fields(body->section + 14);
-			while (*arr != NULL) {
-				buffer_append(buffer, arr, sizeof(*arr));
-				arr++;
-			}
-		}
-		buffer_append(buffer, &null, sizeof(null));
-		ctx.headers_ctx = !ctx.body_fetch_from_cache ? NULL :
-			mailbox_header_lookup_init(box, buffer_get_data(buffer,
-									NULL));
-	}
+static int fetch_uid_init(struct imap_fetch_context *ctx __attr_unused__,
+			  const char *arg __attr_unused__)
+{
+	imap_fetch_add_handler(ctx, fetch_uid, NULL);
+	return TRUE;
+}
 
-	t = mailbox_transaction_begin(box, TRUE);
-	ctx.search_ctx = mailbox_search_init(t, NULL, search_args, NULL,
-					     fetch_data, ctx.headers_ctx);
-	if (ctx.search_ctx == NULL)
-		ctx.failed = TRUE;
-	else {
-		ctx.str = str_new(default_pool, 8192);
-		while ((mail = mailbox_search_next(ctx.search_ctx)) != NULL) {
-			if (!fetch_mail(&ctx, mail)) {
-				ctx.failed = TRUE;
-				break;
-			}
-		}
-		str_free(ctx.str);
-
-		if (mailbox_search_deinit(ctx.search_ctx) < 0)
-			ctx.failed = TRUE;
-	}
-	if (ctx.headers_ctx != NULL)
-		mailbox_header_lookup_deinit(ctx.headers_ctx);
-
-	if (ctx.failed)
-		mailbox_transaction_rollback(t);
-	else {
-		if (mailbox_transaction_commit(t) < 0)
-			ctx.failed = TRUE;
-	}
-	return ctx.failed ? -1 : 0;
-}
+const struct imap_fetch_handler default_handlers[7] = {
+	{ "BODY", fetch_body_init },
+	{ "BODYSTRUCTURE", fetch_bodystructure_init },
+	{ "ENVELOPE", fetch_envelope_init },
+	{ "FLAGS", fetch_flags_init },
+	{ "INTERNALDATE", fetch_internaldate_init },
+	{ "RFC822", fetch_rfc822_init },
+	{ "UID", fetch_uid_init }
+};

--- a/src/imap/imap-fetch.h	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/imap-fetch.h	Sun Aug 15 06:40:30 2004 +0300
@@ -1,52 +1,73 @@
 #ifndef __IMAP_FETCH_H
 #define __IMAP_FETCH_H
 
-enum imap_fetch_field {
-	IMAP_FETCH_UID			= 0x01,
-	IMAP_FETCH_RFC822		= 0x02,
-	IMAP_FETCH_RFC822_HEADER	= 0x04,
-	IMAP_FETCH_RFC822_TEXT		= 0x08
+struct imap_fetch_context;
+
+/* Returns 1 = ok, 0 = client output buffer full, call again, -1 = error.
+   mail = NULL for deinit. */
+typedef int imap_fetch_handler_t(struct imap_fetch_context *ctx,
+				 struct mail *mail, void *context);
+
+struct imap_fetch_handler {
+	const char *name;
+
+	/* Returns FALSE if arg is invalid. */
+	int (*init)(struct imap_fetch_context *ctx, const char *arg);
 };
 
-struct imap_fetch_body_data {
-	struct imap_fetch_body_data *next;
-
-	const char *section; /* NOTE: always uppercased */
-	uoff_t skip, max_size; /* if you don't want max_size,
-	                          set it to (uoff_t)-1 */
-	unsigned int skip_set:1;
-	unsigned int peek:1;
+struct imap_fetch_context_handler {
+	imap_fetch_handler_t *handler;
+	void *context;
 };
 
 struct imap_fetch_context {
+	struct client *client;
+	struct mailbox *box;
+
+	struct mailbox_transaction_context *trans;
 	struct mail_search_context *search_ctx;
 
 	enum mail_fetch_field fetch_data;
-	enum imap_fetch_field imap_data;
-	struct imap_fetch_body_data *bodies;
-	struct mailbox_header_lookup_ctx *headers_ctx;
+	buffer_t *all_headers_buf;
+        struct mailbox_header_lookup_ctx *all_headers_ctx;
+
+	buffer_t *handlers;
 
-	string_t *str;
-	struct ostream *output;
-	const char *prefix;
+	struct mail *cur_mail;
+	unsigned int cur_handler;
+	uoff_t cur_size, cur_offset;
+	string_t *cur_str;
+	struct istream *cur_input;
+	int skip_cr;
+	int (*cont_handler)(struct imap_fetch_context *ctx);
+
 	unsigned int select_counter;
 
-	int update_seen;
 	struct mail_full_flags seen_flag;
 
-	int first, failed, body_fetch_from_cache;
+	unsigned int flags_have_handler:1;
+	unsigned int flags_update_seen:1;
+	unsigned int flags_show_only_seen_changes:1;
+	unsigned int update_partial:1;
+	unsigned int first:1;
+	unsigned int failed:1;
 };
 
-int imap_fetch(struct client *client,
-	       enum mail_fetch_field fetch_data,
-	       enum imap_fetch_field imap_data,
-	       struct imap_fetch_body_data *bodies,
-	       struct mail_search_arg *search_args);
+void imap_fetch_handlers_register(const struct imap_fetch_handler *handlers,
+				  size_t count);
+
+void imap_fetch_add_handler(struct imap_fetch_context *ctx,
+			    imap_fetch_handler_t *handler, void *context);
 
-int imap_fetch_body_section(struct imap_fetch_context *ctx,
-			    const struct imap_fetch_body_data *body,
-			    struct mail *mail);
+struct imap_fetch_context *imap_fetch_init(struct client *client);
+int imap_fetch_deinit(struct imap_fetch_context *ctx);
+int imap_fetch_init_handler(struct imap_fetch_context *ctx, const char *arg);
 
-const char *const *imap_fetch_get_body_fields(const char *fields);
+void imap_fetch_begin(struct imap_fetch_context *ctx,
+		      struct mail_search_arg *search_arg);
+int imap_fetch(struct imap_fetch_context *ctx);
+
+int fetch_body_section_init(struct imap_fetch_context *ctx, const char *arg);
+int fetch_rfc822_init(struct imap_fetch_context *ctx, const char *arg);
 
 #endif

--- a/src/imap/imap-messageset.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/imap-messageset.c	Sun Aug 15 06:40:30 2004 +0300
@@ -28,7 +28,8 @@
 	return num;
 }
 
-struct mail_search_seqset *imap_messageset_parse(const char *messageset)
+struct mail_search_seqset *
+imap_messageset_parse(pool_t pool, const char *messageset)
 {
         struct mail_search_seqset *ret, **next;
 	uint32_t seq1, seq2;
@@ -75,7 +76,7 @@
 			seq2 = temp;
 		}
 
-		*next = t_new(struct mail_search_seqset, 1);
+		*next = p_new(pool, struct mail_search_seqset, 1);
 		(*next)->seq1 = seq1;
 		(*next)->seq2 = seq2;
 		next = &(*next)->next;

--- a/src/imap/imap-messageset.h	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/imap-messageset.h	Sun Aug 15 06:40:30 2004 +0300
@@ -1,6 +1,7 @@
 #ifndef __IMAP_MESSAGESET_H
 #define __IMAP_MESSAGESET_H
 
-struct mail_search_seqset *imap_messageset_parse(const char *messageset);
+struct mail_search_seqset *
+imap_messageset_parse(pool_t pool, const char *messageset);
 
 #endif

--- a/src/imap/imap-search.c	Sun Aug 15 05:54:47 2004 +0300
+++ b/src/imap/imap-search.c	Sun Aug 15 06:40:30 2004 +0300
@@ -14,13 +14,13 @@
 };
 
 static int
-imap_uidset_parse(struct mailbox *box, const char *uidset,
+imap_uidset_parse(pool_t pool, struct mailbox *box, const char *uidset,
 		  struct mail_search_seqset **seqset_r, const char **error_r)
 {
 	struct mail_search_seqset *seqset, **p;
 	int syntax, last;
 
-	*seqset_r = imap_messageset_parse(uidset);
+	*seqset_r = imap_messageset_parse(pool, uidset);
 	if (*seqset_r == NULL) {
 		*error_r = "Invalid UID messageset";
 		return -1;
@@ -99,11 +99,11 @@
 		return FALSE;
 	}
 
-	sarg->value.str = str_ucase(IMAP_ARG_STR(*args));
+	sarg->value.str = str_ucase(p_strdup(data->pool, IMAP_ARG_STR(*args)));
 	*args += 1;
 
 	if (hdr_name != NULL)
-                sarg->hdr_field_name = hdr_name;
+                sarg->hdr_field_name = p_strdup(data->pool, hdr_name);
 
 	return TRUE;
 }
@@ -328,7 +328,7 @@
 			if (!ARG_NEW(SEARCH_SEQSET))
 				return FALSE;
 
-			return imap_uidset_parse(data->box,
+			return imap_uidset_parse(data->pool, data->box,
 						 (*next_sarg)->value.str,
 						 &(*next_sarg)->value.seqset,
 						 &data->error) == 0;
@@ -367,7 +367,7 @@
 	default:
 		if (*str == '*' || (*str >= '0' && *str <= '9')) {
 			/* <message-set> */
-			seqset = imap_messageset_parse(str);
+			seqset = imap_messageset_parse(data->pool, str);
 			if (seqset == NULL) {
 				data->error = "Invalid messageset";
 				return FALSE;
@@ -412,15 +412,15 @@
 	return first_sarg;
 }
 
-int imap_search_get_msgset_arg(const char *messageset,
-			       struct mail_search_arg **arg_r,
-			       const char **error_r)
+static int imap_search_get_msgset_arg(pool_t pool, const char *messageset,
+				      struct mail_search_arg **arg_r,
+				      const char **error_r)
 {
 	struct mail_search_arg *arg;
 
-	arg = t_new(struct mail_search_arg, 1);
+	arg = p_new(pool, struct mail_search_arg, 1);
 	arg->type = SEARCH_SEQSET;
-	arg->value.seqset = imap_messageset_parse(messageset);
+	arg->value.seqset = imap_messageset_parse(pool, messageset);
 	if (arg->value.seqset == NULL) {
 		*error_r = "Invalid messageset";
 		return -1;
@@ -429,16 +429,17 @@
 	return 0;
 }
 
-int imap_search_get_uidset_arg(struct mailbox *box, const char *uidset,
-			       struct mail_search_arg **arg_r,
-			       const char **error_r)
+static int
+imap_search_get_uidset_arg(pool_t pool, struct mailbox *box, const char *uidset,
+			   struct mail_search_arg **arg_r, const char **error_r)
 {
 	struct mail_search_arg *arg;
 
-	arg = t_new(struct mail_search_arg, 1);
+	arg = p_new(pool, struct mail_search_arg, 1);
 	arg->type = SEARCH_SEQSET;
 	*arg_r = arg;
-	return imap_uidset_parse(box, uidset, &arg->value.seqset, error_r);
+	return imap_uidset_parse(pool, box, uidset, &arg->value.seqset,
+				 error_r);
 }
 
 struct mail_search_arg *
@@ -449,14 +450,15 @@
 	int ret;
 
 	if (!uid) {
-		ret = imap_search_get_msgset_arg(set, &search_arg,
-						 &error);
+		ret = imap_search_get_msgset_arg(client->cmd_pool, set,
+						 &search_arg, &error);
 	} else {
-		ret = imap_search_get_uidset_arg(client->mailbox, set,
+		ret = imap_search_get_uidset_arg(client->cmd_pool,
+						 client->mailbox, set,
 						 &search_arg, &error);
 	}
 	if (ret < 0) {
-		client_send_tagline(client, t_strconcat("BAD ", error, NULL));
+		client_send_command_error(client, error);
 		return NULL;
 	}