Skip to content

Commit

Permalink
Added unit test for RegistryEx (SubnauticaNitrox#1617)
Browse files Browse the repository at this point in the history
  • Loading branch information
Measurity authored Oct 6, 2021
1 parent fbd4229 commit 1c5c0d3
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 47 deletions.
8 changes: 4 additions & 4 deletions Nitrox.Bootloader/Main.cs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ namespace Nitrox.Bootloader
{
public static class Main
{
private static readonly Lazy<string> nitroxLauncherDir = new Lazy<string>(() =>
private static readonly Lazy<string> nitroxLauncherDir = new(() =>
{
// Get path from command args.
string[] args = Environment.GetCommandLineArgs();
Expand All @@ -18,7 +18,7 @@ public static class Main
return Path.GetFullPath(args[i + 1]);
}
}

// Get path from environment variable.
string envPath = Environment.GetEnvironmentVariable("NITROX_LAUNCHER_PATH");
if (Directory.Exists(envPath))
Expand Down Expand Up @@ -70,7 +70,7 @@ public static void Execute()
Console.WriteLine(error);
return;
}

Environment.SetEnvironmentVariable("NITROX_LAUNCHER_PATH", nitroxLauncherDir.Value);

AppDomain.CurrentDomain.AssemblyResolve += CurrentDomainOnAssemblyResolve;
Expand Down Expand Up @@ -118,4 +118,4 @@ private static Assembly CurrentDomainOnAssemblyResolve(object sender, ResolveEve
return Assembly.LoadFile(dllPath);
}
}
}
}
147 changes: 104 additions & 43 deletions NitroxModel/Platforms/OS/Windows/Internal/RegistryEx.cs
Original file line number Diff line number Diff line change
Expand Up @@ -10,59 +10,31 @@ namespace NitroxModel.Platforms.OS.Windows.Internal
{
public static class RegistryEx
{
private static (RegistryKey baseKey, string valueKey) GetKey(string path, bool needsWriteAccess = true)
{
if (string.IsNullOrWhiteSpace(path))
{
return (null, null);
}
path = path.Trim();

// Parse path to get the registry key instance and the name of the .
string[] parts = path.Split(Path.DirectorySeparatorChar);
RegistryHive hive = RegistryHive.CurrentUser;
string regPathWithoutHiveOrKey;
if (path.IndexOf("Computer", StringComparison.OrdinalIgnoreCase) < 0)
{
regPathWithoutHiveOrKey = string.Join(Path.DirectorySeparatorChar.ToString(), parts.TakeUntilLast());
}
else
{
regPathWithoutHiveOrKey = string.Join(Path.DirectorySeparatorChar.ToString(), parts.Skip(2).TakeUntilLast());
hive = parts[1].ToLowerInvariant() switch
{
"hkey_classes_root" => RegistryHive.ClassesRoot,
"hkey_local_machine" => RegistryHive.LocalMachine,
"hkey_current_user" => RegistryHive.CurrentUser,
"hkey_users" => RegistryHive.Users,
"hkey_current_config" => RegistryHive.CurrentConfig,
_ => throw new ArgumentException($"Path must contain a valid registry hive but was given '{parts[1]}'", nameof(path))
};
}

return (RegistryKey.OpenBaseKey(hive, RegistryView.Registry64).OpenSubKey(regPathWithoutHiveOrKey, needsWriteAccess), parts[parts.Length - 1]);
}

/// <summary>
/// Reads the value of the registry key or returns the default value of <see cref="T" />.
/// </summary>
/// <param name="pathWithKey">
/// <param name="pathWithValue">
/// Full path to the registry key. If the registry hive is omitted then "current user" is used.
/// </param>
/// <param name="defaultValue">The default value if the registry key is not found or failed to convert to <see cref="T" />.</param>
/// <typeparam name="T">Type of value to read. If the value in the registry key does not match it will try to convert.</typeparam>
/// <returns>Value as read from registry or null if not found.</returns>
public static T Read<T>(string pathWithKey, T defaultValue = default)
public static T Read<T>(string pathWithValue, T defaultValue = default)
{
(RegistryKey baseKey, string valueKey) = GetKey(pathWithKey, false);
(RegistryKey baseKey, string valueKey) = GetKey(pathWithValue, false);
if (baseKey == null)
{
return defaultValue;
}

try
{
return (T)TypeDescriptor.GetConverter(typeof(T)).ConvertFrom(baseKey.GetValue(valueKey));
object value = baseKey.GetValue(valueKey);
if (value == null)
{
return defaultValue;
}
return (T)TypeDescriptor.GetConverter(typeof(T)).ConvertFrom(value);
}
catch (Exception)
{
Expand All @@ -74,14 +46,51 @@ public static T Read<T>(string pathWithKey, T defaultValue = default)
}
}

/// <summary>
/// Deletes the whole subtree or value, whichever exists.
/// </summary>
/// <param name="pathWithOptionalValue">If no value name is given it will delete the key instead.</param>
/// <returns>True if something was deleted.</returns>
public static bool Delete(string pathWithOptionalValue)
{
(RegistryKey key, string valueKey) = GetKey(pathWithOptionalValue);
if (key == null)
{
return false;
}

// Try delete the key.
RegistryKey prev = key;
key = key.OpenSubKey(valueKey);
if (key != null)
{
key.DeleteSubKeyTree(valueKey);
key.Dispose();
prev.Dispose();
return true;
}
key = prev; // Restore state for next step

// Not a key, delete the value if it exists.
if (key.GetValue(valueKey) != null)
{
key.DeleteValue(valueKey);
key.Dispose();
return true;
}

// Nothing to delete.
return false;
}

public static void Write<T>(string pathWithKey, T value)
{
if (value == null)
{
throw new ArgumentNullException(nameof(value));
}
(RegistryKey baseKey, string valueKey) pair = GetKey(pathWithKey, false);
if (pair.baseKey == null)
(RegistryKey baseKey, string valueKey) = GetKey(pathWithKey, true, true);
if (baseKey == null)
{
return;
}
Expand All @@ -99,7 +108,7 @@ public static void Write<T>(string pathWithKey, T value)
{
try
{
kind = pair.baseKey.GetValueKind(pair.valueKey);
kind = baseKey.GetValueKind(valueKey);
}
catch (Exception)
{
Expand All @@ -109,11 +118,11 @@ public static void Write<T>(string pathWithKey, T value)

try
{
pair.baseKey.SetValue(pair.valueKey, value, kind.GetValueOrDefault(RegistryValueKind.String));
baseKey.SetValue(valueKey, value, kind.GetValueOrDefault(RegistryValueKind.String));
}
finally
{
pair.baseKey.Dispose();
baseKey.Dispose();
}
}

Expand Down Expand Up @@ -171,5 +180,57 @@ public static Task CompareAsync<T>(string pathWithKey, Func<T, bool> predicate,
CancellationTokenSource source = new(timeout == default ? TimeSpan.FromSeconds(10) : timeout);
return CompareAsync(pathWithKey, predicate, source.Token);
}

private static (RegistryKey baseKey, string valueKey) GetKey(string path, bool needsWriteAccess = true, bool createIfNotExists = false)
{
if (string.IsNullOrWhiteSpace(path))
{
return (null, null);
}
path = path.Trim();

// Parse path to get the registry key instance and the name of the .
string[] parts = path.Split(Path.DirectorySeparatorChar);
string[] partsWithoutHive;
RegistryHive hive = RegistryHive.CurrentUser;
string regPathWithoutHiveOrKey;
if (path.IndexOf("Computer", StringComparison.OrdinalIgnoreCase) < 0)
{
partsWithoutHive = parts.TakeUntilLast().ToArray();
regPathWithoutHiveOrKey = string.Join(Path.DirectorySeparatorChar.ToString(), partsWithoutHive);
}
else
{
partsWithoutHive = parts.Skip(2).TakeUntilLast().ToArray();
regPathWithoutHiveOrKey = string.Join(Path.DirectorySeparatorChar.ToString(), partsWithoutHive);
hive = parts[1].ToLowerInvariant() switch
{
"hkey_classes_root" => RegistryHive.ClassesRoot,
"hkey_local_machine" => RegistryHive.LocalMachine,
"hkey_current_user" => RegistryHive.CurrentUser,
"hkey_users" => RegistryHive.Users,
"hkey_current_config" => RegistryHive.CurrentConfig,
_ => throw new ArgumentException($"Path must contain a valid registry hive but was given '{parts[1]}'", nameof(path))
};
}

RegistryKey hiveRef = RegistryKey.OpenBaseKey(hive, RegistryView.Registry64);
RegistryKey key = hiveRef.OpenSubKey(regPathWithoutHiveOrKey, needsWriteAccess);
// Should the key (and its path leading to it) be created?
if (key == null && createIfNotExists)
{
key = hiveRef;
foreach (string part in partsWithoutHive)
{
RegistryKey prev = key;
key = key?.OpenSubKey(part, needsWriteAccess) ?? key?.CreateSubKey(part, needsWriteAccess);

// Cleanup old/parent key reference
prev?.Dispose();
}
}

return (key, parts[parts.Length - 1]);
}
}
}
}
1 change: 1 addition & 0 deletions NitroxTest/NitroxTest.csproj
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
<Compile Include="Patcher\Patches\ConstructorInput_OnCraftingBegin_PatchTest.cs" />
<Compile Include="Patcher\Patches\Equipment_RemoveItem_PatchTest.cs" />
<Compile Include="Patcher\Test\PatchTestHelper.cs" />
<Compile Include="Platforms\OS\Windows\RegistryTest.cs" />
<Compile Include="Properties\AssemblyInfo.cs" />
<Compile Include="Serialization\InteractiveChildObjectIdentifiersSerializationTest.cs" />
<Compile Include="Serialization\PdaStateDataSerializationTest.cs" />
Expand Down
41 changes: 41 additions & 0 deletions NitroxTest/Platforms/OS/Windows/RegistryTest.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
using System;
using System.Threading.Tasks;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Microsoft.Win32;
using NitroxModel.Platforms.OS.Windows.Internal;

namespace NitroxTest.Platforms.OS.Windows
{
[TestClass]
public class RegistryTest
{
[TestMethod]
public async Task WaitsForRegistryKeyToExist()
{
const string pathToKey = @"SOFTWARE\Nitrox\test";

RegistryEx.Write(pathToKey, 0);
var readTask = Task.Run(async () =>
{
try
{
await RegistryEx.CompareAsync<int>(pathToKey,
v => v == 1337,
TimeSpan.FromSeconds(5));
return true;
}
catch (TaskCanceledException)
{
return false;
}
});

RegistryEx.Write(pathToKey, 1337);
Assert.IsTrue(await readTask);

// Cleanup (we can keep "Nitrox" key intact).
RegistryEx.Delete(pathToKey);
Assert.IsNull(RegistryEx.Read<string>(pathToKey));
}
}
}

0 comments on commit 1c5c0d3

Please sign in to comment.