godns/hosts/hosts_bolt.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
}