Golang trie tree IP 轉國碼(Country Code)
Toby
Toby

最近因為有人 report 說 Zoraxy 的 GeoIP 功能不準,於是我在重新 review 這個 geoip database module 的時候發現:欸幹真的寫錯了欸!?

然後在花了一整個晚上 debug 之後,我終於找到正確的演算法和做法了。這裡為了省下後人在掘怎樣用 Golang 寫 IP 轉 Country Code 的 resolver,我這裡給大家一個快速上手的教學。

甚麼是 IP 轉 CC?

IP 就是 IP 地址,CC 是 country code,有時稱為 ISO 國碼,例如說 HK TW US GB 之類的。IPv4 跟 v6 的地址區間通常會被 assign 到某一個國家的 ISP 上,而造成說只要知道 IP 地址,某程度上你可以知道那個 request 是從哪個國家出來的(先不要說 VPN 或是 ip spoofing 之類的部分的話)。可是由於每隔一段時間就會有 ISP 倒閉、收購或是重組之類的,所以 IP 轉 CC 的可靠度並非 100% 準確,這個時候我們就需要從一個可靠的來源來定期更新這一個 IP range 到 CC 的 mapping database 了。

這個 database 最出名的應該就是 maxmind 的 geoip database,人家都整理好所有數據到一個 database 檔案裡面,也提供 api 讓你可以快速查詢 cc,很多現代只會拉 library 的開發者通常就會直接選用它。然而,基於它的 license 問題,如果真接用它的方案的話會讓授權和 licensing 變得一團亂,所以基於多種考慮之下這個方案我們最終沒有使用,而是使用 Public domain 或 CC-0 的資料來源。

GeoIP Data Source

最後我們使用了來自 GeoFeed + Whois + ASN (CC-0)的資料來源:

https://github.com/sapics/ip-location-db/tree/main/geolite2-country

問題是它的資料表是一個 csv 檔案,大概長這樣:

 1.0.0.0,1.0.0.255,AU
 1.0.1.0,1.0.3.255,CN
 1.0.4.0,1.0.7.255,AU
 1.0.8.0,1.0.15.255,CN
 1.0.16.0,1.0.31.255,JP
 1.0.32.0,1.0.63.255,CN
 1.0.64.0,1.0.127.255,JP
 1.0.128.0,1.0.255.255,TH
 1.1.0.0,1.1.0.255,CN
 1.1.1.0,1.1.1.255,AU
 1.1.2.0,1.1.63.255,CN
 1.1.64.0,1.1.127.255,JP
....
223.255.248.0,223.255.251.255,HK
223.255.252.0,223.255.253.255,CN
223.255.254.0,223.255.254.255,SG
223.255.255.0,223.255.255.255,AU

那假設我給你一個 ip 地址:A.B.C.D,你要怎樣把 country code 從 csv 裡面抓出來呢?

最直觀的答案:O(n)

當然如果會一點寫程式的人就會想到:那我把整個 csv loop 一次,看看哪一行的 start ip 跟 end ip 是包含我這個 ip 地址的就好了啦?這個的確是我們做給嵌入式裝置的做法(不過是為了省 memory),原理大概這樣:

func isIPv4InRange(startIP, endIP, testIP string) (bool, error) {
	start := net.ParseIP(startIP)
	end := net.ParseIP(endIP)
	test := net.ParseIP(testIP)

	if start == nil || end == nil || test == nil {
		return false, errors.New("invalid IP address format")
	}

	startUint := ipv4ToUInt32(start)
	endUint := ipv4ToUInt32(end)
	testUint := ipv4ToUInt32(test)

	return testUint >= startUint && testUint <= endUint, nil
}

然後就真的把每一行載出來檢查

// Slow country code lookup for
func (s *Store) slowSearchIpv4(ipAddr string) string {
	if isReservedIP(ipAddr) {
		return ""
	}
	for _, ipRange := range s.geodb {
		startIp := ipRange[0]
		endIp := ipRange[1]
		cc := ipRange[2]

		inRange, _ := isIPv4InRange(startIp, endIp, ipAddr)
		if inRange {
			return cc
		}
	}
	return ""
}

啊可是這真的慢,所以後來我們對於較正規的伺服器,使用了另一個資料結構來做到 O(31) (或是 IPv6 的 O(128)) ~ O(1) 的方法來做,就是:

Trie Tree

我都叫它 “Try” tree,但是原作者叫它做 “tree” tree,不過這不是重點,重點是這東西真的很有趣(?)

在我的 geoIP trie tree 定義裡面,我首先建立了兩個 struct:node 跟 head (裡面含一個 node,主要是用來處理 OOP 問題而不是 trie tree 本身需要的東西)

type trie_Node struct {
	childrens [2]*trie_Node
	cc        string
}

