Parsing of custom message attributes

smack-core/src/main/java/org/jivesoftware/smack/packet/IQ.java

@@ -116,7 +116,7 @@ public final String getChildElementNamespace() {
     public final XmlStringBuilder toXML() {
         XmlStringBuilder buf = new XmlStringBuilder();
         buf.halfOpenElement(IQ_ELEMENT);
-        addCommonAttributes(buf);
+        addAttributes(buf);
         if (type == null) {
             buf.attribute("type", "get");
         }

smack-core/src/main/java/org/jivesoftware/smack/packet/Message.java

@@ -422,7 +422,7 @@ else if (language == null) {
     public XmlStringBuilder toXML() {
         XmlStringBuilder buf = new XmlStringBuilder();
         buf.halfOpenElement(ELEMENT);
-        addCommonAttributes(buf);
+        addAttributes(buf);
         buf.optAttribute("type", type);
         buf.rightAngleBracket();
 

smack-core/src/main/java/org/jivesoftware/smack/packet/Presence.java

@@ -224,7 +224,7 @@ public void setMode(Mode mode) {
     public XmlStringBuilder toXML() {
         XmlStringBuilder buf = new XmlStringBuilder();
         buf.halfOpenElement(ELEMENT);
-        addCommonAttributes(buf);
+        addAttributes(buf);
         if (type != Type.available) {
             buf.attribute("type", type);
         }

smack-core/src/main/java/org/jivesoftware/smack/packet/Stanza.java

@@ -28,9 +28,15 @@
 import org.jxmpp.stringprep.XmppStringprepException;
 import org.jxmpp.util.XmppStringUtils;
 
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * Base class for XMPP Stanzas, which are called Stanza(/Packet) in older versions of Smack (i.e. < 4.1).
@@ -51,16 +57,20 @@ public abstract class Stanza implements TopLevelStreamElement {
 
     public static final String TEXT = "text";
     public static final String ITEM = "item";
+    public static final Set<String> REGISTERED_ATTRIBUTES = new HashSet<String>(Arrays.asList("xml:lang", "id", "to", "from", "type"));
 
     protected static final String DEFAULT_LANGUAGE =
             java.util.Locale.getDefault().getLanguage().toLowerCase(Locale.US);
 
+    private static boolean customAttributesEnabled = false;
+
     private final MultiMap<String, ExtensionElement> packetExtensions = new MultiMap<>();
 
     private String id = null;
     private Jid to;
     private Jid from;
     private XMPPError error = null;
+    private LinkedHashMap<String, String> customAttributes;
 
     /**
      * Optional value of the 'xml:lang' attribute of the outermost element of
@@ -92,6 +102,10 @@ protected Stanza(Stanza p) {
         for (ExtensionElement pe : p.getExtensions()) {
             addExtension(pe);
         }
+
+        if (p.customAttributes != null) {
+            customAttributes = new LinkedHashMap<>(p.customAttributes);
+        }
     }
 
     /**
@@ -458,6 +472,152 @@ public ExtensionElement removeExtension(ExtensionElement extension)  {
         return removeExtension(extension.getElementName(), extension.getNamespace());
     }
 
+    /**
+     * Returns <tt>true</tt> if custom attributes management is enabled.
+     * @return <tt>true</tt> if custom attributes management is enabled.
+     */
+    public static boolean customAttributesEnabled() {
+        return customAttributesEnabled;
+    }
+
+    /**
+     * Enables the possibility to add custom attributes to the stanza.
+     * <br/><b>Enabling this feature breaks the XMPP standard and may induce unexpected behaviors!
+     * <br/>If your use-case allow it, prefer using ExtensionElement instead, and if not, be sure to triple-check the result.</b>
+     */
+    public static void enableCustomAttributes() {
+        customAttributesEnabled = true;
+    }
+
+    private static void requireCustomAttributesEnabled() {
+        if (!customAttributesEnabled()) {
+            throw new IllegalStateException(
+                    "You need to enable this feature first by calling enableCustomAttributes. " +
+                    "Ensure you are aware of the consequences of doing so.");
+        }
+    }
+
+    private void requireCustomAttributes() {
+        if (customAttributes == null) {
+            customAttributes = new LinkedHashMap<>();
+        }
+    }
+
+    private void failIfRegistered(String attributeName) {
+        if (REGISTERED_ATTRIBUTES.contains(attributeName)) {
+            throw new IllegalArgumentException(attributeName + " is a registered attribute.");
+        }
+    }
+
+    /**
+     * Replaces all custom attributes of the stanza.
+     * @param attributes a map (name/value) of all the custom attributes of this stanza. Should not be {@null}
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public void setCustomAttributes(Map<String, String> attributes) {
+        requireCustomAttributesEnabled();
+        assert attributes != null;
+        for (String attr : attributes.keySet()) {
+            failIfRegistered(attr);
+        }
+        requireCustomAttributes();
+
+        customAttributes.clear();
+        customAttributes.putAll(attributes);
+    }
+
+    /**
+     * Adds a custom attribute to the stanza. If a custom attribute with the same name is already present, its value is updated.
+     * @param attributeName the custom attribute name.
+     * @param attributeValue the custom attribute value.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public void setCustomAttribute(String attributeName, String attributeValue) {
+        requireCustomAttributesEnabled();
+        requireNotNullOrEmpty(attributeName, "attributeName must not be null or empty");
+        failIfRegistered(attributeName);
+        requireCustomAttributes();
+
+        customAttributes.put(attributeName, attributeValue);
+    }
+
+    /**
+     * Removes all custom attributes of the stanza.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public void removeCustomAttributes() {
+        requireCustomAttributesEnabled();
+        requireCustomAttributes();
+
+        customAttributes.clear();
+    }
+
+    /**
+     * Removes a custom attribute of the stanza if it exists.
+     * @param attributeName the custom attribute name.
+     * @return the removed custom attribute value or <tt>null</tt> if is was not present.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public String removeCustomAttribute(String attributeName) {
+        requireCustomAttributesEnabled();
+        requireNotNullOrEmpty(attributeName, "attributeName must not be null or empty");
+        requireCustomAttributes();
+
+        return customAttributes.remove(attributeName);
+    }
+
+    /**
+     * Returns <tt>true</tt> if this stanza contains custom attributes.
+     * @return <tt>true</tt> if this stanza contains custom attributes.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public boolean hasCustomAttributes() {
+        requireCustomAttributesEnabled();
+        requireCustomAttributes();
+
+        return !customAttributes.isEmpty();
+    }
+
+    /**
+     * Returns <tt>true</tt> if this stanza contains a custom attribute with this name.
+     * @param attributeName the custom attribute name.
+     * @return <tt>true</tt> if this stanza contains a custom attribute with this name.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public boolean hasCustomAttribute(String attributeName) {
+        requireCustomAttributesEnabled();
+        requireNotNullOrEmpty(attributeName, "attributeName must not be null or empty");
+        requireCustomAttributes();
+
+        return customAttributes.containsKey(attributeName);
+    }
+
+    /**
+     * Returns a map (name/value) of all the custom attributes of this stanza.
+     * @return an unmodifiable map (name/value) of all the custom attributes of this stanza.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public Map<String, String> getCustomAttributes() {
+        requireCustomAttributesEnabled();
+        requireCustomAttributes();
+
+        return Collections.unmodifiableMap(customAttributes);
+    }
+
+    /**
+     * Returns the value of a custom attribute, or {@code null} if this stanza does not contain any custom attribute with this name.
+     * @param attributeName the custom attribute name.
+     * @return the value of a custom attribute, or {@code null} if this stanza does not contain any custom attribute with this name.
+     * @throws IllegalStateException if custom attributes management is not enabled.
+     */
+    public String getCustomAttribute(String attributeName) {
+        requireCustomAttributesEnabled();
+        requireNotNullOrEmpty(attributeName, "attributeName must not be null or empty");
+        requireCustomAttributes();
+
+        return customAttributes.get(attributeName);
+    }
+
     @Override
     // NOTE When Smack is using Java 8, then this method should be moved in Element as "Default Method".
     public String toString() {
@@ -494,11 +654,19 @@ public static String getDefaultLanguage() {
      *
      * @param xml
      */
-    protected void addCommonAttributes(XmlStringBuilder xml) {
+    protected void addAttributes(XmlStringBuilder xml) {
         xml.optAttribute("to", getTo());
         xml.optAttribute("from", getFrom());
         xml.optAttribute("id", getStanzaId());
         xml.xmllangAttribute(getLanguage());
+
+        if(customAttributes == null) {
+            return;
+        }
+
+        for (Map.Entry<String, String> entry : customAttributes.entrySet()) {
+            xml.attribute(entry.getKey(), entry.getValue());
+        }
     }
 
     /**

smack-core/src/main/java/org/jivesoftware/smack/util/PacketParserUtils.java

@@ -236,6 +236,10 @@ public static Message parseMessage(XmlPullParser parser)
             defaultLanguage = Stanza.getDefaultLanguage();
         }
 
+        if(Stanza.customAttributesEnabled()) {
+            message.setCustomAttributes(ParserUtils.getCustomAttributes(parser));
+        }
+
         // Parse sub-elements. We include extra logic to make sure the values
         // are only read once. This is because it's possible for the names to appear
         // in arbitrary sub-elements.
@@ -538,6 +542,10 @@ public static Presence parsePresence(XmlPullParser parser)
         // CHECKSTYLE:ON
         }
 
+        if(Stanza.customAttributesEnabled()) {
+            presence.setCustomAttributes(ParserUtils.getCustomAttributes(parser));
+        }
+
         // Parse sub-elements
         outerloop: while (true) {
             int eventType = parser.next();
@@ -612,6 +620,11 @@ public static IQ parseIQ(XmlPullParser parser) throws Exception {
         final Jid from = ParserUtils.getJidAttribute(parser, "from");
         final IQ.Type type = IQ.Type.fromString(parser.getAttributeValue("", "type"));
 
+        Map<String, String> customAttributes = null;
+        if (Stanza.customAttributesEnabled()) {
+            customAttributes = ParserUtils.getCustomAttributes(parser);
+        }
+
         outerloop: while (true) {
             int eventType = parser.next();
 
@@ -668,6 +681,9 @@ public static IQ parseIQ(XmlPullParser parser) throws Exception {
         iqPacket.setFrom(from);
         iqPacket.setType(type);
         iqPacket.setError(error);
+        if (Stanza.customAttributesEnabled()) {
+            iqPacket.setCustomAttributes(customAttributes);
+        }
 
         return iqPacket;
     }

smack-core/src/main/java/org/jivesoftware/smack/util/ParserUtils.java

@@ -21,8 +21,11 @@
 import java.net.URISyntaxException;
 import java.text.ParseException;
 import java.util.Date;
+import java.util.LinkedHashMap;
 import java.util.Locale;
+import java.util.Map;
 
+import org.jivesoftware.smack.packet.Stanza;
 import org.jxmpp.jid.EntityBareJid;
 import org.jxmpp.jid.Jid;
 import org.jxmpp.jid.impl.JidCreate;
@@ -194,4 +197,16 @@ public static URI getUriFromNextText(XmlPullParser parser) throws XmlPullParserE
         return uri;
     }
 
+    public static Map<String, String> getCustomAttributes(XmlPullParser parser) {
+        Map<String, String> customAttributes = new LinkedHashMap<>();
+        for (int i = 0; i < parser.getAttributeCount(); i++) {
+            String attributeName = parser.getAttributeName(i);
+            boolean isLangAttribute = "lang".equals(attributeName) && "xml".equals(parser.getAttributePrefix(i));
+            if (!Stanza.REGISTERED_ATTRIBUTES.contains(attributeName) && !isLangAttribute) {
+                String attributeValue = parser.getAttributeValue(i);
+                customAttributes.put(attributeName, attributeValue);
+            }
+        }
+        return customAttributes;
+    }
 }

smack-core/src/test/java/org/jivesoftware/smack/util/PacketParserUtilsTest.java

@@ -40,7 +40,9 @@
 import org.jivesoftware.smack.sasl.packet.SaslStreamElements.SASLFailure;
 import org.jivesoftware.smack.test.util.TestUtils;
 import org.jivesoftware.smack.test.util.XmlUnitUtils;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.xml.sax.SAXException;
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
@@ -49,6 +51,9 @@
 
 public class PacketParserUtilsTest {
 
+	@Rule
+    public ExpectedException expectedException = ExpectedException.none();
+
     private static Properties outputProperties = new Properties();
     {
         outputProperties.put(javax.xml.transform.OutputKeys.OMIT_XML_DECLARATION, "yes");
@@ -816,6 +821,41 @@ public void parseElementMultipleNamespace()
         assertXMLEqual(stanza, result.toString());
     }
 
+    @Test
+    public void parseMessageWithCustomAttributes() throws FactoryConfigurationError, Exception {
+        final String customAttrName = "customAttrName";
+        final String customAttrValue = "customAttrValue";
+
+        final String stanzaString = XMLBuilder.create("message")
+            .a("from", "romeo@montague.lit/orchard")
+            .a("to", "juliet@capulet.lit/balcony")
+            .a("id", "zid615d9")
+            .a("type", "chat")
+            .a(customAttrName, customAttrValue)
+            .a("xml:lang", Stanza.getDefaultLanguage())
+            .e("body")
+                .t("This is a test of the custom attributes parsing in message stanza")
+            .asString(outputProperties);
+
+        Stanza stanza = PacketParserUtils.parseStanza(PacketParserUtils.getParserFor(stanzaString));
+
+        // Should crash because feature was not enabled
+        expectedException.expect(IllegalStateException.class);
+        stanza.getCustomAttributes();
+
+        // Custom attributes not parsed because feature was enabled after the parsing
+        Stanza.enableCustomAttributes();
+        assertFalse(stanza.hasCustomAttributes());
+
+        // Parse again for getting custom attributes
+        stanza = PacketParserUtils.parseMessage(PacketParserUtils.getParserFor(stanzaString));
+        assertTrue(stanza.hasCustomAttributes());
+        assertTrue(stanza.getCustomAttributes().size() == 1);
+        assertTrue(stanza.hasCustomAttribute(customAttrName));
+        assertTrue(stanza.getCustomAttribute(customAttrName).equals(customAttrValue));
+        assertXMLEqual(stanzaString, stanza.toXML().toString());
+    }
+
     @Test
     public void parseSASLFailureSimple() throws FactoryConfigurationError, SAXException, IOException,
                     TransformerException, ParserConfigurationException, XmlPullParserException {