From 4749f0ff987910a5472c2119df8e937e1da3afd8 Mon Sep 17 00:00:00 2001
From: "Stuart D. Gathman" <stuart@gathman.org>
Date: Tue, 27 Aug 2019 21:47:26 -0400
Subject: [PATCH] Change header callback to bytes, but default Milter to
 convert to str with surrogateescape.

---
 Milter/__init__.py |  7 +++++++
 Milter/testctx.py  | 17 +++++++++++++++--
 miltermodule.c     |  2 +-
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/Milter/__init__.py b/Milter/__init__.py
index 4ff23ce..16c1e73 100755
--- a/Milter/__init__.py
+++ b/Milter/__init__.py
@@ -701,6 +701,13 @@ 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()
+  s = val.decode(encoding='ascii',errors='surrogateescape')
+  return m.header(fld,s)
+
 ## @private
 # @brief Disconnect milterContext and call close method.
 def close_callback(ctx):
diff --git a/Milter/testctx.py b/Milter/testctx.py
index 05d3669..84a595c 100644
--- a/Milter/testctx.py
+++ b/Milter/testctx.py
@@ -15,6 +15,7 @@ except:
 import Milter
 from Milter import utils
 import mime
+import email
 
 ## Milter context for unit testing %milter applications.
 # A substitute for milter.milterContext that can be passed to
@@ -219,7 +220,20 @@ class TestCtx(object):
     return rc
 
   def _header(self,fld,val):
-    return self._priv.header(fld,val)
+    # email.message_from_binary_file uses surrogateescape to 
+    # preserve original bytes in unicode string for decoding errors.
+    # convert str or Header back to original bytes
+    if hasattr(val, '_chunks'):
+      # val is a Header object for invalid header values 
+      v = b''
+      for s,charset in val._chunks:
+        # recover the original bytes
+        b = s.encode(encoding='ascii',errors='surrogateescape')
+        v += b
+    else:
+      v = val.encode(encoding='ascii',errors='surrogateescape')
+    # invoke the Milter header_callback
+    return Milter.header_callback(self,fld,v)
 
   def _eoh(self):
     if self._protocol & Milter.P_NOEOH:
@@ -270,7 +284,6 @@ class TestCtx(object):
     if rc != Milter.CONTINUE: return rc
     # header
     for h,val in msg.items():
-      # val is a Header object for invalid header values 
       rc = self._header(h,val)
       if rc != Milter.CONTINUE: return rc
     # eoh
diff --git a/miltermodule.c b/miltermodule.c
index fb2e989..ba67953 100644
--- a/miltermodule.c
+++ b/miltermodule.c
@@ -674,7 +674,7 @@ milter_wrap_header(SMFICTX *ctx, char *headerf, char *headerv) {
    if (header_callback == NULL) return SMFIS_CONTINUE;
    c = _get_context(ctx);
    if (!c) return SMFIS_TEMPFAIL;
-   arglist = Py_BuildValue("(Oss)", c, headerf, headerv);
+   arglist = Py_BuildValue("(Oyy)", c, headerf, headerv);
    return _generic_wrapper(c, header_callback, arglist);
 }
 
-- 
GitLab