/* This file is part of Mailfromd.
   Copyright (C) 2020 Sergey Poznyakoff

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 3, or (at your option)
   any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#ifdef HAVE_CONFIG_H
# include <config.h>
#endif
#include <stdlib.h>
#include <string.h>
#include <mailutils/errno.h>
#include <mailutils/filter.h>
#include <mailutils/cctype.h>

/*
 * This file defines a mailutils filter implementing the two DKIM
 * canonicalization algrorithms: "simple" and "relaxed" (see
 * RFC 6376, 3.4.  "Canonicalization" (page 13).
 *
 * The filter reads a LF-delimited message stream from its input
 * and outpus a canonicalized message, with header and body parts
 * processed using the selected algorithm.
 *
 * For further processing the output stream must be passed through a
 * CRLF encoder.
 *
 * The filter is governed by the following structure:
 */

struct encoder_state {
	int canon[2];   /* Canonicalization algorithm.  0 - for headers,
			   1 - for body.  See the DKIM_CANON_ constants
			   in dkim.h */
	int state;      /* Encoder state (see below) */
	size_t nlcount; /* Number of contiguous empty lines for "relaxed"
			   body algorithm. */
			   
};

/* Enoder states: */
enum {
	/* Header "simple" canonicalization states */
	HS_INIT,
	HS_NL,
	/* Header "relaxed" canonicalization states */
	HR_NAME,
	HR_SPACE,
	HR_COLON,
	HR_VALUE,
	HR_WS,
	HR_NL,
	/* Delimiter */
	ST_DELIM,
	/* Body "simple" canonicalization states */
	BS_INIT,
	BS_NL,
	/* Body "relaxed" canonicalization states */
	BR_INIT,
	BR_WS,
	BR_NL
};

/* Initial states states */
static int init_state[][2] = {
	{ HS_INIT, HR_NAME },
	{ BS_INIT, BR_INIT }
};

static enum mu_filter_result
dkim_canonicalizer(void *xd,
		   enum mu_filter_command cmd,
		   struct mu_filter_io *iobuf)
{
	struct encoder_state *encoder = xd;
	const char *iptr, *iendptr;
	char *optr, *oendptr;
	
	switch (cmd) {
	case mu_filter_init:
		encoder->state = init_state[0][encoder->canon[0]];
	case mu_filter_done:
		return mu_filter_ok;
	default:
		break;
	}

	iptr = iobuf->input;
	iendptr = iptr + iobuf->isize;
	optr = iobuf->output;
	oendptr = optr + iobuf->osize;
	
