| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | package main |
| 4 | |
| 5 | import ( |
| 6 | "archive/tar" |
| 7 | "archive/zip" |
| 8 | "bytes" |
| 9 | "compress/gzip" |
| 10 | "strings" |
| 11 | "testing" |
| 12 | |
| 13 | "github.com/stretchr/testify/require" |
| 14 | ) |
| 15 | |
| 16 | func TestParseIPToASNCombinedTSVAsn(t *testing.T) { |
| 17 | payload := []byte( |
| 18 | "1.0.0.0\t1.0.0.255\t13335\tUS\tCloudflare\n" + |
| 19 | "2001:db8::\t2001:db8::ffff\t64512\tDE\tExample ASN\n", |
| 20 | ) |
| 21 | asnRanges, err := parseIPToASNCombinedTSVAsn(payload) |
| 22 | require.NoError(t, err) |
| 23 | require.Len(t, asnRanges, 2) |
| 24 | require.EqualValues(t, 13335, asnRanges[0].asn) |
| 25 | require.Equal(t, "Cloudflare", asnRanges[0].org) |
| 26 | require.EqualValues(t, 64512, asnRanges[1].asn) |
| 27 | } |
| 28 | |
| 29 | func TestParseIPToASNCombinedTSVGeo(t *testing.T) { |
| 30 | payload := []byte( |
| 31 | "1.0.0.0\t1.0.0.255\t13335\tUS\tCloudflare\n" + |
| 32 | "2001:db8::\t2001:db8::ffff\t64512\tDE\tExample ASN\n", |
| 33 | ) |
| 34 | geoRanges, err := parseIPToASNCombinedTSVGeo(payload) |
| 35 | require.NoError(t, err) |
| 36 | require.Len(t, geoRanges, 2) |
| 37 | require.Equal(t, "US", geoRanges[0].country) |
| 38 | require.Equal(t, "DE", geoRanges[1].country) |
| 39 | } |
| 40 | |
| 41 | func TestParseDBIPAsnCSV(t *testing.T) { |
| 42 | payload := []byte("\"1.0.0.0\",\"1.0.0.255\",\"13335\",\"Cloudflare\"\n") |
| 43 | ranges, err := parseDBIPAsnCSV(payload) |
| 44 | require.NoError(t, err) |
| 45 | require.Len(t, ranges, 1) |
| 46 | require.EqualValues(t, 13335, ranges[0].asn) |
| 47 | require.Equal(t, "Cloudflare", ranges[0].org) |
| 48 | } |
| 49 | |
| 50 | func TestParseDBIPCountryCSV(t *testing.T) { |
| 51 | payload := []byte("\"1.0.0.0\",\"1.0.0.255\",\"US\"\n") |
| 52 | ranges, err := parseDBIPCountryCSV(payload) |
| 53 | require.NoError(t, err) |
| 54 | require.Len(t, ranges, 1) |
| 55 | require.Equal(t, "US", ranges[0].country) |
| 56 | } |
| 57 | |
| 58 | func TestParseDBIPCityCSV(t *testing.T) { |
| 59 | payload := []byte("\"1.0.0.0\",\"1.0.0.255\",\"NA\",\"US\",\"California\",\"Mountain View\",\"37.3861\",\"-122.0839\"\n") |
| 60 | ranges, err := parseDBIPCityCSV(payload) |
| 61 | require.NoError(t, err) |
| 62 | require.Len(t, ranges, 1) |
| 63 | require.Equal(t, "US", ranges[0].country) |
| 64 | require.Equal(t, "California", ranges[0].state) |
| 65 | require.Equal(t, "Mountain View", ranges[0].city) |
| 66 | require.True(t, ranges[0].hasLocation) |
| 67 | require.InDelta(t, 37.3861, ranges[0].latitude, 0.0001) |
| 68 | require.InDelta(t, -122.0839, ranges[0].longitude, 0.0001) |
| 69 | } |
| 70 | |
| 71 | func TestParseCAIDAPrefix2AS(t *testing.T) { |
| 72 | payload := []byte( |
| 73 | "1.0.0.0\t24\t13335_38803\n" + |
| 74 | "2001:db8::\t48\t64512\n", |
| 75 | ) |
| 76 | ranges, err := parseCAIDAPrefix2AS(payload) |
| 77 | require.NoError(t, err) |
| 78 | require.Len(t, ranges, 2) |
| 79 | require.Equal(t, "1.0.0.0", ranges[0].start.String()) |
| 80 | require.Equal(t, "1.0.0.255", ranges[0].end.String()) |
| 81 | require.EqualValues(t, 13335, ranges[0].asn) |
| 82 | require.Equal(t, "2001:db8::", ranges[1].start.String()) |
| 83 | require.EqualValues(t, 64512, ranges[1].asn) |
| 84 | } |
| 85 | |
| 86 | func TestParseMaxMindCountryCSVZip(t *testing.T) { |
| 87 | payload := buildZip(t, map[string]string{ |
| 88 | "GeoLite2-Country-CSV_20260501/GeoLite2-Country-Blocks-IPv4.csv": strings.Join([]string{ |
| 89 | "network,geoname_id,registered_country_geoname_id,represented_country_geoname_id,is_anonymous_proxy,is_satellite_provider,is_anycast", |
| 90 | "1.0.0.0/24,6252001,,,,0,0", |
| 91 | }, "\n"), |
| 92 | "GeoLite2-Country-CSV_20260501/GeoLite2-Country-Blocks-IPv6.csv": strings.Join([]string{ |
| 93 | "network,geoname_id,registered_country_geoname_id,represented_country_geoname_id,is_anonymous_proxy,is_satellite_provider,is_anycast", |
| 94 | "2001:db8::/48,2921044,,,,0,0", |
| 95 | }, "\n"), |
| 96 | "GeoLite2-Country-CSV_20260501/GeoLite2-Country-Locations-en.csv": strings.Join([]string{ |
| 97 | "geoname_id,locale_code,continent_code,continent_name,country_iso_code,country_name,is_in_european_union", |
| 98 | "6252001,en,NA,North America,US,United States,0", |
| 99 | "2921044,en,EU,Europe,DE,Germany,1", |
| 100 | }, "\n"), |
| 101 | }) |
| 102 | |
| 103 | ranges, err := parseMaxMindCountryCSVZip(payload) |
| 104 | require.NoError(t, err) |
| 105 | require.Len(t, ranges, 2) |
| 106 | require.Equal(t, "US", ranges[0].country) |
| 107 | require.Equal(t, "1.0.0.0", ranges[0].start.String()) |
| 108 | require.Equal(t, "DE", ranges[1].country) |
| 109 | require.Equal(t, "2001:db8::", ranges[1].start.String()) |
| 110 | } |
| 111 | |
| 112 | func TestParseIP2LocationCountryZip(t *testing.T) { |
| 113 | payload := buildZip(t, map[string]string{ |
| 114 | "IP2LOCATION-LITE-DB1.CSV": "\"16777216\",\"16777471\",\"AU\",\"Australia\"\n", |
| 115 | }) |
| 116 | ranges, err := parseIP2LocationCountryZip(payload) |
| 117 | require.NoError(t, err) |
| 118 | require.Len(t, ranges, 1) |
| 119 | require.Equal(t, "1.0.0.0", ranges[0].start.String()) |
| 120 | require.Equal(t, "1.0.0.255", ranges[0].end.String()) |
| 121 | require.Equal(t, "AU", ranges[0].country) |
| 122 | } |
| 123 | |
| 124 | func TestParseIPDenyCountryTarGZ(t *testing.T) { |
| 125 | payload := buildTarGZ(t, map[string]string{ |
| 126 | "./us.zone": "1.0.0.0/24\n", |
| 127 | "./GR.ZONE": "1.0.1.0/24\n", |
| 128 | }) |
| 129 | ranges, err := parseIPDenyCountryTarGZ(payload) |
| 130 | require.NoError(t, err) |
| 131 | require.Len(t, ranges, 2) |
| 132 | byCountry := make(map[string]geoRange, len(ranges)) |
| 133 | for _, rec := range ranges { |
| 134 | byCountry[rec.country] = rec |
| 135 | } |
| 136 | require.Equal(t, "1.0.0.255", byCountry["US"].end.String()) |
| 137 | require.Equal(t, "1.0.1.255", byCountry["GR"].end.String()) |
| 138 | } |
| 139 | |
| 140 | func TestParseIPIPCountryZip(t *testing.T) { |
| 141 | payload := buildZip(t, map[string]string{ |
| 142 | "country.txt": "1.0.0.0/24\tANYCAST\n1.0.1.0/24\tCN\n", |
| 143 | }) |
| 144 | ranges, err := parseIPIPCountryZip(payload) |
| 145 | require.NoError(t, err) |
| 146 | require.Len(t, ranges, 1) |
| 147 | require.Equal(t, "CN", ranges[0].country) |
| 148 | require.Equal(t, "1.0.1.0", ranges[0].start.String()) |
| 149 | require.Equal(t, "1.0.1.255", ranges[0].end.String()) |
| 150 | } |
| 151 | |
| 152 | func TestEstimatedRangeCapacity(t *testing.T) { |
| 153 | require.Equal(t, 0, estimatedRangeCapacity(0, 64, 1<<20)) |
| 154 | require.Equal(t, 0, estimatedRangeCapacity(128, 0, 1<<20)) |
| 155 | require.Equal(t, 2, estimatedRangeCapacity(128, 64, 1<<20)) |
| 156 | require.Equal(t, 4, estimatedRangeCapacity(1024, 64, 4)) |
| 157 | } |
| 158 | |
| 159 | func TestParseIPDecimalIPv4(t *testing.T) { |
| 160 | addr, err := parseIP("16777216") // 1.0.0.0 |
| 161 | require.NoError(t, err) |
| 162 | require.Equal(t, "1.0.0.0", addr.String()) |
| 163 | } |
| 164 | |
| 165 | func TestParseRangeRejectsFamilyMix(t *testing.T) { |
| 166 | _, _, err := parseRangeEndpoints("1.0.0.0", "2001:db8::1") |
| 167 | require.Error(t, err) |
| 168 | } |
| 169 | |
| 170 | func buildZip(t *testing.T, files map[string]string) []byte { |
| 171 | t.Helper() |
| 172 | |
| 173 | var buf bytes.Buffer |
| 174 | zw := zip.NewWriter(&buf) |
| 175 | for name, content := range files { |
| 176 | w, err := zw.Create(name) |
| 177 | require.NoError(t, err) |
| 178 | _, err = w.Write([]byte(content)) |
| 179 | require.NoError(t, err) |
| 180 | } |
| 181 | require.NoError(t, zw.Close()) |
| 182 | return buf.Bytes() |
| 183 | } |
| 184 | |
| 185 | func buildTarGZ(t *testing.T, files map[string]string) []byte { |
| 186 | t.Helper() |
| 187 | |
| 188 | rawTar := buildTar(t, files) |
| 189 | var buf bytes.Buffer |
| 190 | gz := gzip.NewWriter(&buf) |
| 191 | _, err := gz.Write(rawTar) |
| 192 | require.NoError(t, err) |
| 193 | require.NoError(t, gz.Close()) |
| 194 | return buf.Bytes() |
| 195 | } |
| 196 | |
| 197 | func buildTar(t *testing.T, files map[string]string) []byte { |
| 198 | t.Helper() |
| 199 | |
| 200 | var buf bytes.Buffer |
| 201 | tw := tar.NewWriter(&buf) |
| 202 | for name, content := range files { |
| 203 | err := tw.WriteHeader(&tar.Header{ |
| 204 | Name: name, |
| 205 | Mode: 0o644, |
| 206 | Size: int64(len(content)), |
| 207 | }) |
| 208 | require.NoError(t, err) |
| 209 | _, err = tw.Write([]byte(content)) |
| 210 | require.NoError(t, err) |
| 211 | } |
| 212 | require.NoError(t, tw.Close()) |
| 213 | return buf.Bytes() |
| 214 | } |