RoutesQuestion: Add more match criteria

projects/allinone/src/test/java/org/batfish/question/routes/Bgpv4RoutesTest.java

@@ -52,7 +52,8 @@ public void setup() throws IOException {
   @Test
   public void testBestRoutes() {
     RoutesQuestion question =
-        new RoutesQuestion(null, NODE1, null, null, BgpRouteStatus.BEST.name(), RibProtocol.BGP);
+        new RoutesQuestion(
+            null, NODE1, null, null, BgpRouteStatus.BEST.name(), RibProtocol.BGP, null);
     TableAnswerElement answer =
         (TableAnswerElement) new RoutesAnswerer(question, _batfish).answer(_batfish.getSnapshot());
     assertThat(
@@ -66,7 +67,8 @@ public void testBestRoutes() {
   @Test
   public void testBackupRoutes() {
     RoutesQuestion question =
-        new RoutesQuestion(null, NODE1, null, null, BgpRouteStatus.BACKUP.name(), RibProtocol.BGP);
+        new RoutesQuestion(
+            null, NODE1, null, null, BgpRouteStatus.BACKUP.name(), RibProtocol.BGP, null);
     TableAnswerElement answer =
         (TableAnswerElement) new RoutesAnswerer(question, _batfish).answer(_batfish.getSnapshot());
     assertThat(

projects/batfish-common-protocol/src/test/java/org/batfish/datamodel/MockRib.java

@@ -125,7 +125,7 @@ public Set<AnnotatedRoute<AbstractRoute>> longestPrefixMatch(
       Ip address,
       int maxPrefixLength,
       ResolutionRestriction<AnnotatedRoute<AbstractRoute>> restriction) {
-    throw new UnsupportedOperationException();
+    return _longestPrefixMatchResults.getOrDefault(address, ImmutableSet.of());
   }
 
   @Override

projects/question/src/main/java/org/batfish/question/routes/RoutesAnswerer.java

@@ -123,55 +123,41 @@ public AnswerElement answer(NetworkSnapshot snapshot) {
       return new StringAnswerElement(ERROR_NO_MATCHING_VRFS);
     }
 
-    boolean bgpMultipathBest = expandedBgpRouteStatuses.contains(BEST);
-    boolean bgpBackup = expandedBgpRouteStatuses.contains(BACKUP);
     List<Row> rows = new ArrayList<>();
 
     switch (question.getRib()) {
       case BGP:
-        if (bgpBackup) {
-          rows.addAll(
-              getBgpRibRoutes(
-                  dp.getBgpBackupRoutes(),
-                  matchingVrfsByNode,
-                  network,
-                  protocolSpec,
-                  ImmutableSet.of(BACKUP)));
-        }
-        if (bgpMultipathBest) {
-          rows.addAll(
-              getBgpRibRoutes(
-                  dp.getBgpRoutes(),
-                  matchingVrfsByNode,
-                  network,
-                  protocolSpec,
-                  ImmutableSet.of(BEST)));
-        }
+        rows.addAll(
+            getBgpRibRoutes(
+                dp.getBgpRoutes(),
+                dp.getBgpBackupRoutes(),
+                matchingVrfsByNode,
+                network,
+                protocolSpec,
+                expandedBgpRouteStatuses,
+                question.getPrefixMatchType()));
         rows.sort(BGP_COMPARATOR);
         break;
       case EVPN:
-        if (bgpBackup) {
-          rows.addAll(
-              getEvpnRoutes(
-                  dp.getEvpnBackupRoutes(),
-                  matchingVrfsByNode,
-                  network,
-                  protocolSpec,
-                  ImmutableSet.of(BACKUP)));
-        }
-        if (bgpMultipathBest) {
-          rows.addAll(
-              getEvpnRoutes(
-                  dp.getEvpnRoutes(),
-                  matchingVrfsByNode,
-                  network,
-                  protocolSpec,
-                  ImmutableSet.of(BEST)));
-        }
+        rows.addAll(
+            getEvpnRoutes(
+                dp.getEvpnRoutes(),
+                dp.getEvpnBackupRoutes(),
+                matchingVrfsByNode,
+                network,
+                protocolSpec,
+                ImmutableSet.of(BACKUP),
+                question.getPrefixMatchType()));
         rows.sort(BGP_COMPARATOR);
         break;
       case MAIN:
-        rows.addAll(getMainRibRoutes(dp.getRibs(), matchingVrfsByNode, network, protocolSpec));
+        rows.addAll(
+            getMainRibRoutes(
+                dp.getRibs(),
+                matchingVrfsByNode,
+                network,
+                protocolSpec,
+                question.getPrefixMatchType()));
         rows.sort(MAIN_RIB_COMPARATOR);
         break;
       default:

projects/question/src/main/java/org/batfish/question/routes/RoutesAnswererUtil.java

@@ -3,6 +3,10 @@
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Comparator.comparingInt;
+import static org.batfish.datamodel.ResolutionRestriction.alwaysTrue;
+import static org.batfish.datamodel.questions.BgpRouteStatus.BACKUP;
+import static org.batfish.datamodel.questions.BgpRouteStatus.BEST;
 import static org.batfish.datamodel.table.TableDiff.COL_BASE_PREFIX;
 import static org.batfish.datamodel.table.TableDiff.COL_DELTA_PREFIX;
 import static org.batfish.question.routes.RoutesAnswerer.COL_ADMIN_DISTANCE;
@@ -54,6 +58,7 @@
 import java.util.TreeSet;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import org.batfish.common.BatfishException;
 import org.batfish.datamodel.AbstractRoute;
@@ -72,6 +77,7 @@
 import org.batfish.datamodel.table.Row.RowBuilder;
 import org.batfish.datamodel.table.TableDiff;
 import org.batfish.question.routes.DiffRoutesOutput.KeyPresenceStatus;
+import org.batfish.question.routes.RoutesQuestion.PrefixMatchType;
 import org.batfish.question.routes.RoutesQuestion.RibProtocol;
 import org.batfish.specifier.RoutingProtocolSpecifier;
 
@@ -145,90 +151,206 @@ static String computeNextHopNode(
    *     Bgpv4Route}s are to be selected
    * @param network {@link Prefix} of the network used to filter the routes
    * @param protocolSpec {@link RoutingProtocolSpecifier} used to filter the routes
+   * @param prefixMatchType {@link PrefixMatchType} used to select which prefixes are reported
    * @return {@link Multiset} of {@link Row}s representing the routes
    */
   static <T extends AbstractRouteDecorator> Multiset<Row> getMainRibRoutes(
       SortedMap<String, SortedMap<String, GenericRib<T>>> ribs,
       Multimap<String, String> matchingVrfsByNode,
       @Nullable Prefix network,
-      RoutingProtocolSpecifier protocolSpec) {
+      RoutingProtocolSpecifier protocolSpec,
+      PrefixMatchType prefixMatchType) {
     Multiset<Row> rows = HashMultiset.create();
     Map<String, ColumnMetadata> columnMetadataMap =
         getTableMetadata(RibProtocol.MAIN).toColumnMap();
     matchingVrfsByNode.forEach(
         (hostname, vrfName) ->
             Optional.ofNullable(ribs.getOrDefault(hostname, ImmutableSortedMap.of()).get(vrfName))
-                .map(GenericRib::getRoutes)
-                .orElse(ImmutableSet.of())
-                .stream()
-                .filter(
-                    route ->
-                        (network == null || network.equals(route.getNetwork()))
-                            && protocolSpec.getProtocols().contains(route.getProtocol()))
+                .map(rib -> getMatchingPrefixRoutes(prefixMatchType, network, rib))
+                .orElse(Stream.empty())
+                .filter(route -> protocolSpec.getProtocols().contains(route.getProtocol()))
                 .forEach(
                     route ->
                         rows.add(abstractRouteToRow(hostname, vrfName, route, columnMetadataMap))));
     return rows;
   }
 
   /**
-   * Filters a {@link Table} of {@link Bgpv4Route}s to produce a {@link Multiset} of rows
+   * Given the prefixMatchType and network (user input), returns routes from the {@code rib} that
+   * match
+   */
+  @VisibleForTesting
+  static <T extends AbstractRouteDecorator> Stream<AbstractRoute> getMatchingPrefixRoutes(
+      PrefixMatchType prefixMatchType, @Nullable Prefix network, GenericRib<T> rib) {
+    if (network == null) {
+      // everything matches if there is not user input
+      return rib.getRoutes().stream();
+    }
+    if (prefixMatchType == PrefixMatchType.LONGEST_PREFIX_MATCH) {
+      return rib
+          .longestPrefixMatch(network.getStartIp(), network.getPrefixLength(), alwaysTrue())
+          .stream()
+          .map(AbstractRouteDecorator::getAbstractRoute);
+    }
+    return rib.getRoutes().stream()
+        .filter(r -> prefixMatches(prefixMatchType, network, r.getNetwork()));
+  }
+
+  @VisibleForTesting
+  static boolean prefixMatches(
+      PrefixMatchType prefixMatchType, Prefix inputNetwork, Prefix routeNetwork) {
+    switch (prefixMatchType) {
+      case EXACT:
+        return inputNetwork.equals(routeNetwork);
+      case LONGER_PREFIXES:
+        return inputNetwork.containsPrefix(routeNetwork);
+      case SHORTER_PREFIXES:
+        return routeNetwork.containsPrefix(inputNetwork);
+      case LONGEST_PREFIX_MATCH: // handled separately in the caller
+        throw new IllegalArgumentException("Illegal PrefixMatchType " + prefixMatchType);
+      default:
+        throw new UnsupportedOperationException(
+            "Unimplemented prefix match type " + prefixMatchType);
+    }
+  }
+
+  /**
+   * Filters {@link Table} of BEST and BACKUP {@link Bgpv4Route}s to produce a {@link Multiset} of
+   * rows.
    *
    * @param bgpRoutes {@link Table} of all {@link Bgpv4Route}s
    * @param matchingVrfsByNode {@link Multimap} of vrfs grouped by node from which {@link
    *     Bgpv4Route}s are to be selected
    * @param network {@link Prefix} of the network used to filter the routes
    * @param protocolSpec {@link RoutingProtocolSpecifier} used to filter the {@link Bgpv4Route}s
-   * @param statuses BGP route statuses that correspond to routes in {@code bgpRoutes}.
+   * @param routeStatuses BGP route statuses that correspond to routes in {@code bgpRoutes}.
+   * @param prefixMatchType
    * @return {@link Multiset} of {@link Row}s representing the routes
    */
   static Multiset<Row> getBgpRibRoutes(
       Table<String, String, Set<Bgpv4Route>> bgpRoutes,
+      Table<String, String, Set<Bgpv4Route>> bgpBackupRoutes,
       Multimap<String, String> matchingVrfsByNode,
       @Nullable Prefix network,
       RoutingProtocolSpecifier protocolSpec,
-      Set<BgpRouteStatus> statuses) {
+      Set<BgpRouteStatus> routeStatuses,
+      PrefixMatchType prefixMatchType) {
     Multiset<Row> rows = HashMultiset.create();
     Map<String, ColumnMetadata> columnMetadataMap = getTableMetadata(RibProtocol.BGP).toColumnMap();
     matchingVrfsByNode.forEach(
         (hostname, vrfName) ->
-            firstNonNull(bgpRoutes.get(hostname, vrfName), ImmutableSet.<Bgpv4Route>of()).stream()
-                .filter(
-                    route ->
-                        (network == null || network.equals(route.getNetwork()))
-                            && protocolSpec.getProtocols().contains(route.getProtocol()))
+            getMatchingRoutes(
+                    firstNonNull(bgpRoutes.get(hostname, vrfName), ImmutableSet.of()),
+                    firstNonNull(bgpBackupRoutes.get(hostname, vrfName), ImmutableSet.of()),
+                    network,
+                    routeStatuses,
+                    prefixMatchType)
                 .forEach(
-                    route ->
-                        rows.add(
-                            bgpRouteToRow(hostname, vrfName, route, statuses, columnMetadataMap))));
+                    (status, routeStream) ->
+                        routeStream
+                            .filter(r -> protocolSpec.getProtocols().contains(r.getProtocol()))
+                            .forEach(
+                                route ->
+                                    rows.add(
+                                        bgpRouteToRow(
+                                            hostname,
+                                            vrfName,
+                                            route,
+                                            ImmutableSet.of(status),
+                                            columnMetadataMap)))));
     return rows;
   }
 
   static Multiset<Row> getEvpnRoutes(
       Table<String, String, Set<EvpnRoute<?, ?>>> evpnRoutes,
+      Table<String, String, Set<EvpnRoute<?, ?>>> evpnBackupRoutes,
       Multimap<String, String> matchingVrfsByNode,
       @Nullable Prefix network,
       RoutingProtocolSpecifier protocolSpec,
-      Set<BgpRouteStatus> statuses) {
+      Set<BgpRouteStatus> routeStatuses,
+      PrefixMatchType prefixMatchType) {
     Multiset<Row> rows = HashMultiset.create();
     Map<String, ColumnMetadata> columnMetadataMap =
         getTableMetadata(RibProtocol.EVPN).toColumnMap();
     matchingVrfsByNode.forEach(
         (hostname, vrfName) ->
-            firstNonNull(evpnRoutes.get(hostname, vrfName), ImmutableSet.<EvpnRoute<?, ?>>of())
-                .stream()
-                .filter(
-                    route ->
-                        (network == null || network.equals(route.getNetwork()))
-                            && protocolSpec.getProtocols().contains(route.getProtocol()))
+            getMatchingRoutes(
+                    firstNonNull(evpnRoutes.get(hostname, vrfName), ImmutableSet.of()),
+                    firstNonNull(evpnBackupRoutes.get(hostname, vrfName), ImmutableSet.of()),
+                    network,
+                    routeStatuses,
+                    prefixMatchType)
                 .forEach(
-                    route ->
-                        rows.add(
-                            evpnRouteToRow(
-                                hostname, vrfName, route, statuses, columnMetadataMap))));
+                    (status, routeStream) ->
+                        routeStream
+                            .filter(r -> protocolSpec.getProtocols().contains(r.getProtocol()))
+                            .forEach(
+                                route ->
+                                    rows.add(
+                                        evpnRouteToRow(
+                                            hostname,
+                                            vrfName,
+                                            (EvpnRoute<?, ?>) route,
+                                            ImmutableSet.of(status),
+                                            columnMetadataMap)))));
     return rows;
   }
 
+  /**
+   * Filters best and backup routes to those that match the input network, route statuses, and
+   * prefix match type.
+   *
+   * <p>If the network is null, all routes are returned.
+   *
+   * <p>It the prefix match type is LONGEST_PREFIX_MATCH, the returned prefix is decided based on
+   * LPM on the best routes table.
+   */
+  static <T extends AbstractRouteDecorator> Map<BgpRouteStatus, Stream<T>> getMatchingRoutes(
+      Set<T> bestRoutes,
+      Set<T> backupRoutes,
+      @Nullable Prefix network,
+      Set<BgpRouteStatus> routeStatuses,
+      PrefixMatchType prefixMatchType) {
+    ImmutableMap.Builder<BgpRouteStatus, Stream<T>> routes = ImmutableMap.builder();
+    if (prefixMatchType == PrefixMatchType.LONGEST_PREFIX_MATCH && network != null) {
+      Optional<Prefix> lpmMatch = longestMatchingPrefix(network, bestRoutes);
+      if (lpmMatch.isPresent()) {
+        if (routeStatuses.contains(BEST)) {
+          routes.put(
+              BEST, getMatchingPrefixRoutes(bestRoutes, lpmMatch.get(), PrefixMatchType.EXACT));
+        }
+        if (routeStatuses.contains(BACKUP)) {
+          routes.put(
+              BACKUP, getMatchingPrefixRoutes(backupRoutes, lpmMatch.get(), PrefixMatchType.EXACT));
+        }
+      }
+      return routes.build();
+    }
+
+    if (routeStatuses.contains(BEST)) {
+      routes.put(BEST, getMatchingPrefixRoutes(bestRoutes, network, prefixMatchType));
+    }
+    if (routeStatuses.contains(BACKUP)) {
+      routes.put(BACKUP, getMatchingPrefixRoutes(backupRoutes, network, PrefixMatchType.EXACT));
+    }
+    return routes.build();
+  }
+
+  private static <T extends AbstractRouteDecorator> Stream<T> getMatchingPrefixRoutes(
+      Set<T> bgpRoutes, Prefix network, PrefixMatchType prefixMatchType) {
+    return bgpRoutes.stream()
+        .filter(r -> network == null || prefixMatches(prefixMatchType, network, r.getNetwork()));
+  }
+
+  @VisibleForTesting
+  static <T extends AbstractRouteDecorator> Optional<Prefix> longestMatchingPrefix(
+      Prefix network, Set<T> routes) {
+    return routes.stream()
+        .map(AbstractRouteDecorator::getNetwork)
+        .filter(prefix -> prefix.containsPrefix(network))
+        .max(comparingInt(Prefix::getPrefixLength));
+  }
+
   /**
    * Converts a {@link AbstractRoute} to a {@link Row}
    *
@@ -645,7 +767,7 @@ static Map<RouteRowKey, Map<RouteRowSecondaryKey, SortedSet<RouteRowAttribute>>>
     Map<BgpRouteStatus, Table<String, String, Set<Bgpv4Route>>> routesByStatus =
         new EnumMap<>(BgpRouteStatus.class);
     if (bgpBestRoutes != null) {
-      routesByStatus.put(BgpRouteStatus.BEST, bgpBestRoutes);
+      routesByStatus.put(BEST, bgpBestRoutes);
     }
     if (bgpBackupRoutes != null) {
       routesByStatus.put(BgpRouteStatus.BACKUP, bgpBackupRoutes);