resolve dangliing hashes

store/boltdb.go

@@ -16,21 +16,32 @@ type BoltStore struct {
 	db *bbolt.DB
 }
 
-func NewBoltStore(path string) (*BoltStore, error) {
-	db, err := bbolt.Open(path, 0600, nil)
+func NewBoltStore(path string, opts *bbolt.Options) (*BoltStore, error) {
+	db, err := bbolt.Open(path, 0600, opts)
 	if err != nil {
 		return nil, err
 	}
 
-	err = db.Update(func(tx *bbolt.Tx) error {
-		if _, err := tx.CreateBucketIfNotExists(pastesBucket); err != nil {
-			return err
+	bucketsExist := false
+	db.View(func(tx *bbolt.Tx) error {
+		if tx.Bucket(pastesBucket) != nil && tx.Bucket(hashesBucket) != nil {
+			bucketsExist = true
 		}
-		_, err := tx.CreateBucketIfNotExists(hashesBucket)
-		return err
+		return nil
 	})
-	if err != nil {
-		return nil, err
+
+	if !bucketsExist {
+		err = db.Update(func(tx *bbolt.Tx) error {
+			if _, err := tx.CreateBucketIfNotExists(pastesBucket); err != nil {
+				return err
+			}
+			_, err := tx.CreateBucketIfNotExists(hashesBucket)
+			return err
+		})
+		if err != nil {
+			db.Close()
+			return nil, err
+		}
 	}
 
 	return &BoltStore{db: db}, nil
@@ -70,7 +81,7 @@ func (s *BoltStore) GetIDByHash(hash string) (string, bool, error) {
 	return id, exists, err
 }
 
-func (s *BoltStore) Set(id, hash, content string) error {
+func (s *BoltStore) Set(id, hash, content string, metadata map[string]interface{}) error {
 	return s.db.Update(func(tx *bbolt.Tx) error {
 		pb := tx.Bucket(pastesBucket)
 		hb := tx.Bucket(hashesBucket)
@@ -78,6 +89,8 @@ func (s *BoltStore) Set(id, hash, content string) error {
 		paste := Paste{
 			Content:   content,
 			CreatedAt: time.Now(),
+			Hash:      hash,
+			Metadata:  metadata,
 		}
 		encoded, err := json.Marshal(paste)
 		if err != nil {
@@ -93,12 +106,23 @@ func (s *BoltStore) Set(id, hash, content string) error {
 
 func (s *BoltStore) Del(id string) error {
 	return s.db.Update(func(tx *bbolt.Tx) error {
-		return tx.Bucket(pastesBucket).Delete([]byte(id))
+		pb := tx.Bucket(pastesBucket)
+		hb := tx.Bucket(hashesBucket)
+
+		val := pb.Get([]byte(id))
+		if val != nil {
+			var p Paste
+			if err := json.Unmarshal(val, &p); err == nil && p.Hash != "" {
+				hb.Delete([]byte(p.Hash))
+			}
+		}
+		return pb.Delete([]byte(id))
 	})
 }
 
 func (s *BoltStore) Cleanup(maxAge time.Duration) {
 	var keysToDelete [][]byte
+	var hashesToDelete [][]byte
 	s.db.View(func(tx *bbolt.Tx) error {
 		b := tx.Bucket(pastesBucket)
 		c := b.Cursor()
@@ -107,6 +131,9 @@ func (s *BoltStore) Cleanup(maxAge time.Duration) {
 			if err := json.Unmarshal(v, &p); err == nil {
 				if time.Since(p.CreatedAt) > maxAge {
 					keysToDelete = append(keysToDelete, k)
+					if p.Hash != "" {
+						hashesToDelete = append(hashesToDelete, []byte(p.Hash))
+					}
 				}
 			}
 		}
@@ -115,9 +142,13 @@ func (s *BoltStore) Cleanup(maxAge time.Duration) {
 
 	if len(keysToDelete) > 0 {
 		s.db.Update(func(tx *bbolt.Tx) error {
-			b := tx.Bucket(pastesBucket)
+			pb := tx.Bucket(pastesBucket)
+			hb := tx.Bucket(hashesBucket)
 			for _, k := range keysToDelete {
-				b.Delete(k)
+				pb.Delete(k)
+			}
+			for _, h := range hashesToDelete {
+				hb.Delete(h)
 			}
 			return nil
 		})