Add IP-based rules (#14963)

1 files changed, 32 insertions, 0 deletions

diff --git a/app/lib/fast_ip_map.rb b/app/lib/fast_ip_map.rb
new file mode 100644
index 000000000..ba30b45f3
--- /dev/null
+++ b/app/lib/fast_ip_map.rb
@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+class FastIpMap
+  MAX_IPV4_PREFIX = 32
+  MAX_IPV6_PREFIX = 128
+
+  # @param [Enumerable<IPAddr>] addresses
+  def initialize(addresses)
+    @fast_lookup = {}
+    @ranges      = []
+
+    # Hash look-up is faster but only works for exact matches, so we split
+    # exact addresses from non-exact ones
+    addresses.each do |address|
+      if (address.ipv4? && address.prefix == MAX_IPV4_PREFIX) || (address.ipv6? && address.prefix == MAX_IPV6_PREFIX)
+        @fast_lookup[address.to_s] = true
+      else
+        @ranges << address
+      end
+    end
+
+    # We're more likely to hit wider-reaching ranges when checking for
+    # inclusion, so make sure they're sorted first
+    @ranges.sort_by!(&:prefix)
+  end
+
+  # @param [IPAddr] address
+  # @return [Boolean]
+  def include?(address)
+    @fast_lookup[address.to_s] || @ranges.any? { |cidr| cidr.include?(address) }
+  end
+end