ISCSIConsole-Fork/DiskAccessLibrary/Mod/BlockDifferencingDiskImage/BlockDifferencingDiskImage.cs

using DiskAccessLibrary.Mod.BlockDifferencingDiskImage;
using System;
using System.IO;
using System.Security.Authentication;
using System.Text;

namespace DiskAccessLibrary
{
    public class BlockDifferencingDiskImage : DiskImage
    {
        private const int SectorSize = 512;

        private BddInfo BddInfo { get; }

        public bool IsNoBaseImage => null == _basedFileStream;

        public int NumberOfBlocks => BddInfo.NumberOfBlocks;
        public int BlockSize => BddInfo.BlockSize;

        public IReadOnlyIndexer<int, bool> IsBlockAllocated { get; }

        private bool _isExclusiveLock;

        private FileStream _basedFileStream;
        private FileStream _snapshotFileStream;

        public BlockDifferencingDiskImage(string diskImagePath, bool isReadOnly) : base(diskImagePath, isReadOnly)
        {
            BddInfo = new BddInfo(diskImagePath);
            IsBlockAllocated = new IReadOnlyIndexer<int, bool>(i => BddInfo.Allocated[i]);
        }

        public void OpenForRead()
        {
            _snapshotFileStream = File.Open(Path, FileMode.Open, FileAccess.Read, FileShare.ReadWrite);
        }

        public void CloseForRead()
        {
            _snapshotFileStream?.Close();
        }

        public void ReadBlock(byte[] buffer, int offset, int count, int blockIndex, int blockOffset)
        {
            if (false == BddInfo.Allocated[blockIndex])
            {
                if (null == _basedFileStream) return;

                _basedFileStream.Position = (long)blockIndex * BddInfo.BlockSize + blockOffset;

                while (count > 0)
                {
                    var read = _basedFileStream.Read(buffer, offset, count);
                    offset += read;
                    count -= read;
                }
            }
            else
            {
                _snapshotFileStream.Position = BddInfo.Offset[blockIndex] + blockOffset;
                while (count > 0)
                {
                    var read = _snapshotFileStream.Read(buffer, offset, count);
                    offset += read;
                    count -= read;
                }
            }
        }

        private void WriteBlock(byte[] buffer, int offset, int count, int blockIndex, int blockOffset)
        {
            if (false == BddInfo.Allocated[blockIndex])
            {
                //COW:allocate block

                _snapshotFileStream.Position = _snapshotFileStream.Length;
                BddInfo.Offset[blockIndex] = _snapshotFileStream.Position;

                //copy leading block from based image
                if (blockOffset > 0)
                {
                    var readCount = blockOffset;
                    var bufLeading = new byte[readCount];

                    if (null != _basedFileStream)
                    {
                        _basedFileStream.Position = (long)blockIndex * BddInfo.BlockSize;

                        var readOffset = 0;
                        while (readCount > 0)
                        {
                            var read = _basedFileStream.Read(bufLeading, readOffset, readCount);
                            readOffset += read;
                            readCount -= read;
                        }
                    }

                    _snapshotFileStream.Position = BddInfo.Offset[blockIndex];
                    _snapshotFileStream.Write(bufLeading, 0, bufLeading.Length);
                }

                //write to snapshot directly
                _snapshotFileStream.Position = BddInfo.Offset[blockIndex] + blockOffset;
                _snapshotFileStream.Write(buffer, offset, count);

                //copy suffix block from based image
                if (blockOffset + count < BddInfo.BlockSize)
                {
                    var suffixOffset = blockOffset + count;
                    var readCount = BddInfo.BlockSize - suffixOffset;
                    var bufSuffix = new byte[readCount];

                    if (null != _basedFileStream)
                    {
                        _basedFileStream.Position = (long)blockIndex * BddInfo.BlockSize + suffixOffset;
                        var readOffset = 0;
                        while (readCount > 0)
                        {
                            var read = _basedFileStream.Read(bufSuffix, readOffset, readCount);
                            readOffset += read;
                            readCount -= read;
                        }
                    }

                    _snapshotFileStream.Position = BddInfo.Offset[blockIndex] + suffixOffset;
                    _snapshotFileStream.Write(bufSuffix, 0, bufSuffix.Length);
                }

                BddInfo.Allocated[blockIndex] = true;
                _snapshotFileStream.Position = BddInfo.EntryTableOffset + BddInfo.BlockIndexEntrySize * blockIndex;
                var bufEntry = BitConverter.GetBytes(BddInfo.Entries[blockIndex]);
                _snapshotFileStream.Write(bufEntry, 0, bufEntry.Length);
            }
            else
            {
                _snapshotFileStream.Position = BddInfo.Offset[blockIndex] + blockOffset;
                _snapshotFileStream.Write(buffer, offset, count);
            }
        }

