108 lines
2 KiB
Go
108 lines
2 KiB
Go
|
package loader
|
||
|
|
|
||
|
|
import (
|
||
|
|
"path"
|
||
|
|
|
||
|
|
"github.com/anacrolix/torrent/metainfo"
|
||
|
|
"github.com/dgraph-io/badger/v3"
|
||
|
|
dlog "github.com/distribyted/distribyted/log"
|
||
|
|
"github.com/rs/zerolog/log"
|
||
|
|
)
|
||
|
|
|
||
|
|
var _ LoaderAdder = &DB{}
|
||
|
|
|
||
|
|
const routeRootKey = "/route/"
|
||
|
|
|
||
|
|
type DB struct {
|
||
|
|
db *badger.DB
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewDB(path string) (*DB, error) {
|
||
|
|
l := log.Logger.With().Str("component", "torrent-store").Logger()
|
||
|
|
db, err := badger.Open(badger.DefaultOptions(path).WithLogger(&dlog.Badger{L: l}))
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
err = db.RunValueLogGC(0.5)
|
||
|
|
if err != nil && err != badger.ErrNoRewrite {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return &DB{
|
||
|
|
db: db,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *DB) AddMagnet(r, m string) error {
|
||
|
|
err := l.db.Update(func(txn *badger.Txn) error {
|
||
|
|
spec, err := metainfo.ParseMagnetUri(m)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
ih := spec.InfoHash.HexString()
|
||
|
|
|
||
|
|
rp := path.Join(routeRootKey, ih, r)
|
||
|
|
return txn.Set([]byte(rp), []byte(m))
|
||
|
|
})
|
||
|
|
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
return l.db.Sync()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *DB) RemoveFromHash(r, h string) (bool, error) {
|
||
|
|
tx := l.db.NewTransaction(true)
|
||
|
|
defer tx.Discard()
|
||
|
|
|
||
|
|
var mh metainfo.Hash
|
||
|
|
if err := mh.FromHexString(h); err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
|
||
|
|
rp := path.Join(routeRootKey, h, r)
|
||
|
|
if _, err := tx.Get([]byte(rp)); err != nil {
|
||
|
|
return false, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := tx.Delete([]byte(rp)); err != nil {
|
||
|
|
return false, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return true, tx.Commit()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *DB) ListMagnets() (map[string][]string, error) {
|
||
|
|
tx := l.db.NewTransaction(false)
|
||
|
|
defer tx.Discard()
|
||
|
|
|
||
|
|
it := tx.NewIterator(badger.DefaultIteratorOptions)
|
||
|
|
defer it.Close()
|
||
|
|
|
||
|
|
prefix := []byte(routeRootKey)
|
||
|
|
out := make(map[string][]string)
|
||
|
|
for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() {
|
||
|
|
_, r := path.Split(string(it.Item().Key()))
|
||
|
|
i := it.Item()
|
||
|
|
if err := i.Value(func(v []byte) error {
|
||
|
|
out[r] = append(out[r], string(v))
|
||
|
|
return nil
|
||
|
|
}); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return out, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *DB) ListTorrentPaths() (map[string][]string, error) {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (l *DB) Close() error {
|
||
|
|
return l.db.Close()
|
||
|
|
}
|