// Initializing the root of the trie
type trie struct {
	root *trie_Node
}

另外就是一些方便使用的 utilities function

// 把 IP 地址從 string 轉成 byte,ipv4 出來會像 [192,168,1,100] 這樣 4 個 bytes
func ipToBytes(ip string) []byte {
	// Parse the IP address string into a net.IP object
	parsedIP := net.ParseIP(ip)

	// Convert the IP address to a 4-byte slice
	ipBytes := parsedIP.To4()
	if ipBytes == nil {
		//This is an IPv6 address
		ipBytes = parsedIP.To16()
	}

	return ipBytes
}
//檢查是不是特殊 ip 地址,如 loopback 或 LAN
func isReservedIP(ip string) bool {
	parsedIP := net.ParseIP(ip)
	if parsedIP == nil {
		return false
	}
	// Check if the IP address is a loopback address
	if parsedIP.IsLoopback() {
		return true
	}
	// Check if the IP address is in the link-local address range
	if parsedIP.IsLinkLocalUnicast() || parsedIP.IsLinkLocalMulticast() {
		return true
	}
	//Check if the IP is in the reserved private range
	if parsedIP.IsPrivate() {
		return true
	}
	return false
}

之後就是 tree 本身的功能了。首先是建立 tree 和 insert tree 的 function。 順帶一提,因為 IPv4 是 32 bit 而 IPv6 是 128 bit,所以建出來分別就是兩顆 32 + 1層 跟 128 + 1 層的樹(這裡 + 1 是 root node,就像有些地方地面那層是 G/F 也有是 1/F 的,看你怎樣算層數)

// inititlaizing a new trie
func newTrie() *trie {
	t := new(trie)
	t.root = new(trie_Node)
	return t
}

// Passing words to trie
func (t *trie) insert(ipAddr string, cc string) {
	ipBytes := ipToBytes(ipAddr)
	current := t.root
	for _, b := range ipBytes {
		//For each byte in the ip address (4 / 16 bytes)
		//each byte is 8 bit
		for j := 7; j >= 0; j-- {
			bit := int(b >> j & 1)
			if current.childrens[bit] == nil {
				current.childrens[bit] = &trie_Node{
					childrens: [2]*trie_Node{},
					cc:        cc,
				}
			}
			current = current.childrens[bit]
		}
	}
}

然後就是 search 的 function


// Initializing the search for word in node
func (t *trie) search(ipAddr string) string {
	if isReservedIP(ipAddr) {
		return ""
	}

	ipBytes := ipToBytes(ipAddr)
	current := t.root
	for _, b := range ipBytes {
		//For each byte in the ip address
		//each byte is 8 bit
		for j := 7; j >= 0; j-- {
			bit := int(b >> j & 1)
			if current.childrens[bit] == nil {
				return current.cc
			}
			current = current.childrens[bit]
		}
	}

	if len(current.childrens) == 0 {
		return current.cc
	}

	//Not found
	return ""
}

這兩個 function 最重要的部分大概就是個 for loop

for j := 7; j >= 0; j-- {
	bit := int(b >> j & 1)
	if current.childrens[bit] == nil {
		return current.cc
	}
	current = current.childrens[bit]
}

假設我們的 IP 只有 4 個 bit, 1.0.0.0 到 1.0.1.1 是屬於 A 國家的;而 1.1.0.0 到 1.1.1.1 是屬於 B 國家的。要建這個 tree 的時候我們要 insert 4 次:

  • 1.0.0.0 -> A
  • 1.0.1.1 -> A
  • 1.1.0.0 -> B
  • 1.1.1.1 -> B

那我們的 trie tree 建出來會變成這樣,而每個 node 裡面都寫有相關的 country code:

當我們有新的 Ip 進來的時候,假設這是 1.0.1.0 (A)的時候,lookup 過程就會長這樣:

你會發現到第三個 bit 的時候,下一個 node 找不到,這樣我們就可以回傳最後有找到的 node 的 country code 回去(這裡是 A),這樣 country code 就能用接近 O(1) 的 time complexity 被找到了。當然,缺點就是因為整個 tree 都在 memory 裡面,所以這最少需要接近 700mb 的 RAM 才能用,結果如果你是要在嵌入式硬體上面跑這個的話…額…你還是回去用 for loop 好了(

備注

這個是 Zoraxy geodb 的建樹 entry point,如果你是沒看文字只來把 code 幹走的話可以順便把這個也抄一抄:

// Construct the trie data structure for quick lookup
func constrctTrieTree(data [][]string) *trie {
	tt := newTrie()
	for _, entry := range data {
		startIp := entry[0]
		endIp := entry[1]
		cc := entry[2]
		tt.insert(startIp, cc)
		tt.insert(endIp, cc)
	}

	return tt
}