From d5f9f86bba39e5a086d9a7d62627b9d3a03c10b5 Mon Sep 17 00:00:00 2001
From: "Stuart D. Gathman" <stuart@gathman.org>
Date: Thu, 23 Apr 2020 15:52:20 -0400
Subject: [PATCH] Use utf-8 decoding with surrogateescape for invalid utf-8 for
 env and hdr val

---
 Milter/__init__.py | 32 ++++++++++++++++++++++----------
 miltermodule.c     |  2 +-
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/Milter/__init__.py b/Milter/__init__.py
index 4710d82..542be0b 100755
--- a/Milter/__init__.py
+++ b/Milter/__init__.py
@@ -321,6 +321,14 @@ class Base(object):
   # this almost always results in terminating the connection.
   @nocallback
   def hello(self,hostname): return CONTINUE
+  ## Called with bytes by default global envfrom callback.
+  # @since 1.0.5
+  # Converts from utf-8 to unicode with surrogate escape.  Can be overriden 
+  # to pass bytes to @link #header the header callback @endlink instead,
+  # or trap utf-8 conversion exception, etc.
+  def envfrom_bytes(self,*b):
+    s = (v.decode(encoding='utf-8',errors='surrogateescape') for v in b)
+    return self.envfrom(fld,*s)
   ## Called when the SMTP client says MAIL FROM. Called by the
   # <a href="milter_api/xxfi_envfrom.html">
   # xxfi_envfrom</a> callback.  
@@ -331,7 +339,15 @@ class Base(object):
   # <a href="http://tools.ietf.org/html/rfc5322">RFC 5322</a>,
   # see @link #header the header callback @endlink.
   @nocallback
-  def envfrom(self,f,*str): return CONTINUE
+  def envfrom(self,f,*s): return CONTINUE
+  ## Called with bytes by default global envrcpt callback.
+  # @since 1.0.5
+  # Converts from utf-8 to unicode with surrogate escape.  Can be overriden 
+  # to pass bytes to @link #header the header callback @endlink instead,
+  # or trap utf-8 conversion exception, etc.
+  def envrcpt_bytes(self,*b):
+    s = (v.decode(encoding='utf-8',errors='surrogateescape') for v in b)
+    return self.envrcpt(fld,*s)
   ## Called when the SMTP client says RCPT TO. Called by the
   # <a href="milter_api/xxfi_envrcpt.html">
   # xxfi_envrcpt</a> callback.
@@ -713,12 +729,6 @@ def connect_callback(ctx,hostname,family,hostaddr,nr_mask=P_NR_CONN):
     m._setctx(ctx)
   return m.connect(hostname,family,hostaddr)
 
-## @private
-# @brief check str/bytes decorator and invoke header method.
-def header_callback(ctx,fld,val):
-  m = ctx.getpriv()
-  return m.header_bytes(fld,val)
-
 ## @private
 # @brief Disconnect milterContext and call close method.
 def close_callback(ctx):
@@ -782,12 +792,14 @@ def runmilter(name,socketname,timeout = 0,rmsock=True):
   # parms, but then all existing users would have to include **kw to accept
   # arbitrary keywords without crashing.  We do provide envcallback and
   # dictfromlist to make parsing the ESMTP args convenient.
-  milter.set_envfrom_callback(lambda ctx,*str: ctx.getpriv().envfrom(*str))
-  milter.set_envrcpt_callback(lambda ctx,*str: ctx.getpriv().envrcpt(*str))
   if sys.version < '3.0.0':
+    milter.set_envfrom_callback(lambda ctx,*s: ctx.getpriv().envfrom(*s))
+    milter.set_envrcpt_callback(lambda ctx,*s: ctx.getpriv().envrcpt(*s))
     milter.set_header_callback(lambda ctx,f,v: ctx.getpriv().header(f,v))
   else:
-    milter.set_header_callback(header_callback)
+    milter.set_envfrom_callback(lambda ctx,*b: ctx.getpriv().envfrom_bytes(*b))
+    milter.set_envrcpt_callback(lambda ctx,*b: ctx.getpriv().envrcpt_bytes(*b))
+    milter.set_header_callback(lambda ctx,f,v: ctx.getpriv().header_bytes(f,v))
   milter.set_eoh_callback(lambda ctx: ctx.getpriv().eoh())
   milter.set_body_callback(lambda ctx,chunk: ctx.getpriv().body(chunk))
   milter.set_eom_callback(lambda ctx: ctx.getpriv().eom())
diff --git a/miltermodule.c b/miltermodule.c
index ffa1053..f7422e2 100644
--- a/miltermodule.c
+++ b/miltermodule.c
@@ -643,7 +643,7 @@ generic_env_wrapper(SMFICTX *ctx, PyObject*cb, char **argv) {
      /* There's some error checking performed in do_mkvalue() for a string */
      /* that's not currently done here - it probably should be */
 #if PY_MAJOR_VERSION >= 3
-     PyObject *o = PyUnicode_FromStringAndSize(argv[i], strlen(argv[i]));
+     PyObject *o = PyBytes_FromStringAndSize(argv[i], strlen(argv[i]));
 #else
      PyObject *o = PyString_FromStringAndSize(argv[i], strlen(argv[i]));
 #endif