| 1 | package geo |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "encoding/binary" |
| 7 | "errors" |
| 8 | "os" |
| 9 | "path/filepath" |
| 10 | "testing" |
| 11 | "time" |
| 12 | ) |
| 13 | |
| 14 | // pathologicalMMDB builds a structurally valid IPv4 database of 32 nodes |
| 15 | // whose two children are both the next node: 2^32 network paths through 32 |
| 16 | // nodes (code review R8-1). |
| 17 | func pathologicalMMDB(t *testing.T) string { |
| 18 | t.Helper() |
| 19 | const nodes = 32 |
| 20 | var buf bytes.Buffer |
| 21 | rec := func(v uint32) { buf.Write([]byte{byte(v >> 16), byte(v >> 8), byte(v)}) } // 24-bit records |
| 22 | for i := 0; i < nodes; i++ { |
| 23 | next := uint32(i + 1) |
| 24 | if i == nodes-1 { |
| 25 | next = nodes + 16 // data section offset 0 |
| 26 | } |
| 27 | rec(next) |
| 28 | rec(next) |
| 29 | } |
| 30 | buf.Write(make([]byte, 16)) // data section separator |
| 31 | buf.WriteByte(0xE0) // the one record: an empty map |
| 32 | str := func(s string) { buf.WriteByte(0x40 | byte(len(s))); buf.WriteString(s) } |
| 33 | u16 := func(v uint16) { buf.WriteByte(0xA2); binary.Write(&buf, binary.BigEndian, v) } |
| 34 | u32 := func(v uint32) { buf.WriteByte(0xC4); binary.Write(&buf, binary.BigEndian, v) } |
| 35 | buf.WriteString("\xAB\xCD\xEFMaxMind.com") |
| 36 | buf.WriteByte(0xE0 | 9) // metadata map, 9 keys |
| 37 | str("node_count") |
| 38 | u32(nodes) |
| 39 | str("record_size") |
| 40 | u16(24) |
| 41 | str("ip_version") |
| 42 | u16(4) |
| 43 | str("database_type") |
| 44 | str("Pathological") |
| 45 | str("languages") |
| 46 | buf.Write([]byte{0x00, 0x04}) // array of 0 |
| 47 | str("binary_format_major_version") |
| 48 | u16(2) |
| 49 | str("binary_format_minor_version") |
| 50 | u16(0) |
| 51 | str("build_epoch") |
| 52 | buf.Write([]byte{0x08, 0x02}) // uint64, 8 bytes |
| 53 | binary.Write(&buf, binary.BigEndian, uint64(time.Now().Unix())) |
| 54 | str("description") |
| 55 | buf.WriteByte(0xE0) // empty map |
| 56 | p := filepath.Join(t.TempDir(), "dag.mmdb") |
| 57 | if err := os.WriteFile(p, buf.Bytes(), 0o644); err != nil { |
| 58 | t.Fatal(err) |
| 59 | } |
| 60 | return p |
| 61 | } |
| 62 | |
| 63 | func TestOpenMMDBRefusesAPathologicalTreeInTime(t *testing.T) { |
| 64 | p := pathologicalMMDB(t) |
| 65 | start := time.Now() |
| 66 | _, err := OpenMMDBVerified(context.Background(), p) |
| 67 | if !errors.Is(err, ErrSearchTreeTooLarge) { |
| 68 | t.Fatalf("err = %v, want ErrSearchTreeTooLarge", err) |
| 69 | } |
| 70 | if d := time.Since(start); d > 5*time.Second { |
| 71 | t.Fatalf("took %s", d) |
| 72 | } |
| 73 | // The real fixtures still open. |
| 74 | for _, f := range []string{"GeoIP2-City-Test.mmdb", "GeoIP2-Country-Test.mmdb", "MaxMind-DB-test-ipv4-24.mmdb"} { |
| 75 | m, err := OpenMMDBVerified(context.Background(), "testdata/"+f) |
| 76 | if err != nil { |
| 77 | t.Fatalf("%s: %v", f, err) |
| 78 | } |
| 79 | m.Close() |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | // TestRealDatabaseOpensVerified runs only with GEOIP_REAL_MMDB set to a real |
| 84 | // database (DB-IP Lite, GeoLite2): the bounded walk must accept it. |
| 85 | func TestRealDatabaseOpensVerified(t *testing.T) { |
| 86 | p := os.Getenv("GEOIP_REAL_MMDB") |
| 87 | if p == "" { |
| 88 | t.Skip("set GEOIP_REAL_MMDB to a real .mmdb") |
| 89 | } |
| 90 | start := time.Now() |
| 91 | m, err := OpenMMDBVerified(context.Background(), p) |
| 92 | if err != nil { |
| 93 | t.Fatal(err) |
| 94 | } |
| 95 | m.Close() |
| 96 | t.Logf("%s verified in %s", m.Source(), time.Since(start)) |
| 97 | } |
| 98 | |
amadan