geoloader.go 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. package geodb
  2. import (
  3. "bytes"
  4. "encoding/csv"
  5. "io"
  6. "strings"
  7. "imuslab.com/zoraxy/mod/netutils"
  8. )
  9. func (s *Store) search(ip string) string {
  10. if strings.Contains(ip, ",") {
  11. //This is a CF proxied request. We only need the front part
  12. //Example 219.71.102.145, 172.71.139.178
  13. ip = strings.Split(ip, ",")[0]
  14. ip = strings.TrimSpace(ip)
  15. }
  16. //Search in geotrie tree
  17. cc := ""
  18. if netutils.IsIPv6(ip) {
  19. if s.geotrieIpv6 == nil {
  20. cc = s.slowSearchIpv6(ip)
  21. } else {
  22. cc = s.geotrieIpv6.search(ip)
  23. }
  24. } else {
  25. if s.geotrie == nil {
  26. cc = s.slowSearchIpv4(ip)
  27. } else {
  28. cc = s.geotrie.search(ip)
  29. }
  30. }
  31. /*
  32. if cc != "" {
  33. s.geoipCache.Store(ip, cc)
  34. }
  35. */
  36. return cc
  37. }
  38. // Construct the trie data structure for quick lookup
  39. func constrctTrieTree(data [][]string) *trie {
  40. tt := newTrie()
  41. for _, entry := range data {
  42. startIp := entry[0]
  43. endIp := entry[1]
  44. cc := entry[2]
  45. tt.insert(startIp, cc)
  46. tt.insert(endIp, cc)
  47. }
  48. return tt
  49. }
  50. // Parse the embedded csv as ipstart, ipend and country code entries
  51. func parseCSV(content []byte) ([][]string, error) {
  52. var records [][]string
  53. r := csv.NewReader(bytes.NewReader(content))
  54. for {
  55. record, err := r.Read()
  56. if err == io.EOF {
  57. break
  58. }
  59. if err != nil {
  60. return nil, err
  61. }
  62. records = append(records, record)
  63. }
  64. return records, nil
  65. }