fix(homeserver): add a reference counter for chunks to avoid deleting shared blobs

pubky-homeserver/src/database/tables/blobs.rs

@@ -26,7 +26,7 @@ impl DB {
             self.tables
                 .blobs
                 .get(&rtxn, entry.content_hash())?
-                .map(|blob| bytes::Bytes::from(blob.to_vec()))
+                .map(|blob| bytes::Bytes::from(blob[8..].to_vec()))
         } else {
             None
         };

pubky-homeserver/src/database/tables/entries.rs

@@ -43,7 +43,26 @@ impl DB {
 
         let hash = hasher.finalize();
 
-        self.tables.blobs.put(&mut wtxn, hash.as_bytes(), &bytes)?;
+        let key = hash.as_bytes();
+
+        let mut bytes_with_ref_count = Vec::with_capacity(bytes.len() + 8);
+        bytes_with_ref_count.extend_from_slice(&u64::to_be_bytes(0));
+        bytes_with_ref_count.extend_from_slice(&bytes);
+
+        // TODO: For now, we set the first 8 bytes to a reference counter
+        let exists = self
+            .tables
+            .blobs
+            .get(&wtxn, key)?
+            .unwrap_or(bytes_with_ref_count.as_slice());
+
+        let new_count = u64::from_be_bytes(exists[0..8].try_into().unwrap()) + 1;
+
+        bytes_with_ref_count[0..8].copy_from_slice(&u64::to_be_bytes(new_count));
+
+        self.tables
+            .blobs
+            .put(&mut wtxn, hash.as_bytes(), &bytes_with_ref_count)?;
 
         let mut entry = Entry::new();
 
@@ -82,6 +101,29 @@ impl DB {
         let deleted = if let Some(bytes) = self.tables.entries.get(&wtxn, &key)? {
             let entry = Entry::deserialize(bytes)?;
 
+            let mut bytes_with_ref_count = self
+                .tables
+                .blobs
+                .get(&wtxn, entry.content_hash())?
+                .map_or(vec![], |s| s.to_vec());
+
+            let arr: [u8; 8] = bytes_with_ref_count[0..8].try_into().unwrap_or([0; 8]);
+            let reference_count = u64::from_be_bytes(arr);
+
+            if reference_count > 1 {
+                // decrement reference count
+
+                bytes_with_ref_count[0..8].copy_from_slice(&(reference_count - 1).to_be_bytes());
+
+                self.tables
+                    .blobs
+                    .put(&mut wtxn, entry.content_hash(), &bytes_with_ref_count)?;
+
+                let deleted_entry = self.tables.entries.delete(&mut wtxn, &key)?;
+
+                return Ok(deleted_entry);
+            }
+
             // TODO: reference counting of blobs
             let deleted_blobs = self.tables.blobs.delete(&mut wtxn, entry.content_hash())?;
 

pubky/src/shared/public.rs

@@ -765,4 +765,37 @@ mod tests {
         let get = client.get(url.as_str()).await.unwrap();
         dbg!(get);
     }
+
+    #[tokio::test]
+    async fn dont_delete_shared_blobs() {
+        let testnet = Testnet::new(10);
+        let homeserver = Homeserver::start_test(&testnet).await.unwrap();
+        let client = PubkyClient::test(&testnet);
+
+        // Step 1: Create first user (follower)
+        let keypair = Keypair::random();
+
+        let user_id = keypair.public_key().to_z32();
+        client
+            .signup(&keypair, &homeserver.public_key())
+            .await
+            .unwrap();
+
+        // Both files are identical, leads to error
+        let file_1 = vec![1];
+        let file_2 = vec![1];
+
+        let url_1 = format!("pubky://{}/pub/pubky.app/file/file_1", user_id);
+        let url_2 = format!("pubky://{}/pub/pubky.app/file/file_1", user_id);
+
+        client.put(url_1.as_str(), &file_1).await.unwrap();
+        client.put(url_2.as_str(), &file_2).await.unwrap();
+
+        // Delete file 1
+        client.delete(url_1.as_str()).await.unwrap();
+
+        let blob = client.get(url_2.as_str()).await.unwrap().unwrap();
+
+        assert_eq!(blob, vec![1])
+    }
 }