<commit_before>﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using dot10.IO;
using dot10.PE;
using dot10.W32Resources;

namespace dot10.DotNet.Writer {
	/// <summary>
	/// Win32 resources
	/// </summary>
	public sealed class Win32ResourcesChunk : IChunk {
		readonly Win32Resources win32Resources;
		FileOffset offset;
		RVA rva;
		uint length;
		Dictionary<ResourceDirectory, uint> dirDict = new Dictionary<ResourceDirectory, uint>();
		List<ResourceDirectory> dirList = new List<ResourceDirectory>();
		Dictionary<ResourceData, uint> dataHeaderDict = new Dictionary<ResourceData, uint>();
		List<ResourceData> dataHeaderList = new List<ResourceData>();
		Dictionary<string, uint> stringsDict = new Dictionary<string, uint>(StringComparer.Ordinal);
		List<string> stringsList = new List<string>();
		Dictionary<IBinaryReader, uint> dataDict = new Dictionary<IBinaryReader, uint>();
		List<IBinaryReader> dataList = new List<IBinaryReader>();

		/// <inheritdoc/>
		public FileOffset FileOffset {
			get { return offset; }
		}

		/// <inheritdoc/>
		public RVA RVA {
			get { return rva; }
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="win32Resources">Win32 resources</param>
		public Win32ResourcesChunk(Win32Resources win32Resources) {
			this.win32Resources = win32Resources;
		}

		// None of these can be > ModuleWriter.DEFAULT_RESOURCE_ALIGNMENT
		const uint RESOURCE_DIR_ALIGNMENT = 4;
		const uint RESOURCE_DATA_HEADER_ALIGNMENT = 4;
		const uint RESOURCE_DATA_ALIGNMENT = 4;
		const uint RESOURCE_STRING_ALIGNMENT = 2;

		/// <inheritdoc/>
		public void SetOffset(FileOffset offset, RVA rva) {
			this.offset = offset;
			this.rva = rva;
			if (win32Resources == null)
				return;

			FindDirectoryEntries();

			// Place everything in the following order:
			//	1. All resource directories. The root is always first.
			//	2. All resource data headers.
			//	3. All the strings.
			//	4. All resource data.

			uint rsrcOffset = 0;

			uint maxAlignment = 1;
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DIR_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DATA_HEADER_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DATA_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_STRING_ALIGNMENT);
			if (((uint)offset & (maxAlignment - 1)) != 0)
				throw new ModuleWriterException(string.Format("Win32 resources section isn't {0}-byte aligned", maxAlignment));

			foreach (var dir in dirList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DIR_ALIGNMENT);
				dirDict[dir] = rsrcOffset;
				if (dir != dirList[0])
					AddString(dir.Name);
				rsrcOffset += 16 + (uint)(dir.Directories.Count + dir.Data.Count) * 8;
			}

			foreach (var data in dataHeaderList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DATA_HEADER_ALIGNMENT);
				dataHeaderDict[data] = rsrcOffset;
				AddString(data.Name);
				AddData(data.Data);
				rsrcOffset += 16;
			}

			foreach (var s in stringsList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_STRING_ALIGNMENT);
				stringsDict[s] = rsrcOffset;
				rsrcOffset += 2 + (uint)(s.Length * 2);
			}

			foreach (var data in dataList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DATA_ALIGNMENT);
				dataDict[data] = rsrcOffset;
				rsrcOffset += (uint)data.Length;
			}<commit_msg>Add methods to get RVA/FileOffset of Win32 resource headers/data/strings<commit_after>﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using dot10.IO;
using dot10.PE;
using dot10.W32Resources;

namespace dot10.DotNet.Writer {
	/// <summary>
	/// Win32 resources
	/// </summary>
	public sealed class Win32ResourcesChunk : IChunk {
		readonly Win32Resources win32Resources;
		FileOffset offset;
		RVA rva;
		uint length;
		Dictionary<ResourceDirectory, uint> dirDict = new Dictionary<ResourceDirectory, uint>();
		List<ResourceDirectory> dirList = new List<ResourceDirectory>();
		Dictionary<ResourceData, uint> dataHeaderDict = new Dictionary<ResourceData, uint>();
		List<ResourceData> dataHeaderList = new List<ResourceData>();
		Dictionary<string, uint> stringsDict = new Dictionary<string, uint>(StringComparer.Ordinal);
		List<string> stringsList = new List<string>();
		Dictionary<IBinaryReader, uint> dataDict = new Dictionary<IBinaryReader, uint>();
		List<IBinaryReader> dataList = new List<IBinaryReader>();

