IP Address Set Operations with net/netip 4/7

Hey there! In our previous articles, we covered Addr, AddrPort, and Prefix types. Now let's tackle something more advanced - working with sets of IP addresses. While net/netip doesn't provide a built-in set type, we can build powerful abstractions for handling IP ranges and sets.

Why IP Sets?

IP sets are crucial when you need to:

• Manage lists of allowed/blocked IPs
• Handle network ACLs
• Work with IP ranges efficiently
• Perform set operations (union, intersection, difference)

Building an IP Set Type

Let's start by creating a flexible IP set implementation:

package ipset

import (
    "fmt"
    "net/netip"
    "sort"
)

// Range represents a continuous range of IP addresses
type Range struct {
    First netip.Addr
    Last  netip.Addr
}

// IPSet represents a set of IP addresses using ranges
type IPSet struct {
    ranges []Range
}

func NewIPSet() *IPSet {
    return &IPSet{}
}

// Add adds a single IP address to the set
func (s *IPSet) Add(ip netip.Addr) {
    // Convert single IP to range
    s.AddRange(Range{First: ip, Last: ip})
}

// AddRange adds an IP range to the set
func (s *IPSet) AddRange(r Range) {
    if r.First.Compare(r.Last) > 0 {
        r.First, r.Last = r.Last, r.First
    }

    // Find position to insert
    idx := sort.Search(len(s.ranges), func(i int) bool {
        return s.ranges[i].First.Compare(r.First) > 0
    })

    s.ranges = append(s.ranges, Range{}) // Make space
    copy(s.ranges[idx+1:], s.ranges[idx:])
    s.ranges[idx] = r

    // Merge overlapping ranges
    s.optimize()
}

// optimize merges overlapping or adjacent ranges
func (s *IPSet) optimize() {
    if len(s.ranges) <= 1 {
        return
    }

    // Sort ranges by start address
    sort.Slice(s.ranges, func(i, j int) bool {
        return s.ranges[i].First.Compare(s.ranges[j].First) < 0
    })

    // Merge overlapping ranges
    merged := []Range{s.ranges[0]}
    for i := 1; i < len(s.ranges); i++ {
        last := &merged[len(merged)-1]
        current := s.ranges[i]

        // Check if ranges overlap or are adjacent
        if last.Last.Compare(current.First) >= -1 {
            // Extend the range if needed
            if current.Last.Compare(last.Last) > 0 {
                last.Last = current.Last
            }
        } else {
            merged = append(merged, current)
        }
    }

    s.ranges = merged
}

// Contains checks if an IP address is in the set
func (s *IPSet) Contains(ip netip.Addr) bool {
    idx := sort.Search(len(s.ranges), func(i int) bool {
        return s.ranges[i].Last.Compare(ip) >= 0
    })

    if idx < len(s.ranges) {
        return s.ranges[idx].First.Compare(ip) <= 0
    }

    return false
}

Set Operations

Let's implement the basic set operations:

// Union returns a new set containing all IPs from both sets
func (s *IPSet) Union(other *IPSet) *IPSet {
    result := NewIPSet()

    // Add all ranges from both sets
    for _, r := range s.ranges {
        result.AddRange(r)
    }
    for _, r := range other.ranges {
        result.AddRange(r)
    }

    return result
}

// Intersection returns a new set containing IPs present in both sets
func (s *IPSet) Intersection(other *IPSet) *IPSet {
    result := NewIPSet()

    for _, r1 := range s.ranges {
        for _, r2 := range other.ranges {
            // Find overlapping range
            first := r1.First
            if r2.First.Compare(first) > 0 {
                first = r2.First
            }

            last := r1.Last
            if r2.Last.Compare(last) < 0 {
                last = r2.Last
            }

            // Add if there's an overlap
            if first.Compare(last) <= 0 {
                result.AddRange(Range{First: first, Last: last})
            }
        }
    }

    return result
}

