Merge pull request litedb-org#2459 from SpaceCheetah/FixSeek

Fix LiteFileStream.SetReadStreamPosition
jiuaizxc · Jun 11, 2024 · 409f6d5 · 409f6d5
2 parents e60b6ae + 4d58c69
commit 409f6d5
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 7 deletions.
diff --git a/LiteDB.Tests/Issues/Issue2458_Tests.cs b/LiteDB.Tests/Issues/Issue2458_Tests.cs
@@ -0,0 +1,54 @@
+using System;
+using System.IO;
+using Xunit;
+
+namespace LiteDB.Tests.Issues;
+
+public class Issue2458_Tests
+{
+    [Fact]
+    public void NegativeSeekFails()
+    {
+        using var db = new LiteDatabase(":memory:");
+        var fs = db.FileStorage;
+        AddTestFile("test", 1, fs);
+        using Stream stream = fs.OpenRead("test");
+        Assert.Throws<ArgumentOutOfRangeException>(() => stream.Position = -1);
+    }
+
+    //https://learn.microsoft.com/en-us/dotnet/api/system.io.stream.position?view=net-8.0 says seeking to a position
+    //beyond the end of a stream is supported, so implementations should support it (error on read).
+    [Fact]
+    public void SeekPastFileSucceds()
+    {
+        using var db = new LiteDatabase(":memory:");
+        var fs = db.FileStorage;
+        AddTestFile("test", 1, fs);
+        using Stream stream = fs.OpenRead("test");
+        stream.Position = Int32.MaxValue;
+    }
+
+    [Fact]
+    public void SeekShortChunks()
+    {
+        using var db = new LiteDatabase(":memory:");
+        var fs = db.FileStorage;
+        using(Stream writeStream = fs.OpenWrite("test", "test"))
+        {
+            writeStream.WriteByte(0);
+            writeStream.Flush(); //Create single-byte chunk just containing a 0
+            writeStream.WriteByte(1);
+            writeStream.Flush();
+            writeStream.WriteByte(2);
+        }
+        using Stream readStream = fs.OpenRead("test");
+        readStream.Position = 2;
+        Assert.Equal(2, readStream.ReadByte());
+    }
+
+    private void AddTestFile(string id, long length, ILiteStorage<string> fs)
+    {
+        using Stream writeStream = fs.OpenWrite(id, id);
+        writeStream.Write(new byte[length]);
+    }
+}
diff --git a/LiteDB/Client/Storage/LiteFileStream.Read.cs b/LiteDB/Client/Storage/LiteFileStream.Read.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using static LiteDB.Constants;
@@ -7,9 +8,14 @@ namespace LiteDB
 {
     public partial class LiteFileStream<TFileId> : Stream
     {
+        private Dictionary<int, long> _chunkLengths = new Dictionary<int, long>();
         public override int Read(byte[] buffer, int offset, int count)
         {
             if (_mode != FileAccess.Read) throw new NotSupportedException();
+            if (_streamPosition == Length)
+            {
+                return 0;
+            }
 
             var bytesLeft = count;
 
@@ -42,23 +48,54 @@ private byte[] GetChunkData(int index)
                 .FindOne("_id = { f: @0, n: @1 }", _fileId, index);
 
             // if chunk is null there is no more chunks
-            return chunk?["data"].AsBinary;
+            byte[] result = chunk?["data"].AsBinary;
+            if (result != null)
+            {
+                _chunkLengths[index] = result.Length;
+            }
+            return result;
         }
 
         private void SetReadStreamPosition(long newPosition)
         {
-            if (newPosition < 0 && newPosition > Length)
+            if (newPosition < 0)
             {
                 throw new ArgumentOutOfRangeException();
             }
+            if (newPosition >= Length)
+            {
+                _streamPosition = Length;
+                return;
+            }
             _streamPosition = newPosition;
 
             // calculate new chunk position
-            _currentChunkIndex = (int)_streamPosition / MAX_CHUNK_SIZE;
-            _positionInChunk = (int)_streamPosition % MAX_CHUNK_SIZE;
-
-            // get current chunk
-            _currentChunkData = this.GetChunkData(_currentChunkIndex);
+            long seekStreamPosition = 0;
+            int loadedChunk = _currentChunkIndex;
+            int newChunkIndex = 0;
+            while (seekStreamPosition <= _streamPosition)
+            {
+                if (_chunkLengths.TryGetValue(newChunkIndex, out long length))
+                {
+                    seekStreamPosition += length;
+                }
+                else
+                {
+                    loadedChunk = newChunkIndex;
+                    _currentChunkData = GetChunkData(newChunkIndex);
+                    seekStreamPosition += _currentChunkData.Length;
+                }
+                newChunkIndex++;
+            }
+
+            newChunkIndex--;
+            seekStreamPosition -= _chunkLengths[newChunkIndex];
+            _positionInChunk = (int)(_streamPosition - seekStreamPosition);
+            _currentChunkIndex = newChunkIndex;
+            if (loadedChunk != _currentChunkIndex)
+            {
+                _currentChunkData = GetChunkData(_currentChunkIndex);
+            }
         }
     }
 }