        public override byte[] ReadSectors(long sectorIndex, int sectorCount)
        {
            var rawOffset = sectorIndex * SectorSize;
            var rawSize = sectorCount * SectorSize;

            var buffer = new byte[rawSize];

            for (var i = 0; i < rawSize;)
            {
                var blockIndex = (int)(rawOffset / BddInfo.BlockSize);
                var blockOffset = (int)(rawOffset % BddInfo.BlockSize);
                var bytesToRead = BddInfo.BlockSize - blockOffset;

                if (i + bytesToRead > rawSize)
                {
                    bytesToRead = rawSize - i;
                }

                ReadBlock(buffer, i, bytesToRead, blockIndex, blockOffset);

                i += bytesToRead;
                rawOffset += bytesToRead;
            }

            return buffer;
        }

        public override void WriteSectors(long sectorIndex, byte[] data)
        {
            if (IsReadOnly) throw new AuthenticationException("Read only");

            var rawOffset = sectorIndex * SectorSize;

            for (var i = 0; i < data.Length;)
            {
                var blockIndex = (int)(rawOffset / BddInfo.BlockSize);
                var blockOffset = (int)(rawOffset % BddInfo.BlockSize);
                var bytesToWrite = BddInfo.BlockSize - blockOffset;

                if (i + bytesToWrite > data.Length)
                {
                    bytesToWrite = data.Length - i;
                }

                WriteBlock(data, i, bytesToWrite, blockIndex, blockOffset);

                i += bytesToWrite;
                rawOffset += bytesToWrite;
            }
        }

        public override bool ExclusiveLock()
        {
            if (_isExclusiveLock) return true;

            _basedFileStream?.Close();
            _snapshotFileStream?.Close();

            if (null != BddInfo.BasedImagePath) _basedFileStream = new FileStream(BddInfo.BasedImagePath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite);
            _snapshotFileStream = IsReadOnly
                ? File.Open(Path, FileMode.Open, FileAccess.Read, FileShare.Read)
                : File.Open(Path, FileMode.Open, FileAccess.ReadWrite, FileShare.Read);
            _isExclusiveLock = true;

            return true;
        }

        public void Clean()
        {
            if (IsReadOnly) return;
            for (var blockIndex = 0; blockIndex < BddInfo.NumberOfBlocks; blockIndex++)
            {
                BddInfo.Allocated[blockIndex] = false;
            }

            for (var blockIndex = 0; blockIndex < BddInfo.NumberOfBlocks; blockIndex++)
            {
                _snapshotFileStream.Position = BddInfo.EntryTableOffset + BddInfo.BlockIndexEntrySize * blockIndex;
                var bufEntry = BitConverter.GetBytes(BddInfo.Entries[blockIndex]);
                _snapshotFileStream.Write(bufEntry, 0, bufEntry.Length);
            }
            _snapshotFileStream.SetLength(BddInfo.EntryTableOffset + BddInfo.NumberOfBlocks * BddInfo.BlockIndexEntrySize);
        }

        public override bool ReleaseLock()
        {
            _basedFileStream?.Close();
            _snapshotFileStream?.Close();
            _isExclusiveLock = false;
            return true;
        }

        public static BlockDifferencingDiskImage Create(string snapshot, long size, int blockSize)
        {
            var numberOfBlocks = (int)Math.DivRem(size, blockSize, out var rem);
            if (rem != 0) throw new ArgumentException("Size no align to blockSize");

            using var stream = File.Create(snapshot);
            using var writer = new BinaryWriter(stream);

            writer.Write((ushort)0);
            writer.Write(blockSize);
            writer.Write(numberOfBlocks);
            writer.Flush();
            stream.SetLength(stream.Length + numberOfBlocks * BddInfo.BlockIndexEntrySize);
            writer.Close();
            stream.Close();

            //return instance
            return new BlockDifferencingDiskImage(snapshot);
        }

        public static BlockDifferencingDiskImage Create(string based, string snapshot, int blockSize)
        {
            var fileInfo = new FileInfo(based);

            if (fileInfo.Length % blockSize != 0) throw new ArgumentException("Block size no mul to base image size", nameof(blockSize));

            using var stream = File.Create(snapshot);
            using var writer = new BinaryWriter(stream);

            //create bdd struct

            var bufBasedPath = Encoding.UTF8.GetBytes(System.IO.Path.GetFullPath(based));
            var numberOfBlocks = (int)(fileInfo.Length / blockSize);
            writer.Write((ushort)bufBasedPath.Length);
            writer.Write(bufBasedPath);
            writer.Write(blockSize);
            writer.Write(numberOfBlocks);
            writer.Flush();
            stream.SetLength(stream.Length + numberOfBlocks * BddInfo.BlockIndexEntrySize);
            writer.Close();
            stream.Close();

            //return instance
            return new BlockDifferencingDiskImage(snapshot);
        }

        public override long Size => BddInfo.Length;

        public override int BytesPerSector => SectorSize;

        ////////////////////////////////////////////////////////////////////////////////////

        public BlockDifferencingDiskImage(string diskImagePath) : this(diskImagePath, false)
        {
        }

        public override bool ExclusiveLock(bool useOverlappedIO)
        {
            throw new System.NotImplementedException();
        }

        public override void Extend(long numberOfAdditionalBytes)
        {
            throw new System.NotImplementedException();
        }
    }
}