// Difference returns a new set containing IPs in s but not in other
func (s *IPSet) Difference(other *IPSet) *IPSet {
    result := NewIPSet()

    for _, r1 := range s.ranges {
        current := []Range{{First: r1.First, Last: r1.Last}}

        for _, r2 := range other.ranges {
            var newRanges []Range

            for _, cr := range current {
                // Before r2
                if cr.First.Compare(r2.First) < 0 {
                    end := cr.Last
                    if r2.First.Compare(end) < 0 {
                        end = r2.First.Prev()
                    }
                    if cr.First.Compare(end) <= 0 {
                        newRanges = append(newRanges, Range{First: cr.First, Last: end})
                    }
                }

                // After r2
                if cr.Last.Compare(r2.Last) > 0 {
                    start := cr.First
                    if r2.Last.Compare(start) >= 0 {
                        start = r2.Last.Next()
                    }
                    if start.Compare(cr.Last) <= 0 {
                        newRanges = append(newRanges, Range{First: start, Last: cr.Last})
                    }
                }
            }

            current = newRanges
        }

        for _, r := range current {
            result.AddRange(r)
        }
    }

    return result
}

Real-World Applications

1. IP Address List Management

Here's a practical example of managing IP allowlists and blocklists:

type AccessManager struct {
    allowed *IPSet
    blocked *IPSet
}

func NewAccessManager() *AccessManager {
    return &AccessManager{
        allowed: NewIPSet(),
        blocked: NewIPSet(),
    }
}

func (am *AccessManager) AllowIP(ip netip.Addr) {
    am.allowed.Add(ip)
}

func (am *AccessManager) AllowCIDR(cidr string) error {
    prefix, err := netip.ParsePrefix(cidr)
    if err != nil {
        return fmt.Errorf("invalid CIDR: %w", err)
    }

    // Add range for the entire prefix
    first := prefix.Addr()
    last := first
    for i := 0; i < 1<<(32-prefix.Bits()); i++ {
        last = last.Next()
    }
    last = last.Prev() // Go back one to get last valid address

    am.allowed.AddRange(Range{First: first, Last: last})
    return nil
}

func (am *AccessManager) BlockIP(ip netip.Addr) {
    am.blocked.Add(ip)
}

func (am *AccessManager) IsAllowed(ip netip.Addr) bool {
    // Blocked IPs take precedence
    if am.blocked.Contains(ip) {
        return false
    }

    // If no allowlist is defined, allow by default
    if len(am.allowed.ranges) == 0 {
        return true
    }

    return am.allowed.Contains(ip)
}

2. Network Range Calculator

Here's a tool for working with IP ranges and subnets:

type NetworkRange struct {
    ipset *IPSet
}

func NewNetworkRange() *NetworkRange {
    return &NetworkRange{
        ipset: NewIPSet(),
    }
}

func (nr *NetworkRange) AddNetwork(cidr string) error {
    prefix, err := netip.ParsePrefix(cidr)
    if err != nil {
        return fmt.Errorf("invalid CIDR: %w", err)
    }

    // Calculate first and last IP in network
    first := prefix.Addr()
    last := first
    for i := 0; i < 1<<(32-prefix.Bits()); i++ {
        last = last.Next()
    }
    last = last.Prev()

    nr.ipset.AddRange(Range{First: first, Last: last})
    return nil
}

func (nr *NetworkRange) RemoveNetwork(cidr string) error {
    other := NewNetworkRange()
    if err := other.AddNetwork(cidr); err != nil {
        return err
    }

    nr.ipset = nr.ipset.Difference(other.ipset)
    return nil
}

func (nr *NetworkRange) CountAddresses() uint64 {
    var count uint64
    for _, r := range nr.ipset.ranges {
        // This is a simplified count - for real usage you'd need
        // to handle the case where the difference is too large
        first := r.First.As4()
        last := r.Last.As4()

        var firstInt, lastInt uint32
        for i := 0; i < 4; i++ {
            firstInt = firstInt<<8 | uint32(first[i])
            lastInt = lastInt<<8 | uint32(last[i])
        }

        count += uint64(lastInt - firstInt + 1)
    }
    return count
}

