country-geo-locations/pkg/geoloc/database.go

75 lines
1.2 KiB
Go

package geoloc
import (
"fmt"
"github.com/hashicorp/go-memdb"
)
var database *memdb.MemDB
func CreateDatabase(ipinfos []IPInfo) error {
schema := &memdb.DBSchema{
Tables: map[string]*memdb.TableSchema{
"ipinfo": {
Name: "ipinfo",
Indexes: map[string]*memdb.IndexSchema{
"id": {
Name: "id",
Unique: true,
Indexer: &memdb.UintFieldIndex{Field: "IPNumFrom"},
},
},
},
},
}
db, err := memdb.NewMemDB(schema)
if err != nil {
return err
}
txn := db.Txn(true)
defer txn.Abort()
for _, ipinfo := range ipinfos {
if err := txn.Insert("ipinfo", ipinfo); err != nil {
return err
}
}
txn.Commit()
database = db
return nil
}
func SearchIPNet(ipnetnum uint) (*IPInfo, error) {
txn := database.Txn(false)
defer txn.Abort()
raw, err := txn.First("ipinfo", "id", ipnetnum)
if err != nil {
return nil, err
}
if raw == nil {
return nil, nil
}
var ipinfo IPInfo
for {
raw, err := txn.First("ipinfo", "id", ipnetnum)
if err != nil {
return nil, err
}
if raw != nil {
ipinfo = raw.(IPInfo)
break
}
if ipnetnum == 0 {
return nil, fmt.Errorf("SearchIPNet: IP net not found")
}
ipnetnum = ipnetnum - 256
}
return &ipinfo, nil
}