	while (iptr < iendptr && optr < oendptr) {
		switch (encoder->state) {
			/* Header "simple" canonicalization */
		case HS_INIT:
			if (*iptr == '\n')
				encoder->state = HS_NL;
			else
				*optr++ = *iptr;
			iptr++;
			break;

		case HS_NL:
			if (*iptr == '\n') {
				encoder->state = ST_DELIM;
			} else {
				*optr++ = '\n';
				encoder->state = HS_INIT;
			}
			break;

			/* Header "relaxed" canonicalization */
		case HR_NAME:
			if (mu_isblank(*iptr)) {
				iptr++;
				encoder->state = HR_SPACE;
			} else if (*iptr == ':') {
				*optr++ = *iptr++;
				encoder->state = HR_COLON;
			} else if (mu_isheadr(*iptr) ||
				   /* 
				    * Work around the bug in mailutils 3.9:
				    * the MU_CTYPE_HEADR class did not include
				    * underscore.
				    */
				   *iptr == '_') {
				*optr++ = mu_tolower(*iptr++);
			} else {
				iobuf->errcode = MU_ERR_USER0;
				return mu_filter_failure;
			}
			break;

		case HR_SPACE:
			if (mu_isblank(*iptr))
				iptr++;
			else if (*iptr == ':') {
				*optr++ = *iptr++;
				encoder->state = HR_COLON;
			} else {
				iobuf->errcode = MU_ERR_USER0;
				return mu_filter_failure;
			}
			break;

		case HR_COLON:
			if (mu_isblank(*iptr))
				iptr++;
			else
				encoder->state = HR_VALUE;
			break;

		case HR_VALUE:
			if (mu_isblank(*iptr)) {
				iptr++;
				encoder->state = HR_WS;
			} else if (*iptr == '\n') {
				iptr++;
				encoder->state = HR_NL;
			} else
				*optr++ = *iptr++;
			break;

		case HR_WS:
			if (mu_isblank(*iptr))
				iptr++;
			else if (*iptr == '\n') {
				iptr++;
				encoder->state = HR_NL;
			} else {
				*optr++ = ' ';
				encoder->state = HR_VALUE;
			}
			break;

		case HR_NL:
			if (*iptr == '\n') {
				encoder->state = ST_DELIM;
			} else if (mu_isblank(*iptr)) {
				iptr++;
				encoder->state = HR_WS;
			} else {
				*optr++ = '\n';
				encoder->state = HR_NAME;
			}
			break;

			/* Delimiter between header and body */
		case ST_DELIM:
			if (oendptr - optr < 2)
				goto end;
			iptr++;
			*optr++ = '\n';
			*optr++ = '\n';
			encoder->state = init_state[1][encoder->canon[1]];
			break;
	  
			/* Body "simple" canonicalization */
		case BS_INIT:
			if (*iptr == '\n') {
				iptr++;
				encoder->nlcount++;
				encoder->state = BS_NL;
			} else
				*optr++ = *iptr++;
			break;
	  
		case BS_NL:
			if (*iptr == '\n') {
				iptr++;
				encoder->nlcount++;
			} else {
				for (; encoder->nlcount; encoder->nlcount--) {
					if (optr == oendptr)
						goto end;
					*optr++ = '\n';
				}
				encoder->state = BS_INIT;
			}
			break;

			/* Body "relaxed" canonicalization */
		case BR_INIT:
			if (*iptr == '\n') {
				encoder->nlcount++;
				encoder->state = BR_NL;
			} else if (mu_isblank (*iptr))
				encoder->state = BR_WS;
			else
				*optr++ = *iptr;
			iptr++;
			break;

		case BR_WS:
			if (*iptr == '\n') {
				iptr++;				
				encoder->nlcount++;		  
				encoder->state = BR_NL;
			} else if (!mu_isblank(*iptr)) {
				*optr++ = ' ';
				encoder->state = BR_INIT;
			} else {
				iptr++;
			}
			break;

		case BR_NL:
			if (*iptr == '\n') {
				iptr++;
				encoder->nlcount++;
			} else {
				for (; encoder->nlcount; encoder->nlcount--) {
					if (optr == oendptr)
						goto end;
					*optr++ = '\n';
				}
				encoder->state = BR_INIT;
			}
			break;
		}
	}
	
	if (cmd == mu_filter_lastbuf && iptr == iendptr &&
	    encoder->state >= BS_INIT) {
		if (oendptr == optr) {
			iobuf->osize++;
			return mu_filter_moreoutput;
		}
		*optr++ = '\n';
	}
  
end:
	iobuf->isize = iptr - iobuf->input;
	iobuf->osize = optr - iobuf->output;

	return mu_filter_ok;
}

/*
 * Create a message canonicalizer.  Arguments:
 * 
 * pstream       return pointer
 * stream        input stream
 * canon_header  header canonicalization algorithm
 * canon_body    body canonicalization algorithm
 *
 * Return value: mailutils error code.
 */
int
dkim_canonicalizer_create(mu_stream_t *pstream,
			  mu_stream_t stream,
			  int canon_header,
			  int canon_body,
			  int flags)
{
	struct encoder_state *encoder = malloc(sizeof (*encoder));
	if (!encoder)
		return ENOMEM;
	memset(encoder, 0, sizeof(*encoder));
	encoder->canon[0] = canon_header;
	encoder->canon[1] = canon_body;
	return mu_filter_stream_create(pstream, stream,
				       MU_FILTER_ENCODE,
				       dkim_canonicalizer, encoder, flags);
}

#if 0
/*
 * The canonicalizer can be registered as a regular mailutils filter.
 * If such approach is ever needed, uncomment this block.
 */

static int
alloc_state(void **pret, int mode, int argc, const char **argv)
{
	struct encoder_state *encoder;
	switch (mode) {
	case MU_FILTER_ENCODE:
		encoder = malloc(sizeof(*encoder));
		if (!encoder)
			return ENOMEM;
		memset(encoder, 0, sizeof(*encoder));
		*pret = encoder;
		break;
      
	case MU_FILTER_DECODE:;
	}
	return 0;
}

static struct _mu_filter_record dkim_canonicalize_filter_s = {
	"DKIM",
	alloc_state,
	dkim_canonicalizer,
	NULL
};

mu_filter_record_t dkim_canonicalize_filter = &dkim_canonicalize_filter_s;
#endif