125 lines
1.9 KiB
Go
125 lines
1.9 KiB
Go
package hosts
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"github.com/miekg/dns"
|
|
log "github.com/sirupsen/logrus"
|
|
bolt "go.etcd.io/bbolt"
|
|
"strings"
|
|
)
|
|
|
|
const (
|
|
recordBucket = "records"
|
|
)
|
|
|
|
type BoltHosts struct {
|
|
Provider
|
|
|
|
db *bolt.DB
|
|
}
|
|
|
|
func NewBoltProvider(file string) Provider {
|
|
db, err := bolt.Open(file, 0600, &bolt.Options{})
|
|
|
|
if err != nil {
|
|
log.WithError(err).Fatalln("Unable to open database")
|
|
}
|
|
|
|
err = db.Update(func(tx *bolt.Tx) error {
|
|
_, err := tx.CreateBucketIfNotExists([]byte(recordBucket))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
return &BoltHosts{
|
|
db: db,
|
|
}
|
|
}
|
|
|
|
func (b *BoltHosts) Get(queryType uint16, domain string) (*Host, error) {
|
|
log.Debug("Checking bolt provider for %s : %s", queryType, domain)
|
|
|
|
domain = strings.ToLower(domain)
|
|
|
|
var err error
|
|
|
|
key := domain + "_" + dns.TypeToString[queryType]
|
|
var v []byte
|
|
|
|
err = b.db.View(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte("records"))
|
|
|
|
v = b.Get([]byte(key))
|
|
|
|
if string(v) == "" {
|
|
return errors.New( "Record not found, key: " + key)
|
|
}
|
|
|
|
v = b.Get([]byte("*." + key))
|
|
|
|
if string(v) == "" {
|
|
return errors.New( "Record not found, key: " + key)
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var h []Host
|
|
|
|
if err = json.Unmarshal(v, &h); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
for _, host := range h {
|
|
if host.Type == queryType {
|
|
return &host, nil
|
|
}
|
|
}
|
|
|
|
return nil, errRecordNotFound
|
|
}
|
|
|
|
func (b *BoltHosts) Set(domain string, host *Host) error {
|
|
err := b.db.Update(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket([]byte(recordBucket))
|
|
|
|
hosts := []*Host{host}
|
|
|
|
existing := b.Get([]byte(domain))
|
|
|
|
if existing != nil {
|
|
err := json.Unmarshal(existing, &hosts)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
hosts = append(hosts, host)
|
|
}
|
|
|
|
hostBytes, err := json.Marshal(hosts)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = b.Put([]byte(domain), hostBytes)
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
return err
|
|
}
|