3. Firewall Rule Optimizer

Here's how you might optimize overlapping firewall rules:

type Rule struct {
    Network netip.Prefix
    Action  string // "allow" or "deny"
}

type FirewallOptimizer struct {
    allowSet *IPSet
    denySet  *IPSet
}

func NewFirewallOptimizer() *FirewallOptimizer {
    return &FirewallOptimizer{
        allowSet: NewIPSet(),
        denySet:  NewIPSet(),
    }
}

func (fo *FirewallOptimizer) AddRule(rule Rule) error {
    // Convert prefix to range
    first := rule.Network.Addr()
    last := first
    for i := 0; i < 1<<(32-rule.Network.Bits()); i++ {
        last = last.Next()
    }
    last = last.Prev()

    r := Range{First: first, Last: last}

    if rule.Action == "allow" {
        fo.allowSet.AddRange(r)
    } else {
        fo.denySet.AddRange(r)
    }

    return nil
}

func (fo *FirewallOptimizer) GetOptimizedRules() []Rule {
    var rules []Rule

    // First add deny rules (they take precedence)
    for _, r := range fo.denySet.ranges {
        // Convert range back to prefix(es)
        prefixes := rangeToPrefixes(r)
        for _, p := range prefixes {
            rules = append(rules, Rule{
                Network: p,
                Action:  "deny",
            })
        }
    }

    // Then add allow rules, but only for ranges not covered by deny rules
    allowedOnly := fo.allowSet.Difference(fo.denySet)
    for _, r := range allowedOnly.ranges {
        prefixes := rangeToPrefixes(r)
        for _, p := range prefixes {
            rules = append(rules, Rule{
                Network: p,
                Action:  "allow",
            })
        }
    }

    return rules
}

// rangeToPrefixes converts an IP range to a list of CIDR prefixes
func rangeToPrefixes(r Range) []netip.Prefix {
    // This is a simplified implementation
    // In practice, you'd want to implement an algorithm to
    // find the minimal set of prefixes that cover the range
    return []netip.Prefix{
        netip.PrefixFrom(r.First, 24), // Example - real impl would be more complex
    }
}

Performance Considerations

1. Range Optimization The optimize method is crucial for performance. Without it, the ranges list would grow unnecessarily large:

   // Bad - without optimization
   set.Add(ip1)
   set.Add(ip2)
   set.Add(ip3)
   // Results in three separate ranges

   // Good - with optimization
   set.AddRange(range1)
   // Automatically merges overlapping ranges

1. Binary Search Use binary search when looking up IPs in ranges:

   idx := sort.Search(len(ranges), func(i int) bool {
       return ranges[i].Last.Compare(ip) >= 0
   })

1. Memory Usage Store ranges instead of individual IPs:

   // Bad
   map[netip.Addr]bool // One entry per IP

   // Good
   []Range // One entry per continuous range

Best Practices

1. Always Validate Input

   func validateRange(first, last netip.Addr) error {
       if first.Compare(last) > 0 {
           return fmt.Errorf("invalid range: first IP (%s) > last IP (%s)", 
               first, last)
       }

       if first.Is4() != last.Is4() {
           return fmt.Errorf("mixed IP versions not supported")
       }

       return nil
   }

1. Handle IPv4 and IPv6 Separately

   type IPSet struct {
       v4ranges []Range
       v6ranges []Range
   }

1. Use Efficient Set Operations

   // Avoid iterating through individual IPs
   // Use range-based operations instead

What's Next?

In our next article, we'll explore the actual methods available in the net/netip package for working with IP addresses. While we've built our own set implementation here, understanding the built-in capabilities is crucial for effective network programming.

Until then, happy coding! Remember that working with IP sets can be tricky, so always test your implementations thoroughly, especially around edge cases and range boundaries.