From abfc37491cecb0a052d54c3bf0525ac140e94f7e Mon Sep 17 00:00:00 2001
From: R David Murray <rdmurray@bitdance.com>
Date: Tue, 29 May 2012 09:14:44 -0400
Subject: [PATCH] #10839: raise an error on add of duplicate unique headers in
 new email policies

This feature was supposed to be part of the initial email6 checkin, but it got
lost in my big refactoring.

In this patch I'm not providing an easy way to turn off the errors, but they
only happen when a header is added programmatically, and it is almost never
the right thing to do to allow the duplicate to be added.  An application that
needs to add duplicates of unique headers can create a policy subclass to
allow it.
---
 Doc/library/email.policy.rst | 25 +++++++++++++++++++++++++
 Lib/email/_policybase.py     | 19 +++++++++++++++++++
 Lib/email/message.py         | 10 ++++++++++
 Lib/email/policy.py          |  8 ++++++++
 4 files changed, 62 insertions(+)

diff --git a/Doc/library/email.policy.rst b/Doc/library/email.policy.rst
index 9ab949eb788..31b13c15f10 100644
--- a/Doc/library/email.policy.rst
+++ b/Doc/library/email.policy.rst
@@ -196,6 +196,25 @@ added matters.  To illustrate::
       custom ``Message`` objects) should also provide such an attribute,
       otherwise defects in parsed messages will raise unexpected errors.
 
+   .. method:: header_max_count(name)
+
+      Return the maximum allowed number of headers named *name*.
+
+      Called when a header is added to a :class:`~email.message.Message`
+      object.  If the returned value is not ``0`` or ``None``, and there are
+      already a number of headers with the name *name* equal to the value
+      returned, a :exc:`ValueError` is raised.
+
+      Because the default behavior of ``Message.__setitem__`` is to append the
+      value to the list of headers, it is easy to create duplicate headers
+      without realizing it.  This method allows certain headers to be limited
+      in the number of instances of that header that may be added to a
+      ``Message`` programmatically.  (The limit is not observed by the parser,
+      which will faithfully produce as many headers as exist in the message
+      being parsed.)
+
+      The default implementation returns ``None`` for all header names.
+
    .. method:: header_source_parse(sourcelines)
 
       The email package calls this method with a list of strings, each string
@@ -366,6 +385,12 @@ added matters.  To illustrate::
    The class provides the following concrete implementations of the abstract
    methods of :class:`Policy`:
 
+   .. method:: header_max_count(name)
+
+      Returns the value of the
+      :attr:`~email.headerregistry.BaseHeader.max_count` attribute of the
+      specialized class used to represent the header with the given name.
+
    .. method:: header_source_parse(sourcelines)
 
       The implementation of this method is the same as that for the
diff --git a/Lib/email/_policybase.py b/Lib/email/_policybase.py
index 6bc298b945c..d5e8df99404 100644
--- a/Lib/email/_policybase.py
+++ b/Lib/email/_policybase.py
@@ -194,6 +194,25 @@ def register_defect(self, obj, defect):
         """
         obj.defects.append(defect)
 
+    def header_max_count(self, name):
+        """Return the maximum allowed number of headers named 'name'.
+
+        Called when a header is added to a Message object.  If the returned
+        value is not 0 or None, and there are already a number of headers with
+        the name 'name' equal to the value returned, a ValueError is raised.
+
+        Because the default behavior of Message's __setitem__ is to append the
+        value to the list of headers, it is easy to create duplicate headers
+        without realizing it.  This method allows certain headers to be limited
+        in the number of instances of that header that may be added to a
+        Message programmatically.  (The limit is not observed by the parser,
+        which will faithfully produce as many headers as exist in the message
+        being parsed.)
+
+        The default implementation returns None for all header names.
+        """
+        return None
+
     @abc.abstractmethod
     def header_source_parse(self, sourcelines):
         """Given a list of linesep terminated strings constituting the lines of
diff --git a/Lib/email/message.py b/Lib/email/message.py
index 62b82b79c10..9b06207d148 100644
--- a/Lib/email/message.py
+++ b/Lib/email/message.py
@@ -346,6 +346,16 @@ def __setitem__(self, name, val):
         Note: this does not overwrite an existing header with the same field
         name.  Use __delitem__() first to delete any existing headers.
         """
+        max_count = self.policy.header_max_count(name)
+        if max_count:
+            lname = name.lower()
+            found = 0
+            for k, v in self._headers:
+                if k.lower() == lname:
+                    found += 1
+                    if found >= max_count:
+                        raise ValueError("There may be at most {} {} headers "
+                                         "in a message".format(max_count, name))
         self._headers.append(self.policy.header_store_parse(name, val))
 
     def __delitem__(self, name):
diff --git a/Lib/email/policy.py b/Lib/email/policy.py
index 47ed66b642a..bfffb457dcd 100644
--- a/Lib/email/policy.py
+++ b/Lib/email/policy.py
@@ -69,6 +69,14 @@ def __init__(self, **kw):
             object.__setattr__(self, 'header_factory', HeaderRegistry())
         super().__init__(**kw)
 
+    def header_max_count(self, name):
+        """+
+        The implementation for this class returns the max_count attribute from
+        the specialized header class that would be used to construct a header
+        of type 'name'.
+        """
+        return self.header_factory[name].max_count
+
     # The logic of the next three methods is chosen such that it is possible to
     # switch a Message object between a Compat32 policy and a policy derived
     # from this class and have the results stay consistent.  This allows a