Efficient way to check IP in slice of IP addresses in Golang

后端 未结 2 1428
旧时难觅i
旧时难觅i 2021-01-25 17:13

I\'m developing a network application in Golang. I have a slice of IP addresses. Each time a request comes I use net.LookupIP(host) to find out IP address of host w

2条回答
  •  清酒与你
    2021-01-25 17:21

    With a "set"

    Building our set

    There is no builtin Set type in Go, but you can elegantly use a map[Type]bool as a set, e.g.:

    // Create a set with 2 values in it: [1, 2]
    m := map[int]bool{1: true, 2: true}
    
    // Test an element:
    fmt.Println(m[1]) // true
    fmt.Println(m[3]) // false
    
    // Set an element:
    m[3] = true
    fmt.Println(m[3]) // true
    
    // Delete an element:
    delete(m, 1)
    fmt.Println(m[1]) // false
    

    Note: we exploited the fact that if a key is not in the map, indexing the map results in the zero value for the value type, which is false in case of bool, properly telling that the element is not in the map (set).

    Try it on the Go Playground.

    Note #2: there are a few tricks to make the code to handle a map as a set shorter, you can check them in this answer: Check if a value is in a list.

    Using net.IP in the set

    Now we only need a type representing a net.IP which can be used as the key type in a map (see this question about what constitutes a map key type: How can I prevent a type being used as a map key?).

    Unfortunately net.IP itself does not qualify, because it is a slice:

    type IP []byte
    

    And slices are not comparable. See this question for details: Hash with key as an array type and this: Why have arrays in Go?

    An easy way is to convert it to a canonical string value and we're done. For this we may simply convert the bytes of the IP to a hex string. But an IPv4 address may be presented as IPv6, so we should first convert it to IPv6:

    func Key(ip net.IP) string {
        return hex.EncodeToString(ip.To16())
    }
    

    Note: bytes of an IP address may not be a valid UTF-8 encoded string (which is how Go stores strings in memory), but string values in Go represent arbitrary byte sequences, so the following also works, is much simpler and is much more efficient:

    func Key(ip net.IP) string {
        return string(ip.To16())  // Simple []byte => string conversion
    }
    

    We can use such IP strings as the keys. Populate your map with IPs to check against:

    // Populate forbidden IPs:
    forbIPs := map[string]bool{
        Key(ip1): true,
        Key(ip2): true,
    }
    
    // Now check a single IP:
    ipToCheck := ...
    if forbIPs[Key(ipToCheck)] {
        fmt.Println("Forbidden!")
    } else {
        fmt.Println("Allowed.")
    }
    

    If you have multiple IPs to check (as returned by net.LookupIP()), it's a single for loop:

    ips, err := net.LookupIP(host)
    // Check err
    for _, ip := range ips {
        if forbIPs[Key(ip)] {
            // FORBIDDEN!
        }
    }
    

    Alternative Key type

    Note that –as mentioned above– slices are not comparable but arrays are. So we could also use an array as the key. This is how it could look like:

    func Key(ip net.IP) (a [16]byte) {
        copy(a[:], ip)
        return
    }
    
    // And the IP set:
    forbIPs := map[[16]byte]bool{
        // ...
    }
    

    Alternatives

    Sorted slice

    Alternatively we could just simply store the forbidden IPs in a slice []net.IP, and keep it sorted. If it is sorted, we can use binary search to find an IP in it (standard library sort.Search()).

    Yes, binary search has O(log2(n)) complexity compared to the O(1) complexity of the (hash)map solution above. But this alternative has another pro:

    Enumerating individual IPs is not always practical. Sometimes (often) it is easier to list IP ranges. The first solution is not feasible to handle IP ranges, but this solution may be: you can find ranges that cover an IP address also in O(log2(n)) time.

提交回复
热议问题