Fixed multiple extension criteria bug

Change-Id: I57157b83b605e7315c3849743a931f270e8f86a8

@@ -28,8 +28,11 @@ import org.onosproject.net.DeviceId;
 import org.onosproject.net.PortNumber;
 import org.onosproject.net.flow.criteria.Criteria;
 import org.onosproject.net.flow.criteria.Criterion;
+import org.onosproject.net.flow.criteria.ExtensionCriterion;
 import org.onosproject.net.flow.criteria.ExtensionSelector;
+import org.onosproject.net.flow.criteria.ExtensionSelectorType;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
@@ -38,27 +41,38 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.TreeSet;
 
+import static org.onosproject.net.flow.criteria.Criterion.Type.EXTENSION;
+
 /**
  * Default traffic selector implementation.
  */
 public final class DefaultTrafficSelector implements TrafficSelector {
 
     private static final Comparator<? super Criterion> TYPE_COMPARATOR =
-            (c1, c2) -> c1.type().compareTo(c2.type());
+            (c1, c2) -> {
+                if (c1.type() == EXTENSION && c2.type() == EXTENSION) {
+                    return ((ExtensionCriterion) c1).extensionSelector().type().toInt()
+                            - ((ExtensionCriterion) c2).extensionSelector().type().toInt();
+                } else {
+                    return c1.type().compareTo(c2.type());
+                }
+            };
 
     private final Set<Criterion> criteria;
 
     private static final TrafficSelector EMPTY
-            = new DefaultTrafficSelector(Collections.emptySet());
+            = new DefaultTrafficSelector(Collections.emptySet(), Collections.emptySet());
 
     /**
      * Creates a new traffic selector with the specified criteria.
      *
-     * @param criteria criteria
+     * @param criteria    criteria
+     * @param extCriteria extension criteria
      */
-    private DefaultTrafficSelector(Set<Criterion> criteria) {
+    private DefaultTrafficSelector(Collection<Criterion> criteria, Collection<Criterion> extCriteria) {
         TreeSet<Criterion> elements = new TreeSet<>(TYPE_COMPARATOR);
         elements.addAll(criteria);
+        elements.addAll(extCriteria);
         this.criteria = ImmutableSet.copyOf(elements);
     }
 
@@ -137,6 +151,7 @@ public final class DefaultTrafficSelector implements TrafficSelector {
     public static final class Builder implements TrafficSelector.Builder {
 
         private final Map<Criterion.Type, Criterion> selector = new HashMap<>();
+        private final Map<ExtensionSelectorType, Criterion> extSelector = new HashMap<>();
 
         private Builder() {
         }
@@ -149,7 +164,11 @@ public final class DefaultTrafficSelector implements TrafficSelector {
 
         @Override
         public Builder add(Criterion criterion) {
-            selector.put(criterion.type(), criterion);
+            if (criterion.type() == EXTENSION) {
+                extSelector.put(((ExtensionCriterion) criterion).extensionSelector().type(), criterion);
+            } else {
+                selector.put(criterion.type(), criterion);
+            }
             return this;
         }
 
@@ -371,7 +390,7 @@ public final class DefaultTrafficSelector implements TrafficSelector {
 
         @Override
         public TrafficSelector build() {
-            return new DefaultTrafficSelector(ImmutableSet.copyOf(selector.values()));
+            return new DefaultTrafficSelector(selector.values(), extSelector.values());
         }
     }
 }

@@ -73,6 +73,15 @@ public class ExtensionSelectorType {
         this.type = type;
     }
 
+    /**
+     * Returns the integer value associated with this type.
+     *
+     * @return an integer value
+     */
+    public int toInt() {
+        return this.type;
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(type);

@@ -32,6 +32,7 @@ import org.onlab.packet.MplsLabel;
 import org.onlab.packet.TpPort;
 import org.onlab.packet.VlanId;
 import org.onosproject.net.ChannelSpacing;
+import org.onosproject.net.DeviceId;
 import org.onosproject.net.GridType;
 import org.onosproject.net.OchSignal;
 import org.onosproject.net.PortNumber;
@@ -39,10 +40,14 @@ import org.onosproject.net.flow.criteria.Criteria;
 import org.onosproject.net.flow.criteria.Criterion;
 
 import com.google.common.testing.EqualsTester;
+import org.onosproject.net.flow.criteria.ExtensionSelector;
+import org.onosproject.net.flow.criteria.ExtensionSelectorType;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.equalTo;
 import static org.onlab.junit.ImmutableClassChecker.assertThatClassIsImmutable;
 import static org.onosproject.net.flow.criteria.Criterion.Type;
 
@@ -289,5 +294,36 @@ public class DefaultTrafficSelectorTest {
         selector = DefaultTrafficSelector.builder()
                 .add(Criteria.matchLambda(new OchSignal(GridType.DWDM, ChannelSpacing.CHL_100GHZ, 1, 1))).build();
         assertThat(selector, hasCriterionWithType(Type.OCH_SIGID));
+
+        selector = DefaultTrafficSelector.builder()
+                .matchEthDst(macValue)
+                .extension(new MockExtensionSelector(1), DeviceId.NONE)
+                .extension(new MockExtensionSelector(2), DeviceId.NONE)
+                .build();
+        assertThat(selector.criteria().size(), is(equalTo(3)));
+    }
+
+    private class MockExtensionSelector extends AbstractExtension implements ExtensionSelector {
+
+        ExtensionSelectorType type;
+
+        MockExtensionSelector(int typeInt) {
+            this.type = new ExtensionSelectorType(typeInt);
+        }
+
+        @Override
+        public ExtensionSelectorType type() {
+            return type;
+        }
+
+        @Override
+        public byte[] serialize() {
+            return new byte[0];
+        }
+
+        @Override
+        public void deserialize(byte[] data) {
+
+        }
     }
 }