		/// <inheritdoc/>
		public FileOffset FileOffset {
			get { return offset; }
		}

		/// <inheritdoc/>
		public RVA RVA {
			get { return rva; }
		}

		/// <summary>
		/// Constructor
		/// </summary>
		/// <param name="win32Resources">Win32 resources</param>
		public Win32ResourcesChunk(Win32Resources win32Resources) {
			this.win32Resources = win32Resources;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> and <see cref="RVA"/> of a
		/// <see cref="ResourceDirectoryEntry"/>. <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dirEntry">A <see cref="ResourceDirectoryEntry"/></param>
		/// <param name="fileOffset">Updated with the file offset</param>
		/// <param name="rva">Updated with the RVA</param>
		/// <returns><c>true</c> if <paramref name="dirEntry"/> is valid and
		/// <paramref name="fileOffset"/> and <paramref name="rva"/> have been updated. <c>false</c>
		/// if <paramref name="dirEntry"/> is not part of the Win32 resources.</returns>
		public bool GetFileOffsetAndRvaOf(ResourceDirectoryEntry dirEntry, out FileOffset fileOffset, out RVA rva) {
			var dir = dirEntry as ResourceDirectory;
			if (dir != null)
				return GetFileOffsetAndRvaOf(dir, out fileOffset, out rva);

			var dataHeader = dirEntry as ResourceData;
			if (dataHeader != null)
				return GetFileOffsetAndRvaOf(dataHeader, out fileOffset, out rva);

			fileOffset = 0;
			rva = 0;
			return false;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> of a <see cref="ResourceDirectoryEntry"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dirEntry">A <see cref="ResourceDirectoryEntry"/></param>
		/// <returns>The file offset or 0 if <paramref name="dirEntry"/> is invalid</returns>
		public FileOffset GetFileOffset(ResourceDirectoryEntry dirEntry) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dirEntry, out fileOffset, out rva);
			return fileOffset;
		}

