BeatSaber-IPA-Reloaded/IPA.Loader/AntiMalware/WinAPI/AmsiMemoryStream.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 AmsiMemoryStream : IAmsiStream, IDisposable
    {
        private readonly string contentName;
        private readonly byte[] data;
        private readonly GCHandle dataHandle;

        private readonly IntPtr session;
        private bool disposedValue;

        public AmsiMemoryStream(string contentName, byte[] data, IntPtr session)
        {
            this.data = data;
            dataHandle = GCHandle.Alloc(data, GCHandleType.Pinned);
            this.session = session;
            this.contentName = contentName;
        }

        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(contentName, dataSize, buffer);
                    return;

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

                case AmsiAttribute.ContentAddress:
                    // because our data is pinned, it can't move while this object exists so we can pass out the fixed address
                    fixed (byte* dataAddr = data)
                    {
                        *(byte**)buffer = dataAddr;
                    }
                    writtenData = (uint)sizeof(IntPtr);
                    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);
                }
            }
        }

        public unsafe void Read([In] ulong position, [In] uint dataSize, [Out] byte* buffer, out uint readSize)
        {
            if (position >= (ulong)data.Length)
            {
                throw new EndOfStreamException();
            } 

            fixed (byte* dataPtr = data)
            {
                var toRead = Math.Min((ulong)data.Length - position, dataSize);
                Buffer.MemoryCopy(dataPtr + position, buffer, dataSize, toRead);
                readSize = (uint)toRead;
            }
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                    // no managed stae to dispose
                }

                dataHandle.Free();
                disposedValue = true;
            }
        }

        ~AmsiMemoryStream()
        {
            // 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);
        }
    }
}