Release v0.1.0

vendor/github.com/lightninglabs/neutrino/banman/store.go

@@ -0,0 +1,221 @@
+package banman
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/btcsuite/btcwallet/walletdb"
+)
+
+var (
+	// byteOrder is the preferred byte order in which we should write things
+	// to disk.
+	byteOrder = binary.BigEndian
+
+	// banStoreBucket is the top level bucket of the Store that will contain
+	// all relevant sub-buckets.
+	banStoreBucket = []byte("ban-store")
+
+	// banBucket is the main index in which we keep track of IP networks and
+	// their absolute expiration time.
+	//
+	// The key is the IP network host and the value is the absolute
+	// expiration time.
+	banBucket = []byte("ban-index")
+
+	// reasonBucket is an index in which we keep track of why an IP network
+	// was banned.
+	//
+	// The key is the IP network and the value is the Reason.
+	reasonBucket = []byte("reason-index")
+
+	// ErrCorruptedStore is an error returned when we attempt to locate any
+	// of the ban-related buckets in the database but are unable to.
+	ErrCorruptedStore = errors.New("corrupted ban store")
+
+	// ErrUnsupportedIP is an error returned when we attempt to parse an
+	// unsupported IP address type.
+	ErrUnsupportedIP = errors.New("unsupported IP type")
+)
+
+// Status gathers all of the details regarding an IP network's ban status.
+type Status struct {
+	// Banned determines whether the IP network is currently banned.
+	Banned bool
+
+	// Reason is the reason for which the IP network was banned.
+	Reason Reason
+
+	// Expiration is the absolute time in which the ban will expire.
+	Expiration time.Time
+}
+
+// Store is the store responsible for maintaining records of banned IP networks.
+// It uses IP networks, rather than single IP addresses, in order to coalesce
+// multiple IP addresses that are likely to be correlated.
+type Store interface {
+	// BanIPNet creates a ban record for the IP network within the store for
+	// the given duration. A reason can also be provided to note why the IP
+	// network is being banned. The record will exist until a call to Status
+	// is made after the ban expiration.
+	BanIPNet(*net.IPNet, Reason, time.Duration) error
+
+	// Status returns the ban status for a given IP network.
+	Status(*net.IPNet) (Status, error)
+}
+
+// NewStore returns a Store backed by a database.
+func NewStore(db walletdb.DB) (Store, error) {
+	return newBanStore(db)
+}
+
+// banStore is a concrete implementation of the Store interface backed by a
+// database.
+type banStore struct {
+	db walletdb.DB
+}
+
+// A compile-time constraint to ensure banStore satisfies the Store interface.
+var _ Store = (*banStore)(nil)
+
+// newBanStore creates a concrete implementation of the Store interface backed
+// by a database.
+func newBanStore(db walletdb.DB) (*banStore, error) {
+	s := &banStore{db: db}
+
+	// We'll ensure the expected buckets are created upon initialization.
+	err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+		banStore, err := tx.CreateTopLevelBucket(banStoreBucket)
+		if err != nil {
+			return err
+		}
+		_, err = banStore.CreateBucketIfNotExists(banBucket)
+		if err != nil {
+			return err
+		}
+		_, err = banStore.CreateBucketIfNotExists(reasonBucket)
+		return err
+	})
+	if err != nil && err != walletdb.ErrBucketExists {
+		return nil, err
+	}
+
+	return s, nil
+}
+
+// BanIPNet creates a ban record for the IP network within the store for the
+// given duration. A reason can also be provided to note why the IP network is
+// being banned. The record will exist until a call to Status is made after the
+// ban expiration.
+func (s *banStore) BanIPNet(ipNet *net.IPNet, reason Reason, duration time.Duration) error {
+	return walletdb.Update(s.db, func(tx walletdb.ReadWriteTx) error {
+		banStore := tx.ReadWriteBucket(banStoreBucket)
+		if banStore == nil {
+			return ErrCorruptedStore
+		}
+		banIndex := banStore.NestedReadWriteBucket(banBucket)
+		if banIndex == nil {
+			return ErrCorruptedStore
+		}
+		reasonIndex := banStore.NestedReadWriteBucket(reasonBucket)
+		if reasonIndex == nil {
+			return ErrCorruptedStore
+		}
+
+		var ipNetBuf bytes.Buffer
+		if err := encodeIPNet(&ipNetBuf, ipNet); err != nil {
+			return fmt.Errorf("unable to encode %v: %v", ipNet, err)
+		}
+		k := ipNetBuf.Bytes()
+
+		return addBannedIPNet(banIndex, reasonIndex, k, reason, duration)
+	})
+}
+
+// addBannedIPNet adds an entry to the ban store for the given IP network.
+func addBannedIPNet(banIndex, reasonIndex walletdb.ReadWriteBucket,
+	ipNetKey []byte, reason Reason, duration time.Duration) error {
+
+	var v [8]byte
+	banExpiration := time.Now().Add(duration)
+	byteOrder.PutUint64(v[:], uint64(banExpiration.Unix()))
+
+	if err := banIndex.Put(ipNetKey, v[:]); err != nil {
+		return err
+	}
+	return reasonIndex.Put(ipNetKey, []byte{byte(reason)})
+}
+
+// Status returns the ban status for a given IP network.
+func (s *banStore) Status(ipNet *net.IPNet) (Status, error) {
+	var banStatus Status
+	err := walletdb.Update(s.db, func(tx walletdb.ReadWriteTx) error {
+		banStore := tx.ReadWriteBucket(banStoreBucket)
+		if banStore == nil {
+			return ErrCorruptedStore
+		}
+		banIndex := banStore.NestedReadWriteBucket(banBucket)
+		if banIndex == nil {
+			return ErrCorruptedStore
+		}
+		reasonIndex := banStore.NestedReadWriteBucket(reasonBucket)
+		if reasonIndex == nil {
+			return ErrCorruptedStore
+		}
+
+		var ipNetBuf bytes.Buffer
+		if err := encodeIPNet(&ipNetBuf, ipNet); err != nil {
+			return fmt.Errorf("unable to encode %v: %v", ipNet, err)
+		}
+		k := ipNetBuf.Bytes()
+
+		status := fetchStatus(banIndex, reasonIndex, k)
+
+		// If the IP network's ban duration has expired, we can remove
+		// its entry from the store.
+		if !time.Now().Before(status.Expiration) {
+			return removeBannedIPNet(banIndex, reasonIndex, k)
+		}
+
+		banStatus = status
+		return nil
+	})
+	if err != nil {
+		return Status{}, err
+	}
+
+	return banStatus, nil
+}
+
+// fetchStatus retrieves the ban status of the given IP network.
+func fetchStatus(banIndex, reasonIndex walletdb.ReadWriteBucket,
+	ipNetKey []byte) Status {
+
+	v := banIndex.Get(ipNetKey)
+	if v == nil {
+		return Status{}
+	}
+	reason := Reason(reasonIndex.Get(ipNetKey)[0])
+	banExpiration := time.Unix(int64(byteOrder.Uint64(v)), 0)
+
+	return Status{
+		Banned:     true,
+		Reason:     reason,
+		Expiration: banExpiration,
+	}
+}
+
+// removeBannedIPNet removes all references to a banned IP network within the
+// ban store.
+func removeBannedIPNet(banIndex, reasonIndex walletdb.ReadWriteBucket,
+	ipNetKey []byte) error {
+
+	if err := banIndex.Delete(ipNetKey); err != nil {
+		return err
+	}
+	return reasonIndex.Delete(ipNetKey)
+}