From c2a889f884c3e44933812f12d8388c09ab36da12 Mon Sep 17 00:00:00 2001
From: Kale Blankenship <kale@lemnisys.com>
Date: Sun, 11 Feb 2018 14:34:37 -0800
Subject: [PATCH] Restrict annotations to string/symbol or int/int64 keys.

---
 encode.go       | 27 +++++++++++++++++++++++++++
 marshal_test.go | 10 +++++-----
 types.go        | 38 ++++++++++++++++++++++++++++++++++----
 3 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/encode.go b/encode.go
index cbf04076..5ec00fc5 100644
--- a/encode.go
+++ b/encode.go
@@ -491,6 +491,33 @@ func writeMap(wr *buffer, m interface{}) error {
 				return err
 			}
 		}
+	case Annotations:
+		pairs = len(m) * 2
+		for key, val := range m {
+			switch key := key.(type) {
+			case string:
+				err := symbol(key).marshal(wr)
+				if err != nil {
+					return err
+				}
+			case symbol:
+				err := key.marshal(wr)
+				if err != nil {
+					return err
+				}
+			case int64:
+				writeInt64(wr, key)
+			case int:
+				writeInt64(wr, int64(key))
+			default:
+				return errorErrorf("unsupported Annotations key type %T", key)
+			}
+
+			err := marshal(wr, val)
+			if err != nil {
+				return err
+			}
+		}
 	default:
 		return errorErrorf("unsupported map type %T", m)
 	}
diff --git a/marshal_test.go b/marshal_test.go
index 99479d1a..933051e7 100644
--- a/marshal_test.go
+++ b/marshal_test.go
@@ -447,11 +447,11 @@ var (
 				FirstAcquirer: true,
 				DeliveryCount: 32,
 			},
-			DeliveryAnnotations: map[interface{}]interface{}{
-				int8(42): "answer",
+			DeliveryAnnotations: Annotations{
+				int64(42): "answer",
 			},
-			Annotations: map[interface{}]interface{}{
-				int8(42): "answer",
+			Annotations: Annotations{
+				int64(42): "answer",
 			},
 			Properties: &MessageProperties{
 				MessageID:          "yo",
@@ -473,7 +473,7 @@ var (
 			},
 			Data:  []byte("A nice little data payload."),
 			Value: uint8(42),
-			Footer: map[interface{}]interface{}{
+			Footer: Annotations{
 				"hash": []uint8{0, 1, 2, 34, 5, 6, 7, 8, 9, 0},
 			},
 		},
diff --git a/types.go b/types.go
index 2f5dac3e..253ba3ff 100644
--- a/types.go
+++ b/types.go
@@ -1648,7 +1648,7 @@ type Message struct {
 	// The delivery-annotations section is used for delivery-specific non-standard
 	// properties at the head of the message. Delivery annotations convey information
 	// from the sending peer to the receiving peer.
-	DeliveryAnnotations map[interface{}]interface{}
+	DeliveryAnnotations Annotations
 	// If the recipient does not understand the annotation it cannot be acted upon
 	// and its effects (such as any implied propagation) cannot be acted upon.
 	// Annotations might be specific to one implementation, or common to multiple
@@ -1664,7 +1664,7 @@ type Message struct {
 
 	// The message-annotations section is used for properties of the message which
 	// are aimed at the infrastructure.
-	Annotations map[interface{}]interface{}
+	Annotations Annotations
 	// The message-annotations section is used for properties of the message which
 	// are aimed at the infrastructure and SHOULD be propagated across every
 	// delivery step. Message annotations convey information about the message.
@@ -1712,14 +1712,44 @@ type Message struct {
 	// can only be calculated or evaluated once the whole bare message has been
 	// constructed or seen (for example message hashes, HMACs, signatures and
 	// encryption details).
-	Footer map[interface{}]interface{}
-	// TODO: implement custom type with validation
+	Footer Annotations
 
 	receiver *Receiver  // Receiver the message was received from
 	id       deliveryID // used when sending disposition
 	settled  bool       // whether transfer was settled by sender
 }
 
+// Annotations keys must be of type string, int, or int64.
+//
+// String keys are encoded as AMQP Symbols.
+type Annotations map[interface{}]interface{}
+
+func (a Annotations) marshal(wr *buffer) error {
+	return writeMap(wr, a)
+}
+
+func (a *Annotations) unmarshal(r *buffer) error {
+	_, count, err := readMapHeader(r)
+	if err != nil {
+		return err
+	}
+
+	m := make(Annotations, count/2)
+	for i := uint32(0); i < count; i += 2 {
+		key, err := readAny(r)
+		if err != nil {
+			return err
+		}
+		value, err := readAny(r)
+		if err != nil {
+			return err
+		}
+		m[key] = value
+	}
+	*a = m
+	return nil
+}
+
 // Accept notifies the server that the message has been
 // accepted and does not require redelivery.
 func (m *Message) Accept() {