Skip to content

Commit a60828e

Browse files
oschwaldoschwald
committed
Replace LookupNetwork with Network on Result
1 parent 16b8926 commit a60828e

File tree

3 files changed

+132
-141
lines changed

3 files changed

+132
-141
lines changed

‎reader.go

Lines changed: 22 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -131,46 +131,29 @@ func (r *Reader) Lookup(ip netip.Addr) Result {
131131
if r.buffer == nil {
132132
return Result{err: errors.New("cannot call Lookup on a closed database")}
133133
}
134-
pointer, _, _, err := r.lookupPointer(ip)
134+
pointer, prefixLen, err := r.lookupPointer(ip)
135135
if err != nil {
136-
return Result{err: err}
136+
return Result{
137+
ip: ip,
138+
prefixLen: uint8(prefixLen),
139+
err: err,
140+
}
137141
}
138142
if pointer == 0 {
139-
return Result{offset: notFound}
143+
return Result{
144+
ip: ip,
145+
prefixLen: uint8(prefixLen),
146+
offset: notFound,
147+
}
140148
}
141149
offset, err := r.resolveDataPointer(pointer)
142150
return Result{
143-
decoder: r.decoder,
144-
offset: uint(offset),
145-
err: err,
146-
}
147-
}
148-
149-
// LookupNetwork retrieves the database record for ip and stores it in the
150-
// value pointed to by result. The prefix returned is the network associated
151-
// with the data record in the database. The ok return value indicates whether
152-
// the database contained a record for the ip.
153-
//
154-
// If result is nil or not a pointer, an error is returned. If the data in the
155-
// database record cannot be stored in result because of type differences, an
156-
// UnmarshalTypeError is returned. If the database is invalid or otherwise
157-
// cannot be read, an InvalidDatabaseError is returned.
158-
func (r *Reader) LookupNetwork(
159-
ip netip.Addr,
160-
result any,
161-
) (prefix netip.Prefix, ok bool, err error) {
162-
if r.buffer == nil {
163-
return netip.Prefix{}, false, errors.New("cannot call Lookup on a closed database")
164-
}
165-
pointer, prefixLength, ip, err := r.lookupPointer(ip)
166-
// We return this error below as we want to return the prefix it is for
167-
168-
prefix, errP := r.cidr(ip, prefixLength)
169-
if pointer == 0 || err != nil || errP != nil {
170-
return prefix, false, errors.Join(err, errP)
151+
decoder: r.decoder,
152+
ip: ip,
153+
offset: uint(offset),
154+
prefixLen: uint8(prefixLen),
155+
err: err,
171156
}
172-
173-
return prefix, true, r.retrieveData(pointer, result)
174157
}
175158

