Skip to content

Commit 7c52817

Browse files
IncSWIncSW
committed
add enterprise fields & refactoring reader by specific db types
1 parent 1784b32 commit 7c52817

23 files changed

+801
-375
lines changed

‎.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,6 @@ testdata/GeoIP2-City.mmdb
22
testdata/GeoIP2-Connection-Type.mmdb
33
testdata/GeoIP2-Country.mmdb
44
testdata/GeoIP2-ISP.mmdb
5+
testdata/GeoLite2-ASN.mmdb
6+
testdata/GeoLite2-City.mmdb
7+
testdata/GeoLite2-Country.mmdb

‎README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@ Inspired by [oschwald/geoip2-golang](https://github.com/oschwald/geoip2-golang).
1616
```go
1717
import "github.com/IncSW/geoip2"
1818

19-
reader, err := geoip2.NewReaderFromFile("path/to/GeoIP2-City.mmdb")
19+
reader, err := geoip2.NewCityReaderFromFile("path/to/GeoIP2-City.mmdb")
2020
if err != nil {
2121
panic(err)
2222
}
23-
record, err := reader.LookupCity(net.ParseIP("81.2.69.142"))
23+
record, err := reader.Lookup(net.ParseIP("81.2.69.142"))
2424
if err != nil {
2525
panic(err)
2626
}

‎anonymous_ip.go

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,8 @@
11
package geoip2
22

3-
import (
4-
"errors"
5-
"strconv"
6-
)
3+
import "errors"
74

8-
func readAnonymousIPResult(buffer []byte, offset uint) (*AnonymousIPResult, error) {
9-
dataType, size, offset, err := readControl(buffer, offset)
10-
if err != nil {
11-
return nil, err
12-
}
13-
result := &AnonymousIPResult{}
14-
switch dataType {
15-
case dataTypeMap:
16-
_, err = readAnonymousIPMap(result, buffer, size, offset)
17-
if err != nil {
18-
return nil, err
19-
}
20-
case dataTypePointer:
21-
pointer, _, err := readPointer(buffer, size, offset)
22-
if err != nil {
23-
return nil, err
24-
}
25-
dataType, size, offset, err := readControl(buffer, pointer)
26-
if err != nil {
27-
return nil, err
28-
}
29-
if dataType != dataTypeMap {
30-
return nil, errors.New("invalid anonymous ip pointer type: " + strconv.Itoa(int(dataType)))
31-
}
32-
_, err = readAnonymousIPMap(result, buffer, size, offset)
33-
if err != nil {
34-
return nil, err
35-
}
36-
default:
37-
return nil, errors.New("invalid anonymous ip type: " + strconv.Itoa(int(dataType)))
38-
}
39-
return result, nil
40-
}
41-
42-
func readAnonymousIPMap(result *AnonymousIPResult, buffer []byte, mapSize uint, offset uint) (uint, error) {
5+
func readAnonymousIPMap(result *AnonymousIP, buffer []byte, mapSize uint, offset uint) (uint, error) {
436
var key []byte
447
var err error
458
for i := uint(0); i < mapSize; i++ {

‎asn.go

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,8 @@
11
package geoip2
22

3-
import (
4-
"errors"
5-
"strconv"
6-
)
3+
import "errors"
74

8-
func readASNResult(buffer []byte, offset uint) (*ASNResult, error) {
9-
dataType, size, offset, err := readControl(buffer, offset)
10-
if err != nil {
11-
return nil, err
12-
}
13-
result := &ASNResult{}
14-
switch dataType {
15-
case dataTypeMap:
16-
_, err = readASNMap(result, buffer, size, offset)
17-
if err != nil {
18-
return nil, err
19-
}
20-
case dataTypePointer:
21-
pointer, _, err := readPointer(buffer, size, offset)
22-
if err != nil {
23-
return nil, err
24-
}
25-
dataType, size, offset, err := readControl(buffer, pointer)
26-
if err != nil {
27-
return nil, err
28-
}
29-
if dataType != dataTypeMap {
30-
return nil, errors.New("invalid asn pointer type: " + strconv.Itoa(int(dataType)))
31-
}
32-
_, err = readASNMap(result, buffer, size, offset)
33-
if err != nil {
34-
return nil, err
35-
}
36-
default:
37-
return nil, errors.New("invalid asn type: " + strconv.Itoa(int(dataType)))
38-
}
39-
return result, nil
40-
}
41-
42-
func readASNMap(result *ASNResult, buffer []byte, mapSize uint, offset uint) (uint, error) {
5+
func readASNMap(result *ASN, buffer []byte, mapSize uint, offset uint) (uint, error) {
436
var key []byte
447
var err error
458
for i := uint(0); i < mapSize; i++ {

‎city.go

Lines changed: 5 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -5,74 +5,6 @@ import (
55
"strconv"
66
)
77

8-
func readCityResult(buffer []byte, offset uint) (*CityResult, error) {
9-
dataType, size, offset, err := readControl(buffer, offset)
10-
if err != nil {
11-
return nil, err
12-
}
13-
if dataType != dataTypeMap {
14-
return nil, errors.New("invalid city type: " + strconv.Itoa(int(dataType)))
15-
}
16-
var key []byte
17-
result := &CityResult{}
18-
for i := uint(0); i < size; i++ {
19-
key, offset, err = readMapKey(buffer, offset)
20-
if err != nil {
21-
return nil, err
22-
}
23-
switch b2s(key) {
24-
case "city":
25-
offset, err = readCity(&result.City, buffer, offset)
26-
if err != nil {
27-
return nil, err
28-
}
29-
case "continent":
30-
offset, err = readContinent(&result.Continent, buffer, offset)
31-
if err != nil {
32-
return nil, err
33-
}
34-
case "country":
35-
offset, err = readCountry(&result.Country, buffer, offset)
36-
if err != nil {
37-
return nil, err
38-
}
39-
case "location":
40-
offset, err = readLocation(&result.Location, buffer, offset)
41-
if err != nil {
42-
return nil, err
43-
}
44-
case "postal":
45-
offset, err = readPostal(&result.Postal, buffer, offset)
46-
if err != nil {
47-
return nil, err
48-
}
49-
case "registered_country":
50-
offset, err = readCountry(&result.RegisteredCountry, buffer, offset)
51-
if err != nil {
52-
return nil, err
53-
}
54-
case "represented_country":
55-
offset, err = readCountry(&result.RepresentedCountry, buffer, offset)
56-
if err != nil {
57-
return nil, err
58-
}
59-
case "subdivisions":
60-
result.Subdivisions, offset, err = readSubdivisions(buffer, offset)
61-
if err != nil {
62-
return nil, err
63-
}
64-
case "traits":
65-
offset, err = readTraits(&result.Traits, buffer, offset)
66-
if err != nil {
67-
return nil, err
68-
}
69-
default:
70-
return nil, errors.New("unknown city response key: " + string(key) + ", type: " + strconv.Itoa(int(dataType)))
71-
}
72-
}
73-
return result, nil
74-
}
75-
768
func readCity(city *City, buffer []byte, offset uint) (uint, error) {
779
dataType, size, offset, err := readControl(buffer, offset)
7810
if err != nil {
@@ -122,6 +54,11 @@ func readCityMap(city *City, buffer []byte, mapSize uint, offset uint) (uint, er
12254
if err != nil {
12355
return 0, err
12456
}
57+
case "confidence":
58+
city.Confidence, offset, err = readUInt8(buffer, offset)
59+
if err != nil {
60+
return 0, err
61+
}
12562
default:
12663
return 0, errors.New("unknown city key: " + string(key))
12764
}

‎common.go

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,33 @@ func readFloat64(buffer []byte, offset uint) (float64, uint, error) {
8787
}
8888
}
8989

90+
func readUInt8(buffer []byte, offset uint) (uint8, uint, error) { // TODO: fix dataType
91+
dataType, size, offset, err := readControl(buffer, offset)
92+
if err != nil {
93+
return 0, 0, err
94+
}
95+
switch dataType {
96+
case dataTypeUint16, dataTypeUint32:
97+
newOffset := offset + size
98+
return uint8(bytesToUInt64(buffer[offset:newOffset])), newOffset, nil
99+
case dataTypePointer:
100+
pointer, newOffset, err := readPointer(buffer, size, offset)
101+
if err != nil {
102+
return 0, 0, err
103+
}
104+
dataType, size, offset, err := readControl(buffer, pointer)
105+
if err != nil {
106+
return 0, 0, err
107+
}
108+
if dataType != dataTypeUint16 && dataType != dataTypeUint32 {
109+
return 0, 0, errors.New("invalid uint8 pointer type: " + strconv.Itoa(int(dataType)))
110+
}
111+
return uint8(bytesToUInt64(buffer[offset : offset+size])), newOffset, nil
112+
default:
113+
return 0, 0, errors.New("invalid uint8 type: " + strconv.Itoa(int(dataType)))
114+
}
115+
}
116+
90117
func readUInt16(buffer []byte, offset uint) (uint16, uint, error) {
91118
dataType, size, offset, err := readControl(buffer, offset)
92119
if err != nil {

‎connection_type.go

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,8 @@
11
package geoip2
22

3-
import (
4-
"errors"
5-
"strconv"
6-
)
3+
import "errors"
74

8-
func readConnectionType(buffer []byte, offset uint) (string, error) {
9-
dataType, size, offset, err := readControl(buffer, offset)
10-
if err != nil {
11-
return "", err
12-
}
13-
result := &ConnectionTypeResult{}
14-
switch dataType {
15-
case dataTypeMap:
16-
_, err = readConnectionTypeMap(result, buffer, size, offset)
17-
if err != nil {
18-
return "", err
19-
}
20-
case dataTypePointer:
21-
pointer, _, err := readPointer(buffer, size, offset)
22-
if err != nil {
23-
return "", err
24-
}
25-
dataType, size, offset, err := readControl(buffer, pointer)
26-
if err != nil {
27-
return "", err
28-
}
29-
if dataType != dataTypeMap {
30-
return "", errors.New("invalid connectionType pointer type: " + strconv.Itoa(int(dataType)))
31-
}
32-
_, err = readConnectionTypeMap(result, buffer, size, offset)
33-
if err != nil {
34-
return "", err
35-
}
36-
default:
37-
return "", errors.New("invalid connectionType type: " + strconv.Itoa(int(dataType)))
38-
}
39-
return result.ConnectionType, nil
40-
}
41-
42-
func readConnectionTypeMap(result *ConnectionTypeResult, buffer []byte, mapSize uint, offset uint) (uint, error) {
5+
func readConnectionTypeMap(result *ConnectionType, buffer []byte, mapSize uint, offset uint) (uint, error) {
436
var key []byte
447
var err error
458
for i := uint(0); i < mapSize; i++ {

‎country.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,11 @@ func readCountryMap(country *Country, buffer []byte, mapSize uint, offset uint)
6969
if err != nil {
7070
return 0, err
7171
}
72+
case "confidence":
73+
country.Confidence, offset, err = readUInt8(buffer, offset)
74+
if err != nil {
75+
return 0, err
76+
}
7277
default:
7378
return 0, errors.New("unknown country key: " + string(key))
7479
}

‎domain.go

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,8 @@
11
package geoip2
22

3-
import (
4-
"errors"
5-
"strconv"
6-
)
3+
import "errors"
74

8-
func readDomain(buffer []byte, offset uint) (string, error) {
9-
dataType, size, offset, err := readControl(buffer, offset)
10-
if err != nil {
11-
return "", err
12-
}
13-
result := &DomainResult{}
14-
switch dataType {
15-
case dataTypeMap:
16-
_, err = readDomainMap(result, buffer, size, offset)
17-
if err != nil {
18-
return "", err
19-
}
20-
case dataTypePointer:
21-
pointer, _, err := readPointer(buffer, size, offset)
22-
if err != nil {
23-
return "", err
24-
}
25-
dataType, size, offset, err := readControl(buffer, pointer)
26-
if err != nil {
27-
return "", err
28-
}
29-
if dataType != dataTypeMap {
30-
return "", errors.New("invalid domain pointer type: " + strconv.Itoa(int(dataType)))
31-
}
32-
_, err = readDomainMap(result, buffer, size, offset)
33-
if err != nil {
34-
return "", err
35-
}
36-
default:
37-
return "", errors.New("invalid domain type: " + strconv.Itoa(int(dataType)))
38-
}
39-
return result.Domain, nil
40-
}
41-
42-
func readDomainMap(result *DomainResult, buffer []byte, mapSize uint, offset uint) (uint, error) {
5+
func readDomainMap(result *Domain, buffer []byte, mapSize uint, offset uint) (uint, error) {
436
var key []byte
447
var err error
458
for i := uint(0); i < mapSize; i++ {

0 commit comments

Comments
 (0)