		/// <summary>
		/// Returns the <see cref="RVA"/> of a <see cref="ResourceDirectoryEntry"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dirEntry">A <see cref="ResourceDirectoryEntry"/></param>
		/// <returns>The RVA or 0 if <paramref name="dirEntry"/> is invalid</returns>
		public RVA GetRVA(ResourceDirectoryEntry dirEntry) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dirEntry, out fileOffset, out rva);
			return rva;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> and <see cref="RVA"/> of a
		/// <see cref="ResourceDirectory"/>. <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dir">A <see cref="ResourceDirectory"/></param>
		/// <param name="fileOffset">Updated with the file offset</param>
		/// <param name="rva">Updated with the RVA</param>
		/// <returns><c>true</c> if <paramref name="dir"/> is valid and
		/// <paramref name="fileOffset"/> and <paramref name="rva"/> have been updated. <c>false</c>
		/// if <paramref name="dir"/> is not part of the Win32 resources.</returns>
		public bool GetFileOffsetAndRvaOf(ResourceDirectory dir, out FileOffset fileOffset, out RVA rva) {
			uint offs;
			if (dir == null || !dirDict.TryGetValue(dir, out offs)) {
				fileOffset = 0;
				rva = 0;
				return false;
			}

			fileOffset = offset + offs;
			rva = this.rva + offs;
			return true;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> of a <see cref="ResourceDirectory"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dir">A <see cref="ResourceDirectory"/></param>
		/// <returns>The file offset or 0 if <paramref name="dir"/> is invalid</returns>
		public FileOffset GetFileOffset(ResourceDirectory dir) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dir, out fileOffset, out rva);
			return fileOffset;
		}

		/// <summary>
		/// Returns the <see cref="RVA"/> of a <see cref="ResourceDirectory"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dir">A <see cref="ResourceDirectory"/></param>
		/// <returns>The RVA or 0 if <paramref name="dir"/> is invalid</returns>
		public RVA GetRVA(ResourceDirectory dir) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dir, out fileOffset, out rva);
			return rva;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> and <see cref="RVA"/> of a
		/// <see cref="ResourceData"/>. <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dataHeader">A <see cref="ResourceData"/></param>
		/// <param name="fileOffset">Updated with the file offset</param>
		/// <param name="rva">Updated with the RVA</param>
		/// <returns><c>true</c> if <paramref name="dataHeader"/> is valid and
		/// <paramref name="fileOffset"/> and <paramref name="rva"/> have been updated. <c>false</c>
		/// if <paramref name="dataHeader"/> is not part of the Win32 resources.</returns>
		public bool GetFileOffsetAndRvaOf(ResourceData dataHeader, out FileOffset fileOffset, out RVA rva) {
			uint offs;
			if (dataHeader == null || !dataHeaderDict.TryGetValue(dataHeader, out offs)) {
				fileOffset = 0;
				rva = 0;
				return false;
			}

			fileOffset = offset + offs;
			rva = this.rva + offs;
			return true;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> of a <see cref="ResourceData"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dataHeader">A <see cref="ResourceData"/></param>
		/// <returns>The file offset or 0 if <paramref name="dataHeader"/> is invalid</returns>
		public FileOffset GetFileOffset(ResourceData dataHeader) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dataHeader, out fileOffset, out rva);
			return fileOffset;
		}

		/// <summary>
		/// Returns the <see cref="RVA"/> of a <see cref="ResourceData"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="dataHeader">A <see cref="ResourceData"/></param>
		/// <returns>The RVA or 0 if <paramref name="dataHeader"/> is invalid</returns>
		public RVA GetRVA(ResourceData dataHeader) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(dataHeader, out fileOffset, out rva);
			return rva;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> and <see cref="RVA"/> of the raw data
		/// owned by a <see cref="ResourceData"/>. <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="data">A <see cref="ResourceData"/>'s <see cref="IBinaryReader"/></param>
		/// <param name="fileOffset">Updated with the file offset</param>
		/// <param name="rva">Updated with the RVA</param>
		/// <returns><c>true</c> if <paramref name="data"/> is valid and
		/// <paramref name="fileOffset"/> and <paramref name="rva"/> have been updated. <c>false</c>
		/// if <paramref name="data"/> is not part of the Win32 resources.</returns>
		public bool GetFileOffsetAndRvaOf(IBinaryReader data, out FileOffset fileOffset, out RVA rva) {
			uint offs;
			if (data == null || !dataDict.TryGetValue(data, out offs)) {
				fileOffset = 0;
				rva = 0;
				return false;
			}

			fileOffset = offset + offs;
			rva = this.rva + offs;
			return true;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> of the raw data owned by a
		/// <see cref="ResourceData"/>. <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="data">A <see cref="ResourceData"/>'s <see cref="IBinaryReader"/></param>
		/// <returns>The file offset or 0 if <paramref name="data"/> is invalid</returns>
		public FileOffset GetFileOffset(IBinaryReader data) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(data, out fileOffset, out rva);
			return fileOffset;
		}

		/// <summary>
		/// Returns the <see cref="RVA"/> of the raw data owned by a <see cref="ResourceData"/>.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="data">A <see cref="ResourceData"/>'s <see cref="IBinaryReader"/></param>
		/// <returns>The RVA or 0 if <paramref name="data"/> is invalid</returns>
		public RVA GetRVA(IBinaryReader data) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(data, out fileOffset, out rva);
			return rva;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> and <see cref="RVA"/> of a
		/// <see cref="ResourceDirectoryEntry"/>'s name. <see cref="SetOffset"/> must have been
		/// called.
		/// </summary>
		/// <param name="name">The name of a <see cref="ResourceDirectoryEntry"/></param>
		/// <param name="fileOffset">Updated with the file offset</param>
		/// <param name="rva">Updated with the RVA</param>
		/// <returns><c>true</c> if <paramref name="name"/> is valid and
		/// <paramref name="fileOffset"/> and <paramref name="rva"/> have been updated. <c>false</c>
		/// if <paramref name="name"/> is not part of the Win32 resources.</returns>
		public bool GetFileOffsetAndRvaOf(string name, out FileOffset fileOffset, out RVA rva) {
			uint offs;
			if (name == null || !stringsDict.TryGetValue(name, out offs)) {
				fileOffset = 0;
				rva = 0;
				return false;
			}

			fileOffset = offset + offs;
			rva = this.rva + offs;
			return true;
		}

		/// <summary>
		/// Returns the <see cref="FileOffset"/> of a <see cref="ResourceDirectoryEntry"/>'s name.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="name">The name of a <see cref="ResourceDirectoryEntry"/></param>
		/// <returns>The file offset or 0 if <paramref name="name"/> is invalid</returns>
		public FileOffset GetFileOffset(string name) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(name, out fileOffset, out rva);
			return fileOffset;
		}

		/// <summary>
		/// Returns the <see cref="RVA"/> of a <see cref="ResourceDirectoryEntry"/>'s name.
		/// <see cref="SetOffset"/> must have been called.
		/// </summary>
		/// <param name="name">The name of a <see cref="ResourceDirectoryEntry"/></param>
		/// <returns>The RVA or 0 if <paramref name="name"/> is invalid</returns>
		public RVA GetRVA(string name) {
			FileOffset fileOffset;
			RVA rva;
			GetFileOffsetAndRvaOf(name, out fileOffset, out rva);
			return rva;
		}

		// None of these can be > ModuleWriter.DEFAULT_RESOURCE_ALIGNMENT
		const uint RESOURCE_DIR_ALIGNMENT = 4;
		const uint RESOURCE_DATA_HEADER_ALIGNMENT = 4;
		const uint RESOURCE_STRING_ALIGNMENT = 2;
		const uint RESOURCE_DATA_ALIGNMENT = 4;

		/// <inheritdoc/>
		public void SetOffset(FileOffset offset, RVA rva) {
			this.offset = offset;
			this.rva = rva;
			if (win32Resources == null)
				return;

			FindDirectoryEntries();

			// Place everything in the following order:
			//	1. All resource directories. The root is always first.
			//	2. All resource data headers.
			//	3. All the strings.
			//	4. All resource data.

			uint rsrcOffset = 0;

			uint maxAlignment = 1;
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DIR_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DATA_HEADER_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_STRING_ALIGNMENT);
			maxAlignment = Math.Max(maxAlignment, RESOURCE_DATA_ALIGNMENT);
			if (((uint)offset & (maxAlignment - 1)) != 0)
				throw new ModuleWriterException(string.Format("Win32 resources section isn't {0}-byte aligned", maxAlignment));

			foreach (var dir in dirList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DIR_ALIGNMENT);
				dirDict[dir] = rsrcOffset;
				if (dir != dirList[0])
					AddString(dir.Name);
				rsrcOffset += 16 + (uint)(dir.Directories.Count + dir.Data.Count) * 8;
			}

			foreach (var data in dataHeaderList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DATA_HEADER_ALIGNMENT);
				dataHeaderDict[data] = rsrcOffset;
				AddString(data.Name);
				AddData(data.Data);
				rsrcOffset += 16;
			}

			foreach (var s in stringsList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_STRING_ALIGNMENT);
				stringsDict[s] = rsrcOffset;
				rsrcOffset += 2 + (uint)(s.Length * 2);
			}

			foreach (var data in dataList) {
				rsrcOffset = Utils.AlignUp(rsrcOffset, RESOURCE_DATA_ALIGNMENT);
				dataDict[data] = rsrcOffset;
				rsrcOffset += (uint)data.Length;
			}
{'fraction_non_alphanumeric': 0.1452463733365304, 'fraction_numerical': 0.008751948207648963, 'mean_word_length': 3.64837001950404, 'pattern_counts': {'":': 0, '<': 281, '<?xml version=': 0, '>': 283, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '238451', 'n_tokens_mistral': 6007, 'n_tokens_neox': 5257, 'n_words': 1470}