176159
// LookupOffset maps an argument net.IP to a corresponding record offset in the
@@ -182,7 +165,7 @@ func (r *Reader) LookupOffset(ip netip.Addr) (uintptr, error) {
182165
if r.buffer == nil {
183166
return 0, errors.New("cannot call LookupOffset on a closed database")
184167
}
185-
pointer, _, _, err := r.lookupPointer(ip)
168+
pointer, _, err := r.lookupPointer(ip)
186169
if pointer == 0 || err != nil {
187170
return NotFound, err
188171
}
@@ -191,28 +174,6 @@ func (r *Reader) LookupOffset(ip netip.Addr) (uintptr, error) {
191174

192175
var zeroIP = netip.MustParseAddr("::")
193176

194-
func (r *Reader) cidr(ip netip.Addr, prefixLength int) (netip.Prefix, error) {
195-
if ip.Is4() {
196-
// This is necessary as the node that the IPv4 start is at may
197-
// be at a bit depth that is less that 96, i.e., ipv4Start points
198-
// to a leaf node. For instance, if a record was inserted at ::/8,
199-
// the ipv4Start would point directly at the leaf node for the
200-
// record and would have a bit depth of 8. This would not happen
201-
// with databases currently distributed by MaxMind as all of them
202-
// have an IPv4 subtree that is greater than a single node.
203-
if r.Metadata.IPVersion == 6 && r.ipv4StartBitDepth != 96 {
204-
return netip.PrefixFrom(zeroIP, r.ipv4StartBitDepth), nil
205-
}
206-
prefixLength -= 96
207-
}
208-
209-
prefix, err := ip.Prefix(prefixLength)
210-
if err != nil {
211-
return netip.Prefix{}, fmt.Errorf("creating prefix from %s/%d: %w", ip, prefixLength, err)
212-
}
213-
return prefix, nil
214-
}
215-
216177
// Decode the record at |offset| into |result|. The result value pointed to
217178
// must be a data value that corresponds to a record in the database. This may
218179
// include a struct representation of the data, a map capable of holding the
@@ -237,9 +198,9 @@ func (r *Reader) Decode(offset uintptr, result any) error {
237198
return Result{decoder: r.decoder, offset: uint(offset)}.Decode(result)
238199
}
239200

240-
func (r *Reader) lookupPointer(ip netip.Addr) (uint, int, netip.Addr, error) {
201+
func (r *Reader) lookupPointer(ip netip.Addr) (uint, int, error) {
241202
if r.Metadata.IPVersion == 4 && ip.Is6() {
242-
return 0, 0, ip, fmt.Errorf(
203+
return 0, 0, fmt.Errorf(
243204
"error looking up '%s': you attempted to look up an IPv6 address in an IPv4-only database",
244205
ip.String(),
245206
)
@@ -250,12 +211,12 @@ func (r *Reader) lookupPointer(ip netip.Addr) (uint, int, netip.Addr, error) {
250211
nodeCount := r.Metadata.NodeCount
251212
if node == nodeCount {
252213
// Record is empty
253-
return 0, prefixLength, ip, nil
214+
return 0, prefixLength, nil
254215
} else if node > nodeCount {
255-
return node, prefixLength, ip, nil
216+
return node, prefixLength, nil
256217
}
257218

258-
return 0, prefixLength, ip, newInvalidDatabaseError("invalid node in search tree")
219+
return 0, prefixLength, newInvalidDatabaseError("invalid node in search tree")
259220
}
260221

261222
func (r *Reader) traverseTree(ip netip.Addr, node uint, stopBit int) (uint, int) {

‎reader_test.go

Lines changed: 80 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -100,95 +100,95 @@ func TestLookupNetwork(t *testing.T) {
100100
}
101101

102102
tests := []struct {
103-
IP netip.Addr
104-
DBFile string
105-
ExpectedCIDR string
106-
ExpectedRecord any
107-
ExpectedOK bool
103+
IP netip.Addr
104+
DBFile string
105+
ExpectedNetwork string
106+
ExpectedRecord any
107+
ExpectedFound bool
108108
}{
109109
{
110-
IP: netip.MustParseAddr("1.1.1.1"),
111-
DBFile: "MaxMind-DB-test-ipv6-32.mmdb",
112-
ExpectedCIDR: "1.0.0.0/8",
113-
ExpectedRecord: nil,
114-
ExpectedOK: false,
110+
IP: netip.MustParseAddr("1.1.1.1"),
111+
DBFile: "MaxMind-DB-test-ipv6-32.mmdb",
112+
ExpectedNetwork: "1.0.0.0/8",
113+
ExpectedRecord: nil,
114+
ExpectedFound: false,
115115
},
116116
{
117-
IP: netip.MustParseAddr("::1:ffff:ffff"),
118-
DBFile: "MaxMind-DB-test-ipv6-24.mmdb",
119-
ExpectedCIDR: "::1:ffff:ffff/128",
120-
ExpectedRecord: map[string]any{"ip": "::1:ffff:ffff"},
121-
ExpectedOK: true,
117+
IP: netip.MustParseAddr("::1:ffff:ffff"),
118+
DBFile: "MaxMind-DB-test-ipv6-24.mmdb",
119+
ExpectedNetwork: "::1:ffff:ffff/128",
120+
ExpectedRecord: map[string]any{"ip": "::1:ffff:ffff"},
121+
ExpectedFound: true,
122122
},
123123
{
124-
IP: netip.MustParseAddr("::2:0:1"),
125-
DBFile: "MaxMind-DB-test-ipv6-24.mmdb",
126-
ExpectedCIDR: "::2:0:0/122",
127-
ExpectedRecord: map[string]any{"ip": "::2:0:0"},
128-
ExpectedOK: true,
124+
IP: netip.MustParseAddr("::2:0:1"),
125+
DBFile: "MaxMind-DB-test-ipv6-24.mmdb",
126+
ExpectedNetwork: "::2:0:0/122",
127+
ExpectedRecord: map[string]any{"ip": "::2:0:0"},
128+
ExpectedFound: true,
129129
},
130130
{
131-
IP: netip.MustParseAddr("1.1.1.1"),
132-
DBFile: "MaxMind-DB-test-ipv4-24.mmdb",
133-
ExpectedCIDR: "1.1.1.1/32",
134-
ExpectedRecord: map[string]any{"ip": "1.1.1.1"},
135-
ExpectedOK: true,
131+
IP: netip.MustParseAddr("1.1.1.1"),
132+
DBFile: "MaxMind-DB-test-ipv4-24.mmdb",
133+
ExpectedNetwork: "1.1.1.1/32",
134+
ExpectedRecord: map[string]any{"ip": "1.1.1.1"},
135+
ExpectedFound: true,
136136
},
137137
{
138-
IP: netip.MustParseAddr("1.1.1.3"),
139-
DBFile: "MaxMind-DB-test-ipv4-24.mmdb",
140-
ExpectedCIDR: "1.1.1.2/31",
141-
ExpectedRecord: map[string]any{"ip": "1.1.1.2"},
142-
ExpectedOK: true,
138+
IP: netip.MustParseAddr("1.1.1.3"),
139+
DBFile: "MaxMind-DB-test-ipv4-24.mmdb",
140+
ExpectedNetwork: "1.1.1.2/31",
141+
ExpectedRecord: map[string]any{"ip": "1.1.1.2"},
142+
ExpectedFound: true,
143143
},
144144
{
145-
IP: netip.MustParseAddr("1.1.1.3"),
146-
DBFile: "MaxMind-DB-test-decoder.mmdb",
147-
ExpectedCIDR: "1.1.1.0/24",
148-
ExpectedRecord: decoderRecord,
149-
ExpectedOK: true,
145+
IP: netip.MustParseAddr("1.1.1.3"),
146+
DBFile: "MaxMind-DB-test-decoder.mmdb",
147+
ExpectedNetwork: "1.1.1.0/24",
148+
ExpectedRecord: decoderRecord,
149+
ExpectedFound: true,
150150
},
151151
{
152-
IP: netip.MustParseAddr("::ffff:1.1.1.128"),
153-
DBFile: "MaxMind-DB-test-decoder.mmdb",
154-
ExpectedCIDR: "::ffff:1.1.1.0/120",
155-
ExpectedRecord: decoderRecord,
156-
ExpectedOK: true,
152+
IP: netip.MustParseAddr("::ffff:1.1.1.128"),
153+
DBFile: "MaxMind-DB-test-decoder.mmdb",
154+
ExpectedNetwork: "::ffff:1.1.1.0/120",
155+
ExpectedRecord: decoderRecord,
156+
ExpectedFound: true,
157157
},
158158
{
159-
IP: netip.MustParseAddr("::1.1.1.128"),
160-
DBFile: "MaxMind-DB-test-decoder.mmdb",
161-
ExpectedCIDR: "::101:100/120",
162-
ExpectedRecord: decoderRecord,
163-
ExpectedOK: true,
159+
IP: netip.MustParseAddr("::1.1.1.128"),
160+
DBFile: "MaxMind-DB-test-decoder.mmdb",
161+
ExpectedNetwork: "::101:100/120",
162+
ExpectedRecord: decoderRecord,
163+
ExpectedFound: true,
164164
},
165165
{
166-
IP: netip.MustParseAddr("200.0.2.1"),
167-
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
168-
ExpectedCIDR: "::/64",
169-
ExpectedRecord: "::0/64",
170-
ExpectedOK: true,
166+
IP: netip.MustParseAddr("200.0.2.1"),
167+
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
168+
ExpectedNetwork: "::/64",
169+
ExpectedRecord: "::0/64",
170+
ExpectedFound: true,
171171
},
172172
{
173-
IP: netip.MustParseAddr("::200.0.2.1"),
174-
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
175-
ExpectedCIDR: "::/64",
176-
ExpectedRecord: "::0/64",
177-
ExpectedOK: true,
173+
IP: netip.MustParseAddr("::200.0.2.1"),
174+
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
175+
ExpectedNetwork: "::/64",
176+
ExpectedRecord: "::0/64",
177+
ExpectedFound: true,
178178
},
179179
{
180-
IP: netip.MustParseAddr("0:0:0:0:ffff:ffff:ffff:ffff"),
181-
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
182-
ExpectedCIDR: "::/64",
183-
ExpectedRecord: "::0/64",
184-
ExpectedOK: true,
180+
IP: netip.MustParseAddr("0:0:0:0:ffff:ffff:ffff:ffff"),
181+
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
182+
ExpectedNetwork: "::/64",
183+
ExpectedRecord: "::0/64",
184+
ExpectedFound: true,
185185
},
186186
{
187-
IP: netip.MustParseAddr("ef00::"),
188-
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
189-
ExpectedCIDR: "8000::/1",
190-
ExpectedRecord: nil,
191-
ExpectedOK: false,
187+
IP: netip.MustParseAddr("ef00::"),
188+
DBFile: "MaxMind-DB-no-ipv4-search-tree.mmdb",
189+
ExpectedNetwork: "8000::/1",
190+
ExpectedRecord: nil,
191+
ExpectedFound: false,
192192
},
193193
}
194194

@@ -198,10 +198,12 @@ func TestLookupNetwork(t *testing.T) {
198198
reader, err := Open(testFile(test.DBFile))
199199
require.NoError(t, err)
200200

201-
network, ok, err := reader.LookupNetwork(test.IP, &record)
202-
require.NoError(t, err)
203-
assert.Equal(t, test.ExpectedOK, ok)
204-
assert.Equal(t, test.ExpectedCIDR, network.String())
201+
result := reader.Lookup(test.IP)
202+
require.NoError(t, result.Err())
203+
assert.Equal(t, test.ExpectedFound, result.Found())
204+
assert.Equal(t, test.ExpectedNetwork, result.Network().String())
205+
206+
require.NoError(t, result.Decode(&record))
205207
assert.Equal(t, test.ExpectedRecord, record)
206208
})
207209
}
@@ -819,21 +821,23 @@ func BenchmarkInterfaceLookup(b *testing.B) {
819821
require.NoError(b, db.Close(), "error on close")
820822
}
821823

822-
func BenchmarkInterfaceLookupNetwork(b *testing.B) {
824+
func BenchmarkLookupNetwork(b *testing.B) {
823825
db, err := Open("GeoLite2-City.mmdb")
824826
require.NoError(b, err)
825827

826828
//nolint:gosec // this is a test
827829
r := rand.New(rand.NewSource(time.Now().UnixNano()))
828-
var result any
829830

830831
s := make(net.IP, 4)
831832
for i := 0; i < b.N; i++ {
832833
ip := randomIPv4Address(r, s)
833-
_, _, err = db.LookupNetwork(ip, &result)
834-
if err != nil {
834+
res := db.Lookup(ip)
835+
if err := res.Err(); err != nil {
835836
b.Error(err)
836837
}
838+
if !res.Network().IsValid() {
839+
b.Fatalf("invalid network for %s", ip)
840+
}
837841
}
838842
require.NoError(b, db.Close(), "error on close")
839843
}
@@ -907,19 +911,18 @@ func BenchmarkCityLookup(b *testing.B) {
907911
require.NoError(b, db.Close(), "error on close")
908912
}
909913

910-
func BenchmarkCityLookupNetwork(b *testing.B) {
914+
func BenchmarkCityLookupOnly(b *testing.B) {
911915
db, err := Open("GeoLite2-City.mmdb")
912916
require.NoError(b, err)
913917

914918
//nolint:gosec // this is a test
915919
r := rand.New(rand.NewSource(time.Now().UnixNano()))
916-
var result fullCity
917920

918921
s := make(net.IP, 4)
919922
for i := 0; i < b.N; i++ {
920923
ip := randomIPv4Address(r, s)
921-
_, _, err = db.LookupNetwork(ip, &result)
922-
if err != nil {
924+
result := db.Lookup(ip)
925+
if err := result.Err(); err != nil {
923926
b.Error(err)
924927
}
925928
}

0 commit comments

Comments
 (0)