observer-bot/db.go

package main

import (
	"encoding/json"
	"git.gammaspectra.live/P2Pool/consensus/v3/monero/address"
	bolt "go.etcd.io/bbolt"
	"log"
	"slices"
	"strings"
	"time"
)

type Subscription struct {
	Address *address.Address `json:"address"`
	Nick    string           `json:"nick"`
}

type DB struct {
	db *bolt.DB
}

var refByNick = []byte("refByNick")
var refByAddr = []byte("refByAddr")

func NewDB(path string) (*DB, error) {
	if db, err := bolt.Open(path, 0666, &bolt.Options{Timeout: time.Second * 5}); err != nil {
		return nil, err
	} else {
		if err = db.Update(func(tx *bolt.Tx) error {
			if _, err := tx.CreateBucketIfNotExists(refByNick); err != nil {
				return err
			} else if _, err := tx.CreateBucketIfNotExists(refByAddr); err != nil {
				return err
			}
			return nil
		}); err != nil {
			return nil, err
		}
		return &DB{
			db: db,
		}, nil
	}
}

func (db *DB) GetByNick(nick string) (result []*Subscription) {
	_ = db.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(refByNick)
		if buf := b.Get([]byte(strings.ToLower(nick))); buf != nil {
			var addrs []*address.Address
			if err := json.Unmarshal(buf, &addrs); err != nil {
				return err
			}
			for _, a := range addrs {
				result = append(result, &Subscription{
					Address: a,
					Nick:    nick,
				})
			}
		}
		return nil
	})
	return result
}

func (db *DB) GetByAddress(a *address.Address) (result []*Subscription) {
	_ = db.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(refByAddr)
		if buf := b.Get([]byte(a.ToBase58())); buf != nil {
			var nicks []string
			if err := json.Unmarshal(buf, &nicks); err != nil {
				return err
			}
			for _, nick := range nicks {
				result = append(result, &Subscription{
					Address: a,
					Nick:    nick,
				})
			}
		}
		return nil
	})
	return result
}

func (db *DB) Unsubscribe(sub *Subscription) error {

	if err := db.db.Update(func(tx *bolt.Tx) error {

		b1 := tx.Bucket(refByAddr)
		var nicks []string
		if k := b1.Get([]byte(sub.Address.ToBase58())); k != nil {
			if err := json.Unmarshal(k, &nicks); err != nil {
				return err
			}
		}
		if i := slices.Index(nicks, strings.ToLower(sub.Nick)); i != -1 {
			nicks = slices.Delete(nicks, i, i+1)
			if len(nicks) == 0 {
				if err := b1.Delete([]byte(sub.Address.ToBase58())); err != nil {
					return err
				}
			} else {
				if buf, err := json.Marshal(nicks); err != nil {
					return err
				} else if err = b1.Put([]byte(sub.Address.ToBase58()), buf); err != nil {
					return err
				}
			}
		}

		b2 := tx.Bucket(refByNick)
		var addresses []*address.Address
		if k := b2.Get([]byte(strings.ToLower(sub.Nick))); k != nil {
			if err := json.Unmarshal(k, &addresses); err != nil {
				return err
			}
		}
		if i := slices.IndexFunc(addresses, func(a *address.Address) bool {
			return a.Compare(sub.Address) == 0
		}); i != -1 {
			addresses = slices.Delete(addresses, i, i+1)

			if len(addresses) == 0 {
				if err := b2.Delete([]byte(strings.ToLower(sub.Nick))); err != nil {
					return err
				}
			} else {
				if buf, err := json.Marshal(addresses); err != nil {
					return err
				} else if err = b2.Put([]byte(strings.ToLower(sub.Nick)), buf); err != nil {
					return err
				}
			}
		}

		return nil
	}); err != nil {
		log.Printf("[DB] bolt error: %s", err)
		return err
	}
	return nil
}

func (db *DB) Store(sub *Subscription) error {

	if err := db.db.Update(func(tx *bolt.Tx) error {

		b1 := tx.Bucket(refByAddr)
		var nicks []string
		if k := b1.Get([]byte(sub.Address.ToBase58())); k != nil {
			if err := json.Unmarshal(k, &nicks); err != nil {
				return err
			}
		}
		if !slices.Contains(nicks, strings.ToLower(sub.Nick)) {
			nicks = append(nicks, strings.ToLower(sub.Nick))
		}
		if buf, err := json.Marshal(nicks); err != nil {
			return err
		} else if err = b1.Put([]byte(sub.Address.ToBase58()), buf); err != nil {
			return err
		}

		b2 := tx.Bucket(refByNick)
		var addresses []*address.Address
		if k := b2.Get([]byte(strings.ToLower(sub.Nick))); k != nil {
			if err := json.Unmarshal(k, &addresses); err != nil {
				return err
			}
		}
		if !slices.ContainsFunc(addresses, func(a *address.Address) bool {
			return a.Compare(sub.Address) == 0
		}) {
			addresses = append(addresses, sub.Address)
		}
		if buf, err := json.Marshal(addresses); err != nil {
			return err
		} else if err = b2.Put([]byte(strings.ToLower(sub.Nick)), buf); err != nil {
			return err
		}

		return nil
	}); err != nil {
		log.Printf("[DB] bolt error: %s", err)
		return err
	}
	return nil
}