89 lines
1.5 KiB
Go
89 lines
1.5 KiB
Go
package database
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"git.ar21.de/yolokube/country-geo-locations/pkg/geoloc"
|
|
"github.com/hashicorp/go-memdb"
|
|
)
|
|
|
|
var (
|
|
ErrUnknownInterface = errors.New("unknown interface structure")
|
|
ErrIPNetNotFound = errors.New("IP net not found")
|
|
)
|
|
|
|
type Database struct {
|
|
db *memdb.MemDB
|
|
}
|
|
|
|
func NewDatabase() (*Database, error) {
|
|
database, err := memdb.NewMemDB(
|
|
&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"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Database{
|
|
db: database,
|
|
}, nil
|
|
}
|
|
|
|
func (d *Database) Load(ipinfos []geoloc.IPInfo) error {
|
|
txn := d.db.Txn(true)
|
|
defer txn.Abort()
|
|
|
|
for _, ipinfo := range ipinfos {
|
|
if err := txn.Insert("ipinfo", ipinfo); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
txn.Commit()
|
|
return nil
|
|
}
|
|
|
|
func (d *Database) SearchIPNet(ipnetnum uint) (*geoloc.IPInfo, error) {
|
|
txn := d.db.Txn(false)
|
|
defer txn.Abort()
|
|
|
|
var (
|
|
ipinfo geoloc.IPInfo
|
|
ok bool
|
|
)
|
|
for {
|
|
raw, err := txn.First("ipinfo", "id", ipnetnum)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if raw != nil {
|
|
ipinfo, ok = raw.(geoloc.IPInfo)
|
|
if !ok {
|
|
return nil, ErrUnknownInterface
|
|
}
|
|
break
|
|
}
|
|
|
|
if ipnetnum == 0 {
|
|
return nil, ErrIPNetNotFound
|
|
}
|
|
|
|
ipnetnum -= geoloc.CalculationValue
|
|
}
|
|
|
|
return &ipinfo, nil
|
|
}
|