BeatSaber-IPA-Reloaded/IPA.Loader/AntiMalware/WinAPI/AmsiFileStream.cs

#nullable enable
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace IPA.AntiMalware.WinAPI
{
    internal class AmsiFileStream : IAmsiStream, IDisposable
    {
        private readonly FileInfo file;
        private readonly IntPtr session;

        public AmsiFileStream(FileInfo file, IntPtr session)
        {
            this.file = file;
            this.session = session;
        }

        public unsafe void GetAttribute([In] AmsiAttribute attribute, [In] uint dataSize, [Out] byte* buffer, out uint writtenData)
        {
            switch (attribute)
            {
                case AmsiAttribute.AppName:
                    writtenData = WriteWString(AmsiConstants.AppName, dataSize, buffer);
                    return;
                case AmsiAttribute.Session:
                    *(IntPtr*)buffer = session;
                    writtenData = (uint)sizeof(IntPtr);
                    return;

                case AmsiAttribute.ContentName:
                    writtenData = WriteWString(file.FullName, dataSize, buffer);
                    return;

                case AmsiAttribute.ContentSize:
                    *(ulong*)buffer = (ulong)file.Length;
                    writtenData = sizeof(ulong);
                    return;

                default:
                    throw new NotImplementedException(); // return e_notimpl
            }

            static unsafe uint WriteWString(string str, uint dataSize, byte* buffer)
            {
                fixed (char* name = str)
                {
                    return (uint)Encoding.Unicode.GetBytes(name, str.Length, buffer, (int)dataSize);
                }
            }
        }

        private FileStream? stream;
        private bool disposedValue;
        private readonly byte[] readBuffer = new byte[1024];

        public unsafe void Read([In] ulong position, [In] uint dataSize, [Out] byte* buffer, out uint readSize)
        {
            stream ??= file.OpenRead();

            stream.Position = (long)position;

            var bytesToRead = dataSize;
            readSize = 0;

            while (bytesToRead > 0)
            {
                var bytesRead = stream.Read(readBuffer, 0, (int)Math.Min(readBuffer.Length, bytesToRead));
                if (bytesRead == 0)
                {
                    break;
                }
                fixed (byte* readBufferPtr = readBuffer)
                {
                    Buffer.MemoryCopy(readBufferPtr, buffer + readSize, dataSize - readSize, bytesRead);
                }
                bytesToRead -= (uint)bytesRead;
                readSize += (uint)bytesRead;
            }
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    stream?.Dispose();
                }

                disposedValue = true;
            }
        }

        // This does not have unmanagd resources, so it doesn't need to exist
        // ~AmsiFileStream()
        // {
        //     // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
        //     Dispose(disposing: false);
        // }

        public void Dispose()
        {
            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
            Dispose(disposing: true);
            GC.SuppressFinalize(this);
        }
    }
}