From f18745f052e3c235acd47519b1664a5e1cfd5f1d Mon Sep 17 00:00:00 2001
From: Klimov Paul <klimov.paul@gmail.com>
Date: Mon, 29 Aug 2016 16:06:14 +0300
Subject: [PATCH] Added ability to specify custom mail header at
 `yii\swiftmailer\Message`

---
 CHANGELOG.md          |  1 +
 Message.php           | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++
 tests/MessageTest.php | 35 +++++++++++++++++++++++++++++++-
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d33cede..3b63c72 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@ Yii Framework 2 swiftmailer extension Change Log
 2.0.6 under development
 -----------------------
 
+- Enh #6: Added ability to specify custom mail header at `yii\swiftmailer\Message` (klimov-paul)
 - Enh #27: Added ability to specify message signature (klimov-paul)
 
 
diff --git a/Message.php b/Message.php
index 3972967..4884843 100644
--- a/Message.php
+++ b/Message.php
@@ -399,6 +399,62 @@ class Message extends BaseMessage
     }
 
     /**
+     * Adds custom header value to the message.
+     * Several invocations of this method with the same name will add multiple header values.
+     * @param string $name header name.
+     * @param string $value header value.
+     * @return $this self reference.
+     * @since 2.0.6
+     */
+    public function addHeader($name, $value)
+    {
+        $this->getSwiftMessage()->getHeaders()->addTextHeader($name, $value);
+        return $this;
+    }
+
+    /**
+     * Sets custom header value to the message.
+     * @param string $name header name.
+     * @param string|array $value header value or values.
+     * @return $this self reference.
+     * @since 2.0.6
+     */
+    public function setHeader($name, $value)
+    {
+        $headerSet = $this->getSwiftMessage()->getHeaders();
+
+        if ($headerSet->has($name)) {
+            $headerSet->remove($name);
+        }
+
+        foreach ((array)$value as $v) {
+            $headerSet->addTextHeader($name, $v);
+        }
+
+        return $this;
+    }
+
+    /**
+     * Returns all values for the specified header.
+     * @param string $name header name.
+     * @return array header values list.
+     * @since 2.0.6
+     */
+    public function getHeader($name)
+    {
+        $headerSet = $this->getSwiftMessage()->getHeaders();
+        if (!$headerSet->has($name)) {
+            return [];
+        }
+
+        $headers = [];
+        foreach ($headerSet->getAll($name) as $header) {
+            $headers[] = $header->getValue();
+        }
+        return $headers;
+    }
+
+    /**
      * @inheritdoc
      */
     public function toString()
diff --git a/tests/MessageTest.php b/tests/MessageTest.php
index bd3543f..37833b8 100644
--- a/tests/MessageTest.php
+++ b/tests/MessageTest.php
@@ -176,7 +176,7 @@ class MessageTest extends TestCase
     /**
      * @depends testGetSwiftMessage
      */
-    public function testSetupHeaders()
+    public function testSetupHeaderShortcuts()
     {
         $charset = 'utf-16';
         $subject = 'Test Subject';
@@ -458,4 +458,37 @@ U41eAdnQ3dDGzUNedIJkSh6Z0A4VMZIEOag9hPNYqQXZBQgfobvPKw==
         $content = $message->toString();
         $this->assertEquals(2, substr_count($content, $charset), 'Wrong charset for alternative body override.');
     }
+
+    /**
+     * @depends testGetSwiftMessage
+     */
+    public function testSetupHeaders()
+    {
+        $messageString = $this->createTestMessage()
+            ->addHeader('Some', 'foo')
+            ->addHeader('Multiple', 'value1')
+            ->addHeader('Multiple', 'value2')
+            ->toString();
+
+        $this->assertContains('Some: foo', $messageString, 'Unable to add header!');
+        $this->assertContains('Multiple: value1', $messageString, 'First value of multiple header lost!');
+        $this->assertContains('Multiple: value2', $messageString, 'Second value of multiple header lost!');
+
+        $messageString = $this->createTestMessage()
+            ->setHeader('Some', 'foo')
+            ->setHeader('Some', 'override')
+            ->setHeader('Multiple', ['value1', 'value2'])
+            ->toString();
+
+        $this->assertContains('Some: override', $messageString, 'Unable to set header!');
+        $this->assertNotContains('Some: foo', $messageString, 'Unable to override header!');
+        $this->assertContains('Multiple: value1', $messageString, 'First value of multiple header lost!');
+        $this->assertContains('Multiple: value2', $messageString, 'Second value of multiple header lost!');
+
+        $message = $this->createTestMessage();
+        $message->setHeader('Some', 'foo');
+        $this->assertEquals(['foo'], $message->getHeader('Some'));
+        $message->setHeader('Multiple', ['value1', 'value2']);
+        $this->assertEquals(['value1', 'value2'], $message->getHeader('Multiple'));
+    }
 }