fix: add file ref count

GZTimeWalker · Jul 14, 2022 · f78f17c · f78f17c
1 parent 1027fb8
commit f78f17c
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 16 deletions.
diff --git a/...rations/20220713033940_InitDb.Designer.cs → ...rations/20220714063743_InitDb.Designer.cs b/...rations/20220713033940_InitDb.Designer.cs → ...rations/20220714063743_InitDb.Designer.cs
diff --git a/GZCTF/Migrations/20220713033940_InitDb.cs → GZCTF/Migrations/20220714063743_InitDb.cs b/GZCTF/Migrations/20220713033940_InitDb.cs → GZCTF/Migrations/20220714063743_InitDb.cs
@@ -53,7 +53,8 @@ protected override void Up(MigrationBuilder migrationBuilder)
                     Id = table.Column<int>(type: "integer", nullable: false)
                         .Annotation("Npgsql:ValueGenerationStrategy", NpgsqlValueGenerationStrategy.IdentityByDefaultColumn),
                     Hash = table.Column<string>(type: "character varying(64)", maxLength: 64, nullable: false),
-                    Name = table.Column<string>(type: "text", nullable: false)
+                    Name = table.Column<string>(type: "text", nullable: false),
+                    ReferenceCount = table.Column<long>(type: "bigint", nullable: false)
                 },
                 constraints: table =>
                 {

diff --git a/GZCTF/Migrations/AppDbContextModelSnapshot.cs b/GZCTF/Migrations/AppDbContextModelSnapshot.cs
@@ -344,6 +344,9 @@ protected override void BuildModel(ModelBuilder modelBuilder)
                         .IsRequired()
                         .HasColumnType("text");
 
+                    b.Property<long>("ReferenceCount")
+                        .HasColumnType("bigint");
+
                     b.HasKey("Id");
 
                     b.HasIndex("Hash");

diff --git a/GZCTF/Models/Data/LocalFile.cs b/GZCTF/Models/Data/LocalFile.cs
@@ -24,6 +24,12 @@ public class LocalFile
     [Required]
     public string Name { get; set; } = string.Empty;
 
+    /// <summary>
+    /// 引用计数
+    /// </summary>
+    [JsonIgnore]
+    public uint ReferenceCount { get; set; } = 1;
+
     /// <summary>
     /// 文件存储位置
     /// </summary>

diff --git a/GZCTF/Repositories/FileRepository.cs b/GZCTF/Repositories/FileRepository.cs
@@ -30,29 +30,33 @@ public async Task<LocalFile> CreateOrUpdateFile(IFormFile file, string? fileName
         await file.CopyToAsync(tmp, token);
 
         tmp.Position = 0;
-        var hash = await SHA256.Create().ComputeHashAsync(tmp, token);
+        var hash = await SHA384.Create().ComputeHashAsync(tmp, token);
         var fileHash = BitConverter.ToString(hash).Replace("-", "").ToLower();
 
         var localFile = await GetFileByHash(fileHash, token);
 
-        if (localFile is null)
-        {
-            localFile = new() { Hash = fileHash, Name = fileName ?? file.FileName };
-            await context.AddAsync(localFile, token);
-        }
-        else
+        if (localFile is not null)
         {
             localFile.Name = fileName ?? file.FileName;
+
+            localFile.ReferenceCount++; // same hash, add ref count
+
+            logger.SystemLog($"文件引用计数 [{localFile.Hash[..8]}] {localFile.Name} => {localFile.ReferenceCount}", TaskStatus.Success, LogLevel.Debug);
+
             context.Update(localFile);
         }
+        else
+        {
+            localFile = new() { Hash = fileHash, Name = fileName ?? file.FileName };
+            await context.AddAsync(localFile, token);
 
-        var path = Path.Combine(uploadPath, localFile.Location);
+            var path = Path.Combine(uploadPath, localFile.Location);
 
-        if (!Directory.Exists(path))
-            Directory.CreateDirectory(path);
+            if (!Directory.Exists(path))
+                Directory.CreateDirectory(path);
+
+            using var fileStream = File.Create(Path.Combine(path, localFile.Hash));
 
-        using (var fileStream = File.Create(Path.Combine(path, localFile.Hash)))
-        {
             tmp.Position = 0;
             await tmp.CopyToAsync(fileStream, token);
         }
@@ -64,14 +68,26 @@ public async Task<LocalFile> CreateOrUpdateFile(IFormFile file, string? fileName
 
     public async Task<TaskStatus> DeleteFileByHash(string fileHash, CancellationToken token = default)
     {
-        var file = await context.Files.Where(f => f.Hash == fileHash).FirstOrDefaultAsync(token);
+        var file = await GetFileByHash(fileHash, token);
 
         if (file is null)
             return TaskStatus.NotFound;
 
         var path = Path.Combine(uploadPath, file.Location, file.Hash);
 
-        logger.SystemLog($"删除文件 [{file.Hash[..8]}] {file.Name}...", TaskStatus.Pending, LogLevel.Information);
+        if (file.ReferenceCount > 1)
+        {
+            file.ReferenceCount--; // other ref exists, decrease ref count
+
+            logger.SystemLog($"文件引用计数 [{file.Hash[..8]}] {file.Name} => {file.ReferenceCount}", TaskStatus.Success, LogLevel.Debug);
+
+            context.Update(file);
+            await context.SaveChangesAsync(token);
+
+            return TaskStatus.Success;
+        }
+
+        logger.SystemLog($"删除文件 [{file.Hash[..8]}] {file.Name}", TaskStatus.Pending, LogLevel.Information);
 
         if (File.Exists(path))
         {