Adding support for IGMPv2

Change-Id: Iba3fbdfaed1b91cda8c6c0ef19fe69b5e5d7f900
diff --git a/utils/misc/src/main/java/org/onlab/packet/IGMP.java b/utils/misc/src/main/java/org/onlab/packet/IGMP.java
index 7244b02..42b71ba 100644
--- a/utils/misc/src/main/java/org/onlab/packet/IGMP.java
+++ b/utils/misc/src/main/java/org/onlab/packet/IGMP.java
@@ -15,14 +15,13 @@
  */
 package org.onlab.packet;
 
+import com.google.common.collect.ImmutableList;
 import org.slf4j.Logger;
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -32,24 +31,21 @@
 /**
  * Implements IGMP control packet format.
  */
-public class IGMP extends BasePacket {
-    private static final Logger log = getLogger(IGMP.class);
+public abstract class IGMP extends BasePacket {
+    protected static final Logger log = getLogger(IGMP.class);
 
     public static final byte TYPE_IGMPV3_MEMBERSHIP_QUERY = 0x11;
     public static final byte TYPE_IGMPV1_MEMBERSHIP_REPORT = 0x12;
     public static final byte TYPE_IGMPV2_MEMBERSHIP_REPORT = 0x16;
     public static final byte TYPE_IGMPV2_LEAVE_GROUP = 0x17;
     public static final byte TYPE_IGMPV3_MEMBERSHIP_REPORT = 0x22;
-    public static final Map<Byte, Deserializer<? extends IPacket>> PROTOCOL_DESERIALIZER_MAP = new HashMap<>();
-
-    public static final int MINIMUM_HEADER_LEN = 12;
 
     List<IGMPGroup> groups = new ArrayList<>();
 
     // Fields contained in the IGMP header
-    private byte igmpType;
-    private byte resField = 0;
-    private short checksum = 0;
+    protected byte igmpType;
+    protected byte resField = 0;
+    protected short checksum = 0;
 
     private byte[] unsupportTypeData;
 
@@ -97,12 +93,7 @@
      *
      * @param respCode the Maximum Response Code.
      */
-    public void setMaxRespCode(byte respCode) {
-        if (igmpType != IGMP.TYPE_IGMPV3_MEMBERSHIP_QUERY) {
-            log.debug("Requesting the max response code for an incorrect field: ");
-        }
-        this.resField = respCode;
-    }
+    public abstract void setMaxRespCode(byte respCode);
 
     /**
      * Get the list of IGMPGroups.  The group objects will be either IGMPQuery or IGMPMembership
@@ -112,7 +103,7 @@
      * @return The list of IGMP groups.
      */
     public List<IGMPGroup> getGroups() {
-        return groups;
+        return ImmutableList.copyOf(groups);
     }
 
     /**
@@ -121,32 +112,7 @@
      * @param group the IGMPGroup will be IGMPQuery or IGMPMembership depending on the message type.
      * @return true if group was valid and added, false otherwise.
      */
-    public boolean addGroup(IGMPGroup group) {
-        checkNotNull(group);
-        switch (this.igmpType) {
-            case TYPE_IGMPV3_MEMBERSHIP_QUERY:
-                if (group instanceof IGMPMembership) {
-                    return false;
-                }
-
-                if (group.sources.size() > 1) {
-                    return false;
-                }
-                break;
-
-            case TYPE_IGMPV3_MEMBERSHIP_REPORT:
-                if (group instanceof IGMPMembership) {
-                    return false;
-                }
-                break;
-
-            default:
-                log.debug("Warning no IGMP message type has been set");
-        }
-
-        this.groups.add(group);
-        return true;
-    }
+    public abstract boolean addGroup(IGMPGroup group);
 
     /**
      * Serialize this IGMP packet.  This will take care
@@ -169,27 +135,35 @@
         // Must calculate checksum
         bb.putShort((short) 0);
 
+        if (this instanceof IGMPv3) {
+            switch (this.igmpType) {
 
+                case IGMP.TYPE_IGMPV3_MEMBERSHIP_REPORT:
+                    // reserved
+                    bb.putShort((short) 0);
+                    // Number of groups
+                    bb.putShort((short) groups.size());
+                    // Fall through
 
-        switch (this.igmpType) {
+                case IGMP.TYPE_IGMPV3_MEMBERSHIP_QUERY:
 
-            case IGMP.TYPE_IGMPV3_MEMBERSHIP_REPORT:
-                // reserved
-                bb.putShort((short) 0);
-                // Number of groups
-                bb.putShort((short) groups.size());
-                // Fall through
+                    for (IGMPGroup grp : groups) {
+                        grp.serialize(bb);
+                    }
+                    break;
 
-            case IGMP.TYPE_IGMPV3_MEMBERSHIP_QUERY:
-
-                for (IGMPGroup grp : groups) {
-                    grp.serialize(bb);
-                }
-                break;
-
-            default:
-                bb.put(this.unsupportTypeData);
-                break;
+                default:
+                    bb.put(this.unsupportTypeData);
+                    break;
+            }
+        } else if (this instanceof IGMPv2) {
+            if (this.groups.isEmpty()) {
+                bb.putInt(0);
+            } else {
+                bb.putInt(groups.get(0).getGaddr().getIp4Address().toInt());
+            }
+        } else {
+            throw new UnsupportedOperationException();
         }
 
         int size = bb.position();
@@ -226,11 +200,11 @@
     public IPacket deserialize(final byte[] data, final int offset,
                                final int length) {
 
-        IGMP igmp = new IGMP();
+        final IGMP igmp;
         try {
             igmp = IGMP.deserializer().deserialize(data, offset, length);
         } catch (DeserializationException e) {
-            log.error(e.getStackTrace().toString());
+            log.error("Deserialization exception", e);
             return this;
         }
         this.igmpType = igmp.igmpType;
@@ -247,15 +221,29 @@
      */
     public static Deserializer<IGMP> deserializer() {
         return (data, offset, length) -> {
-            checkInput(data, offset, length, MINIMUM_HEADER_LEN);
+            checkInput(data, offset, length, IGMPv2.HEADER_LENGTH);
 
-            IGMP igmp = new IGMP();
+            // we will assume that this is IGMPv2 if the length is 8
+            boolean isV2 = length == IGMPv2.HEADER_LENGTH;
+
+            IGMP igmp = isV2 ? new IGMPv2() : new IGMPv3();
 
             final ByteBuffer bb = ByteBuffer.wrap(data, offset, length);
             igmp.igmpType = bb.get();
             igmp.resField = bb.get();
             igmp.checksum = bb.getShort();
 
+            if (isV2) {
+                igmp.addGroup(new IGMPQuery(IpAddress.valueOf(bb.getInt()), 0));
+                if (igmp.validChecksum()) {
+                    return igmp;
+                }
+                throw new DeserializationException("invalid checksum");
+            }
+
+            // second check for IGMPv3
+            checkInput(data, offset, length, IGMPv3.MINIMUM_HEADER_LEN);
+
             String msg;
 
             switch (igmp.igmpType) {
@@ -300,6 +288,13 @@
         };
     }
 
+    /**
+     * Validates the message's checksum.
+     *
+     * @return true if valid, false if not
+     */
+    protected abstract boolean validChecksum();
+
     /*
      * (non-Javadoc)
      *
@@ -363,4 +358,82 @@
                 .toString();
         // TODO: need to handle groups
     }
+
+    public static class IGMPv3 extends IGMP {
+        public static final int MINIMUM_HEADER_LEN = 12;
+
+        @Override
+        public void setMaxRespCode(byte respCode) {
+            if (igmpType != IGMP.TYPE_IGMPV3_MEMBERSHIP_QUERY) {
+                log.debug("Requesting the max response code for an incorrect field: ");
+            }
+            this.resField = respCode;
+        }
+
+        @Override
+        public boolean addGroup(IGMPGroup group) {
+            checkNotNull(group);
+            switch (this.igmpType) {
+                case TYPE_IGMPV3_MEMBERSHIP_QUERY:
+                    if (group instanceof IGMPMembership) {
+                        return false;
+                    }
+
+                    if (group.sources.size() > 1) {
+                        return false;
+                    }
+                    break;
+
+                case TYPE_IGMPV3_MEMBERSHIP_REPORT:
+                    if (group instanceof IGMPMembership) {
+                        return false;
+                    }
+                    break;
+
+                default:
+                    log.debug("Warning no IGMP message type has been set");
+            }
+
+            this.groups.add(group);
+            return true;
+        }
+
+        @Override
+        protected boolean validChecksum() {
+            return true; //FIXME
+        }
+    }
+
+    public static class IGMPv2 extends IGMP {
+        public static final int HEADER_LENGTH = 8;
+
+        @Override
+        public void setMaxRespCode(byte respCode) {
+            this.resField = respCode;
+        }
+
+        @Override
+        public boolean addGroup(IGMPGroup group) {
+            if (groups.isEmpty()) {
+                groups = ImmutableList.of(group);
+                return true;
+            }
+            return false;
+        }
+
+        @Override
+        protected boolean validChecksum() {
+            int accumulation = (((int) this.igmpType) & 0xff) << 8;
+            accumulation += ((int) this.resField) & 0xff;
+            if (!groups.isEmpty()) {
+                int ipaddr = groups.get(0).getGaddr().getIp4Address().toInt();
+                accumulation += (ipaddr >> 16) & 0xffff;
+                accumulation += ipaddr & 0xffff;
+            }
+            accumulation = (accumulation >> 16 & 0xffff)
+                    + (accumulation & 0xffff);
+            short checksum = (short) (~accumulation & 0xffff);
+            return checksum == this.checksum;
+        }
+    }
 }
diff --git a/utils/misc/src/test/java/org/onlab/packet/IGMPTest.java b/utils/misc/src/test/java/org/onlab/packet/IGMPTest.java
index 53b80cd..122afaa 100644
--- a/utils/misc/src/test/java/org/onlab/packet/IGMPTest.java
+++ b/utils/misc/src/test/java/org/onlab/packet/IGMPTest.java
@@ -18,6 +18,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertTrue;
 
 /**
@@ -44,7 +45,7 @@
         deserializer = IGMP.deserializer();
 
         // Create an IGMP Query object
-        igmpQuery = new IGMP();
+        igmpQuery = new IGMP.IGMPv3();
         igmpQuery.setIgmpType(IGMP.TYPE_IGMPV3_MEMBERSHIP_QUERY);
         igmpQuery.setMaxRespCode((byte) 0x7f);
         IGMPQuery q = new IGMPQuery(gaddr1, (byte) 0x7f);
@@ -54,7 +55,7 @@
         igmpQuery.groups.add(q);
 
         // Create an IGMP Membership Object
-        igmpMembership = new IGMP();
+        igmpMembership = new IGMP.IGMPv3();
         igmpMembership.setIgmpType(IGMP.TYPE_IGMPV3_MEMBERSHIP_REPORT);
         IGMPMembership g1 = new IGMPMembership(gaddr1);
         g1.addSource(saddr1);
@@ -94,6 +95,19 @@
         assertTrue(igmp.equals(igmpMembership));
     }
 
+    @Test
+    public void testIGMPv2() throws Exception {
+        IGMP igmp = new IGMP.IGMPv2();
+        igmp.setIgmpType((byte) 0x11);
+        igmp.setMaxRespCode((byte) 0x64);
+        igmp.addGroup(new IGMPQuery(IpAddress.valueOf(0), 0));
+
+        byte[] data = igmp.serialize();
+        assertEquals("Packet length is not 8 bytes", data.length, IGMP.IGMPv2.HEADER_LENGTH);
+        IGMP deserialized = deserializer.deserialize(data, 0, data.length);
+        assertTrue(igmp.equals(deserialized));
+    }
+
     /**
      * Tests toString.
      */