git.iko.soy/typeable/oauth2-proxy
on ea261ca help sign in
typeable/oauth2-proxy/pkg/ip/net_set.go
112 lines 2.9 kB history raw
1package ip
3import (
4 "fmt"
5 "net"
6)
8// Fast lookup table for intersection of a single IP address within a collection of CIDR networks.
9//
10// Supports 4-byte (IPv4) and 16-byte (IPv6) networks.
11//
12// Provides O(1) best-case, O(log(n)) worst-case performance.
13// In practice netmasks included will generally only be of standard lengths:
14// - /8, /16, /24, and /32 for IPv4
15// - /64 and /128 for IPv6.
16// As a result, typical lookup times will lean closer to best-case rather than worst-case even when most of the internet
17// is included.
18type NetSet struct {
19 ip4NetMaps []ipNetMap
20 ip6NetMaps []ipNetMap
23// Create a new NetSet with all of the provided networks.
24func NewNetSet() *NetSet {
25 return &NetSet{
26 ip4NetMaps: make([]ipNetMap, 0),
27 ip6NetMaps: make([]ipNetMap, 0),
28 }
31// Check if `ip` is in the set, true if within the set otherwise false.
32func (w *NetSet) Has(ip net.IP) bool {
33 netMaps := w.getNetMaps(ip)
35 // Check all ipNetMaps for intersection with `ip`.
36 for _, netMap := range *netMaps {
37 if netMap.has(ip) {
38 return true
39 }
40 }
41 return false
44// Add an CIDR network to the set.
45func (w *NetSet) AddIPNet(ipNet net.IPNet) {
46 netMaps := w.getNetMaps(ipNet.IP)
48 // Determine the size / number of ones in the CIDR network mask.
49 ones, _ := ipNet.Mask.Size()
51 var netMap *ipNetMap
53 // Search for the ipNetMap containing networks with the same number of ones.
54 for i := 0; len(*netMaps) > i; i++ {
55 if netMapOnes, _ := (*netMaps)[i].mask.Size(); netMapOnes == ones {
56 netMap = &(*netMaps)[i]
57 break
58 }
59 }
61 // Create a new ipNetMap if none with this number of ones have been created yet.
62 if netMap == nil {
63 netMap = &ipNetMap{
64 mask: ipNet.Mask,
65 ips: make(map[string]bool),
66 }
67 *netMaps = append(*netMaps, *netMap)
68 // Recurse once now that there exists an netMap.
69 w.AddIPNet(ipNet)
70 return
71 }
73 // Add the IP to the ipNetMap.
74 netMap.ips[ipNet.IP.String()] = true
77// Get the appropriate array of networks for the given IP version.
78func (w *NetSet) getNetMaps(ip net.IP) (netMaps *[]ipNetMap) {
79 switch {
80 case ip.To4() != nil:
81 netMaps = &w.ip4NetMaps
82 case ip.To16() != nil:
83 netMaps = &w.ip6NetMaps
84 default:
85 panic(fmt.Sprintf("IP (%s) is neither 4-byte nor 16-byte?", ip.String()))
86 }
88 return netMaps
91// Hash-set of CIDR networks with the same mask size.
92type ipNetMap struct {
93 mask net.IPMask
94 ips map[string]bool
97// Check if the IP is in any of the CIDR networks contained in this map.
98func (m ipNetMap) has(ip net.IP) bool {
99 // Apply the mask to the IP to remove any irrelevant bits in the IP.
100 ipMasked := ip.Mask(m.mask)
101 if ipMasked == nil {
102 panic(fmt.Sprintf(
103 "Mismatch in net.IPMask and net.IP protocol version, cannot apply mask %s to %s",
104 m.mask.String(), ip.String()))
105 }
107 // Check if the masked IP is the same as any of the networks.
108 if _, ok := m.ips[ipMasked.String()]; ok {
109 return true
110 }
111 return false