Repository: FluuxIO/go-xmpp
Pull request #27 (2019-06-04 16:47:10)

Add support for parsing message extensions
- Add a registry structure and merge it with IQ Payload registry
- Add support for parsing receipts, chat markers and out of band image urls.
- Clean up the IQPayload interface / marker for IQ payload


Patch:
diff --git a/auth.go b/auth.go
index de38cac..36ce19e 100644
--- a/auth.go
+++ b/auth.go
@@ -105,14 +105,12 @@ type auth struct {
 }
 
 type BindBind struct {
+	IQPayload
 	XMLName  xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-bind bind"`
 	Resource string   `xml:"resource,omitempty"`
 	Jid      string   `xml:"jid,omitempty"`
 }
 
-func (*BindBind) IsIQPayload() {
-}
-
 // Session is obsolete in RFC 6121.
 // Added for compliance with RFC 3121.
 // Remove when ejabberd purely conforms to RFC 6121.
diff --git a/cmd/xmpp-check/xmpp-check.go b/cmd/xmpp-check/xmpp-check.go
index 6e7ceab..77af605 100644
--- a/cmd/xmpp-check/xmpp-check.go
+++ b/cmd/xmpp-check/xmpp-check.go
@@ -29,7 +29,6 @@ func main() {
 
 func runCheck(address, domain string) {
 	client, err := xmpp.NewChecker(address, domain)
-	// client, err := xmpp.NewChecker("mickael.m.in-app.io:5222", "mickael.m.in-app.io")
 
 	if err != nil {
 		log.Fatal("Error: ", err)
diff --git a/cmd/xmpp_jukebox/xmpp_jukebox.go b/cmd/xmpp_jukebox/xmpp_jukebox.go
index bd3e610..a60e164 100644
--- a/cmd/xmpp_jukebox/xmpp_jukebox.go
+++ b/cmd/xmpp_jukebox/xmpp_jukebox.go
@@ -12,8 +12,6 @@ import (
 	"github.com/processone/mpg123"
 	"github.com/processone/soundcloud"
 	"gosrc.io/xmpp"
-	"gosrc.io/xmpp/iot"
-	"gosrc.io/xmpp/pep"
 )
 
 // Get the actual song Stream URL from SoundCloud website song URL and play it with mpg123 player.
@@ -65,7 +63,7 @@ func processMessage(client *xmpp.Client, p *mpg123.Player, packet *xmpp.Message)
 func processIq(client *xmpp.Client, p *mpg123.Player, packet *xmpp.IQ) {
 	switch payload := packet.Payload[0].(type) {
 	// We support IOT Control IQ
-	case *iot.ControlSet:
+	case *xmpp.ControlSet:
 		var url string
 		for _, element := range payload.Fields {
 			if element.XMLName.Local == "string" && element.Name == "url" {
@@ -75,7 +73,7 @@ func processIq(client *xmpp.Client, p *mpg123.Player, packet *xmpp.IQ) {
 		}
 
 		playSCURL(p, url)
-		setResponse := new(iot.ControlSetResponse)
+		setResponse := new(xmpp.ControlSetResponse)
 		reply := xmpp.IQ{PacketAttrs: xmpp.PacketAttrs{To: packet.From, Type: "result", Id: packet.Id}, Payload: []xmpp.IQPayload{setResponse}}
 		_ = client.Send(reply)
 		// TODO add Soundclound artist / title retrieval
@@ -86,8 +84,11 @@ func processIq(client *xmpp.Client, p *mpg123.Player, packet *xmpp.IQ) {
 }
 
 func sendUserTune(client *xmpp.Client, artist string, title string) {
-	tune := pep.Tune{Artist: artist, Title: title}
-	_ = client.SendRaw(tune.XMPPFormat())
+	tune := xmpp.Tune{Artist: artist, Title: title}
+	iq := xmpp.NewIQ("set", "", "", "usertune-1", "en")
+	payload := xmpp.PubSub{Publish: xmpp.Publish{Node: "http://jabber.org/protocol/tune", Item: xmpp.Item{Tune: tune}}}
+	iq.AddPayload(&payload)
+	_ = client.Send(iq)
 }
 
 func playSCURL(p *mpg123.Player, rawURL string) {
diff --git a/iot/control.go b/iot_control.go
similarity index 74%
rename from iot/control.go
rename to iot_control.go
index 580f26a..795e5e3 100644
--- a/iot/control.go
+++ b/iot_control.go
@@ -1,15 +1,15 @@
-package iot // import "gosrc.io/xmpp/iot"
+package xmpp // import "gosrc.io/xmpp/iot"
 
-import "encoding/xml"
+import (
+	"encoding/xml"
+)
 
 type ControlSet struct {
+	IQPayload
 	XMLName xml.Name       `xml:"urn:xmpp:iot:control set"`
 	Fields  []ControlField `xml:",any"`
 }
 
-func (*ControlSet) IsIQPayload() {
-}
-
 type ControlGetForm struct {
 	XMLName xml.Name `xml:"urn:xmpp:iot:control getForm"`
 }
@@ -21,8 +21,6 @@ type ControlField struct {
 }
 
 type ControlSetResponse struct {
+	IQPayload
 	XMLName xml.Name `xml:"urn:xmpp:iot:control setResponse"`
 }
-
-func (*ControlSetResponse) IsIQPayload() {
-}
diff --git a/iq.go b/iq.go
index ecd6e49..bfc13e3 100644
--- a/iq.go
+++ b/iq.go
@@ -2,10 +2,7 @@ package xmpp // import "gosrc.io/xmpp"
 
 import (
 	"encoding/xml"
-	"reflect"
 	"strconv"
-
-	"gosrc.io/xmpp/iot"
 )
 
 /*
@@ -19,6 +16,7 @@ TODO support ability to put Raw payload inside IQ
 // presence or iq stanza.
 // It is intended to be added in the payload of the erroneous stanza.
 type Err struct {
+	IQPayload
 	XMLName xml.Name `xml:"error"`
 	Code    int      `xml:"code,attr,omitempty"`
 	Type    string   `xml:"type,attr,omitempty"`
@@ -26,8 +24,6 @@ type Err struct {
 	Text    string `xml:"urn:ietf:params:xml:ns:xmpp-stanzas text,omitempty"`
 }
 
-func (*Err) IsIQPayload() {}
-
 // UnmarshalXML implements custom parsing for IQs
 func (x *Err) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
 	x.XMLName = start.Name
@@ -208,21 +204,16 @@ func (iq *IQ) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
 		case xml.StartElement:
 			level++
 			if level <= 1 {
-				var elt interface{}
-				payloadType := tt.Name.Space + " " + tt.Name.Local
-				if payloadType := typeRegistry[payloadType]; payloadType != nil {
-					val := reflect.New(payloadType)
-					elt = val.Interface()
-				} else {
-					elt = new(Node)
-				}
-
-				if iqPl, ok := elt.(IQPayload); ok {
-					err = d.DecodeElement(elt, &tt)
+				if iqExt := typeRegistry.GetIQExtension(tt.Name); iqExt != nil {
+					// Decode payload extension
+					err = d.DecodeElement(iqExt, &tt)
 					if err != nil {
 						return err
 					}
-					iq.Payload = append(iq.Payload, iqPl)
+					iq.Payload = append(iq.Payload, iqExt)
+				} else {
+					// TODO: Fix me. We do nothing of that element here.
+					// elt = new(Node)
 				}
 			}
 
@@ -238,13 +229,12 @@ func (iq *IQ) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
 // ============================================================================
 // Generic IQ Payload
 
-type IQPayload interface {
-	IsIQPayload()
-}
+type IQPayload interface{}
 
 // Node is a generic structure to represent XML data. It is used to parse
 // unreferenced or custom stanza payload.
 type Node struct {
+	IQPayload
 	XMLName xml.Name
 	Attrs   []xml.Attr `xml:"-"`
 	Content string     `xml:",innerxml"`
@@ -283,8 +273,6 @@ func (n Node) MarshalXML(e *xml.Encoder, start xml.StartElement) (err error) {
 	return e.EncodeToken(xml.EndElement{Name: start.Name})
 }
 
-func (*Node) IsIQPayload() {}
-
 // ============================================================================
 // Disco
 
@@ -294,14 +282,13 @@ const (
 )
 
 type DiscoInfo struct {
+	IQPayload
 	XMLName  xml.Name  `xml:"http://jabber.org/protocol/disco#info query"`
 	Node     string    `xml:"node,attr,omitempty"`
 	Identity Identity  `xml:"identity"`
 	Features []Feature `xml:"feature"`
 }
 
-func (*DiscoInfo) IsIQPayload() {}
-
 type Identity struct {
 	XMLName  xml.Name `xml:"identity,omitempty"`
 	Name     string   `xml:"name,attr,omitempty"`
@@ -317,13 +304,12 @@ type Feature struct {
 // ============================================================================
 
 type DiscoItems struct {
+	IQPayload
 	XMLName xml.Name    `xml:"http://jabber.org/protocol/disco#items query"`
 	Node    string      `xml:"node,attr,omitempty"`
 	Items   []DiscoItem `xml:"item"`
 }
 
-func (*DiscoItems) IsIQPayload() {}
-
 type DiscoItem struct {
 	XMLName xml.Name `xml:"item"`
 	Name    string   `xml:"name,attr,omitempty"`
@@ -331,13 +317,9 @@ type DiscoItem struct {
 	Node    string   `xml:"node,attr,omitempty"`
 }
 
-// ============================================================================
-
-var typeRegistry = make(map[string]reflect.Type)
-
 func init() {
-	typeRegistry["http://jabber.org/protocol/disco#info query"] = reflect.TypeOf(DiscoInfo{})
-	typeRegistry["http://jabber.org/protocol/disco#items query"] = reflect.TypeOf(DiscoItems{})
-	typeRegistry["urn:ietf:params:xml:ns:xmpp-bind bind"] = reflect.TypeOf(BindBind{})
-	typeRegistry["urn:xmpp:iot:control set"] = reflect.TypeOf(iot.ControlSet{})
+	typeRegistry.MapExtension(PKTIQ, xml.Name{"http://jabber.org/protocol/disco#info", "query"}, DiscoInfo{})
+	typeRegistry.MapExtension(PKTIQ, xml.Name{"http://jabber.org/protocol/disco#items", "query"}, DiscoItems{})
+	typeRegistry.MapExtension(PKTIQ, xml.Name{"urn:ietf:params:xml:ns:xmpp-bind", "bind"}, BindBind{})
+	typeRegistry.MapExtension(PKTIQ, xml.Name{"urn:xmpp:iot:control", "set"}, ControlSet{})
 }
diff --git a/message.go b/message.go
index a9e108e..5b99348 100644
--- a/message.go
+++ b/message.go
@@ -11,10 +11,11 @@ import (
 type Message struct {
 	XMLName xml.Name `xml:"message"`
 	PacketAttrs
-	Subject string `xml:"subject,omitempty"`
-	Body    string `xml:"body,omitempty"`
-	Thread  string `xml:"thread,omitempty"`
-	Error   Err    `xml:"error,omitempty"`
+	Subject    string         `xml:"subject,omitempty"`
+	Body       string         `xml:"body,omitempty"`
+	Thread     string         `xml:"thread,omitempty"`
+	Error      Err            `xml:"error,omitempty"`
+	Extensions []MsgExtension `xml:",omitempty"`
 }
 
 func (Message) Name() string {
@@ -44,9 +45,76 @@ func (messageDecoder) decode(p *xml.Decoder, se xml.StartElement) (Message, erro
 	return packet, err
 }
 
+// TODO: Support missing element (thread, extensions) by using proper marshaller
 func (msg *Message) XMPPFormat() string {
 	return fmt.Sprintf("<message to='%s' type='chat' xml:lang='en'>"+
 		"<body>%s</body></message>",
 		msg.To,
 		xmlEscape(msg.Body))
 }
+
+// UnmarshalXML implements custom parsing for IQs
+func (msg *Message) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {
+	msg.XMLName = start.Name
+
+	// Extract packet attributes
+	for _, attr := range start.Attr {
+		if attr.Name.Local == "id" {
+			msg.Id = attr.Value
+		}
+		if attr.Name.Local == "type" {
+			msg.Type = attr.Value
+		}
+		if attr.Name.Local == "to" {
+			msg.To = attr.Value
+		}
+		if attr.Name.Local == "from" {
+			msg.From = attr.Value
+		}
+		if attr.Name.Local == "lang" {
+			msg.Lang = attr.Value
+		}
+	}
+
+	// decode inner elements
+	for {
+		t, err := d.Token()
+		if err != nil {
+			return err
+		}
+
+		switch tt := t.(type) {
+
+		case xml.StartElement:
+			if msgExt := typeRegistry.GetMsgExtension(tt.Name); msgExt != nil {
+				// Decode message extension
+				err = d.DecodeElement(msgExt, &tt)
+				if err != nil {
+					return err
+				}
+				msg.Extensions = append(msg.Extensions, msgExt)
+			} else {
+				// Decode standard message sub-elements
+				var err error
+				switch tt.Name.Local {
+				case "body":
+					err = d.DecodeElement(&msg.Body, &tt)
+				case "thread":
+					err = d.DecodeElement(&msg.Thread, &tt)
+				case "subject":
+					err = d.DecodeElement(&msg.Subject, &tt)
+				case "error":
+					err = d.DecodeElement(&msg.Error, &tt)
+				}
+				if err != nil {
+					return err
+				}
+			}
+
+		case xml.EndElement:
+			if tt == start.End() {
+				return nil
+			}
+		}
+	}
+}
diff --git a/msg_chat_markers.go b/msg_chat_markers.go
new file mode 100644
index 0000000..2940c1f
--- /dev/null
+++ b/msg_chat_markers.go
@@ -0,0 +1,24 @@
+package xmpp // import "gosrc.io/xmpp"
+
+import "encoding/xml"
+
+/*
+Support for:
+- XEP-0333 - Chat Markers: https://xmpp.org/extensions/xep-0333.html
+*/
+
+type Markable struct {
+	MsgExtension
+	XMLName xml.Name `xml:"urn:xmpp:chat-markers:0 markable"`
+}
+
+type MarkReceived struct {
+	MsgExtension
+	XMLName xml.Name `xml:"urn:xmpp:chat-markers:0 received"`
+	ID      string
+}
+
+func init() {
+	typeRegistry.MapExtension(PKTMessage, xml.Name{"urn:xmpp:chat-markers:0", "markable"}, Markable{})
+	typeRegistry.MapExtension(PKTMessage, xml.Name{"urn:xmpp:chat-markers:0", "received"}, MarkReceived{})
+}
diff --git a/msg_oob.go b/msg_oob.go
new file mode 100644
index 0000000..4b2188e
--- /dev/null
+++ b/msg_oob.go
@@ -0,0 +1,19 @@
+package xmpp // import "gosrc.io/xmpp"
+
+import "encoding/xml"
+
+/*
+Support for:
+- XEP-0066 - Out of Band Data: https://xmpp.org/extensions/xep-0066.html
+*/
+
+type OOB struct {
+	MsgExtension
+	XMLName xml.Name `xml:"jabber:x:oob x"`
+	URL     string   `xml:"url"`
+	Desc    string   `xml:"desc,omitempty"`
+}
+
+func init() {
+	typeRegistry.MapExtension(PKTMessage, xml.Name{"jabber:x:oob", "x"}, OOB{})
+}
diff --git a/msg_receipts.go b/msg_receipts.go
new file mode 100644
index 0000000..76958a8
--- /dev/null
+++ b/msg_receipts.go
@@ -0,0 +1,25 @@
+package xmpp // import "gosrc.io/xmpp"
+
+import "encoding/xml"
+
+/*
+Support for:
+- XEP-0184 - Message Delivery Receipts: https://xmpp.org/extensions/xep-0184.html
+*/
+
+// Used on outgoing message, to tell the recipient that you are requesting a message receipt / ack.
+type ReceiptRequest struct {
+	MsgExtension
+	XMLName xml.Name `xml:"urn:xmpp:receipts request"`
+}
+
+type ReceiptReceived struct {
+	MsgExtension
+	XMLName xml.Name `xml:"urn:xmpp:receipts received"`
+	ID      string
+}
+
+func init() {
+	typeRegistry.MapExtension(PKTMessage, xml.Name{"urn:xmpp:receipts", "request"}, ReceiptRequest{})
+	typeRegistry.MapExtension(PKTMessage, xml.Name{"urn:xmpp:receipts", "received"}, ReceiptReceived{})
+}
diff --git a/pep/user_tune.go b/pep.go
similarity index 72%
rename from pep/user_tune.go
rename to pep.go
index c1fe5ad..1870ca8 100644
--- a/pep/user_tune.go
+++ b/pep.go
@@ -1,29 +1,21 @@
-package pep // import "gosrc.io/xmpp/pep"
+package xmpp // import "gosrc.io/xmpp/pep"
 
 import (
 	"encoding/xml"
-
-	"gosrc.io/xmpp"
 )
 
-type iq struct {
-	XMLName          xml.Name `xml:"jabber:client iq"`
-	C                pubSub   // c for "contains"
-	xmpp.PacketAttrs          // Rename h for "header" ?
-}
-
-type pubSub struct {
+type PubSub struct {
 	XMLName xml.Name `xml:"http://jabber.org/protocol/pubsub pubsub"`
-	Publish publish
+	Publish Publish
 }
 
-type publish struct {
+type Publish struct {
 	XMLName xml.Name `xml:"publish"`
 	Node    string   `xml:"node,attr"`
-	Item    item
+	Item    Item
 }
 
-type item struct {
+type Item struct {
 	XMLName xml.Name `xml:"item"`
 	Tune    Tune
 }
@@ -67,11 +59,6 @@ type Tune struct {
 }
 */
 
-func (t *Tune) XMPPFormat() (s string) {
-	packet, _ := xml.Marshal(iq{PacketAttrs: xmpp.PacketAttrs{Id: "tunes", Type: "set"}, C: pubSub{Publish: publish{Node: "http://jabber.org/protocol/tune", Item: item{Tune: *t}}}})
-	return string(packet)
-}
-
 /*
 func (*Tune) XMPPFormat() string {
 	return fmt.Sprintf(
diff --git a/registry.go b/registry.go
new file mode 100644
index 0000000..49e6178
--- /dev/null
+++ b/registry.go
@@ -0,0 +1,105 @@
+package xmpp
+
+import (
+	"encoding/xml"
+	"reflect"
+	"sync"
+)
+
+type MsgExtension interface{}
+
+// The Registry for msg and IQ types is a global variable.
+// TODO: Move to the client init process to remove the dependency on a global variable.
+//   That should make it possible to be able to share the decoder.
+// TODO: Ensure that a client can add its own custom namespace to the registry (or overload existing ones).
+
+type packetType uint8
+
+const (
+	PKTPresence packetType = iota
+	PKTMessage
+	PKTIQ
+)
+
+var typeRegistry = newRegistry()
+
+// We store different registries per packet type and namespace.
+type registryKey struct {
+	packetType packetType
+	namespace  string
+}
+
+type registryForNamespace map[string]reflect.Type
+
+type registry struct {
+	// We store different registries per packet type and namespace.
+	msgTypes map[registryKey]registryForNamespace
+	// Handle concurrent access
+	msgTypesLock *sync.RWMutex
+}
+
+func newRegistry() *registry {
+	return &registry{
+		msgTypes:     make(map[registryKey]registryForNamespace),
+		msgTypesLock: &sync.RWMutex{},
+	}
+}
+
+// MapExtension stores extension type for packet payload.
+// The match is done per packetType (iq, message, or presence) and XML tag name.
+// You can use the alias "*" as local XML name to be able to match all unknown tag name for that
+// packet type and namespace.
+func (r *registry) MapExtension(pktType packetType, name xml.Name, extension MsgExtension) {
+	key := registryKey{pktType, name.Space}
+	r.msgTypesLock.RLock()
+	store := r.msgTypes[key]
+	r.msgTypesLock.RUnlock()
+
+	r.msgTypesLock.Lock()
+	defer r.msgTypesLock.Unlock()
+	if store == nil {
+		store = make(map[string]reflect.Type)
+	}
+	store[name.Local] = reflect.TypeOf(extension)
+	r.msgTypes[key] = store
+}
+
+// GetExtensionType returns extension type for packet payload, based on packet type and tag name.
+func (r *registry) GetExtensionType(pktType packetType, name xml.Name) reflect.Type {
+	key := registryKey{pktType, name.Space}
+
+	r.msgTypesLock.RLock()
+	defer r.msgTypesLock.RUnlock()
+	store := r.msgTypes[key]
+	result := store[name.Local]
+	if result == nil && name.Local != "*" {
+		return store["*"]
+	}
+	return result
+}
+
+// GetMsgExtension returns an instance of MsgExtension, by matching packet type and XML
+// tag name against the registry.
+func (r *registry) GetMsgExtension(name xml.Name) MsgExtension {
+	if extensionType := r.GetExtensionType(PKTMessage, name); extensionType != nil {
+		val := reflect.New(extensionType)
+		elt := val.Interface()
+		if msgExt, ok := elt.(MsgExtension); ok {
+			return msgExt
+		}
+	}
+	return nil
+}
+
+// GetIQExtension returns an instance of IQPayload, by matching packet type and XML
+// tag name against the registry.
+func (r *registry) GetIQExtension(name xml.Name) IQPayload {
+	if extensionType := r.GetExtensionType(PKTIQ, name); extensionType != nil {
+		val := reflect.New(extensionType)
+		elt := val.Interface()
+		if iqExt, ok := elt.(IQPayload); ok {
+			return iqExt
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/control_test.go b/control_test.go
index f52f663..2a2a027 100644
--- a/control_test.go
+++ b/control_test.go
@@ -3,8 +3,6 @@ package xmpp // import "gosrc.io/xmpp"
 import (
 	"encoding/xml"
 	"testing"
-
-	"gosrc.io/xmpp/iot"
 )
 
 func TestControlSet(t *testing.T) {
@@ -22,7 +20,7 @@ func TestControlSet(t *testing.T) {
 		t.Errorf("Unmarshal(%s) returned error", data)
 	}
 
-	if cs, ok := parsedIQ.Payload[0].(*iot.ControlSet); !ok {
+	if cs, ok := parsedIQ.Payload[0].(*ControlSet); !ok {
 		t.Errorf("Paylod is not an iot control set: %v", cs)
 	}
 }
diff --git a/message_test.go b/message_test.go
index 21ee882..fb30029 100644
--- a/message_test.go
+++ b/message_test.go
@@ -27,3 +27,23 @@ func TestGenerateMessage(t *testing.T) {
 		t.Errorf("non matching items\n%s", cmp.Diff(parsedMessage, message))
 	}
 }
+
+func TestDecodeError(t *testing.T) {
+	str := `<message from='juliet@capulet.com'
+         id='msg_1'
+         to='romeo@montague.lit'
+         type='error'>
+  <error type='cancel'>
+    <not-acceptable xmlns='urn:ietf:params:xml:ns:xmpp-stanzas'/>
+  </error>
+</message>`
+
+	parsedMessage := xmpp.Message{}
+	if err := xml.Unmarshal([]byte(str), &parsedMessage); err != nil {
+		t.Errorf("message error stanza unmarshall error: %v", err)
+		return
+	}
+	if parsedMessage.Error.Type != "cancel" {
+		t.Errorf("incorrect error type: %s", parsedMessage.Error.Type)
+	}
+}
diff --git a/msg_receipts_test.go b/msg_receipts_test.go
new file mode 100644
index 0000000..0db0c64
--- /dev/null
+++ b/msg_receipts_test.go
@@ -0,0 +1,42 @@
+package xmpp_test
+
+import (
+	"encoding/xml"
+	"testing"
+
+	"gosrc.io/xmpp"
+)
+
+func TestDecodeRequest(t *testing.T) {
+	str := `<message
+    from='northumberland@shakespeare.lit/westminster'
+    id='richard2-4.1.247'
+    to='kingrichard@royalty.england.lit/throne'>
+  <body>My lord, dispatch; read o'er these articles.</body>
+  <request xmlns='urn:xmpp:receipts'/>
+</message>`
+	parsedMessage := xmpp.Message{}
+	if err := xml.Unmarshal([]byte(str), &parsedMessage); err != nil {
+		t.Errorf("message receipt unmarshall error: %v", err)
+		return
+	}
+
+	if parsedMessage.Body != "My lord, dispatch; read o'er these articles." {
+		t.Errorf("Unexpected body: '%s'", parsedMessage.Body)
+	}
+
+	if len(parsedMessage.Extensions) < 1 {
+		t.Errorf("no extension found on parsed message")
+		return
+	}
+
+	switch ext := parsedMessage.Extensions[0].(type) {
+	case *xmpp.ReceiptRequest:
+		if ext.XMLName.Local != "request" {
+			t.Errorf("unexpected extension: %s:%s", ext.XMLName.Space, ext.XMLName.Local)
+		}
+	default:
+		t.Errorf("could not find receipts extension")
+	}
+
+}
diff --git a/registry_test.go b/registry_test.go
new file mode 100644
index 0000000..4b3ba81
--- /dev/null
+++ b/registry_test.go
@@ -0,0 +1,47 @@
+package xmpp // import "gosrc.io/xmpp"
+
+import (
+	"encoding/xml"
+	"reflect"
+	"testing"
+)
+
+func TestRegistry_RegisterMsgExt(t *testing.T) {
+	// Setup registry
+	typeRegistry := newRegistry()
+
+	// Register an element
+	name := xml.Name{Space: "urn:xmpp:receipts", Local: "request"}
+	typeRegistry.MapExtension(PKTMessage, name, ReceiptRequest{})
+
+	// Match that element
+	receipt := typeRegistry.GetMsgExtension(name)
+	if receipt == nil {
+		t.Error("cannot read element type from registry")
+		return
+	}
+
+	switch r := receipt.(type) {
+	case *ReceiptRequest:
+	default:
+		t.Errorf("Registry did not return expected type ReceiptRequest: %v", reflect.TypeOf(r))
+	}
+}
+
+func BenchmarkRegistryGet(b *testing.B) {
+	// Setup registry
+	typeRegistry := newRegistry()
+
+	// Register an element
+	name := xml.Name{Space: "urn:xmpp:receipts", Local: "request"}
+	typeRegistry.MapExtension(PKTMessage, name, ReceiptRequest{})
+
+	for i := 0; i < b.N; i++ {
+		// Match that element
+		receipt := typeRegistry.GetExtensionType(PKTMessage, name)
+		if receipt == nil {
+			b.Error("cannot read element type from registry")
+			return
+		}
+	}
+}