Added support for user atoms.

Author: tyranid

NtApiDotNet/NtAtom.cs

@@ -26,17 +26,86 @@ public sealed class NtAtom
     {
         #region Constructors
 
-        internal NtAtom(ushort atom)
+        internal NtAtom(ushort atom, bool global)
         {
             Atom = atom;
+            Global = global;
+        }
+
+        #endregion
+
+        #region Private Members
+        private static IEnumerable<NtAtom> GetGlobalAtoms()
+        {
+            int size = 1024;
+            while (size < 5 * 1024 * 1024)
+            {
+                using (var buffer = new SafeStructureInOutBuffer<AtomTableInformation>(size, true))
+                {
+                    NtStatus status = NtSystemCalls.NtQueryInformationAtom(0,
+                        AtomInformationClass.AtomTableInformation, buffer, buffer.Length, out int return_length);
+                    if (status.IsSuccess())
+                    {
+                        AtomTableInformation table = buffer.Result;
+                        IntPtr data = buffer.Data.DangerousGetHandle();
+                        ushort[] atoms = new ushort[table.NumberOfAtoms];
+                        buffer.Data.ReadArray(0, atoms, 0, atoms.Length);
+                        return atoms.Select(a => new NtAtom(a, true)).ToArray();
+                    }
+                    else if (status != NtStatus.STATUS_INFO_LENGTH_MISMATCH)
+                    {
+                        throw new NtException(status);
+                    }
+                    size *= 2;
+                }
+            }
+            return new NtAtom[0];
+        }
+
+        private static IEnumerable<NtAtom> GetUserAtoms()
+        {
+            List<NtAtom> atoms = new List<NtAtom>();
+            for (int atom = 0xC000; atom < 0x10000; ++atom)
+            {
+                var next_atom = new NtAtom((ushort)atom, false);
+                if (next_atom.GetName(false).IsSuccess)
+                {
+                    atoms.Add(next_atom);
+                }
+            }
+            return atoms.AsReadOnly();
+        }
+
+        private NtResult<string> GetGlobalName(bool throw_on_error)
+        {
+            using (var buffer = new SafeStructureInOutBuffer<AtomBasicInformation>(2048, false))
+            {
+                return NtSystemCalls.NtQueryInformationAtom(Atom, AtomInformationClass.AtomBasicInformation,
+                     buffer, buffer.Length, out int return_length)
+                     .CreateResult(throw_on_error, () => buffer.Data.ReadUnicodeString(buffer.Result.NameLength / 2));
+            }
+        }
+
+        private NtResult<string> GetUserName(bool throw_on_error)
+        {
+            using (UnicodeStringAllocated str = new UnicodeStringAllocated(2048))
+            {
+                int length = NtSystemCalls.NtUserGetAtomName(Atom, str);
+                if (length == 0)
+                {
+                    return NtStatus.STATUS_OBJECT_NAME_NOT_FOUND.CreateResultFromError<string>(throw_on_error);
+                }
+                str.String.Length = (ushort)(length * 2);
+                return str.ToString().CreateResult();
+            }
         }
 
         #endregion
 
         #region Static Methods
 
         /// <summary>
-        /// Add an atom name
+        /// Add a global atom name
         /// </summary>
         /// <param name="name">The name to add</param>
         /// <param name="flags">Flags for the add.</param>
@@ -47,17 +116,17 @@ public static NtResult<NtAtom> Add(string name, AddAtomFlags flags, bool throw_o
             if (flags == AddAtomFlags.None)
             {
                 return NtSystemCalls.NtAddAtom(name + "\0", (name.Length + 1) * 2,
-                    out ushort atom).CreateResult(throw_on_error, () => new NtAtom(atom));
+                    out ushort atom).CreateResult(throw_on_error, () => new NtAtom(atom, true));
             }
             else
             {
                 return NtSystemCalls.NtAddAtomEx(name + "\0", (name.Length + 1) * 2,
-                    out ushort atom, flags).CreateResult(throw_on_error, () => new NtAtom(atom));
+                    out ushort atom, flags).CreateResult(throw_on_error, () => new NtAtom(atom, true));
             }
         }
 
         /// <summary>
-        /// Add an atom name
+        /// Add a global atom name
         /// </summary>
         /// <param name="name">The name to add</param>
         /// <param name="flags">Flags for the add.</param>
@@ -68,7 +137,7 @@ public static NtAtom Add(string name, AddAtomFlags flags)
         }
 
         /// <summary>
-        /// Add an atom name
+        /// Add a global atom name
         /// </summary>
         /// <param name="name">The name to add</param>
         /// <param name="throw_on_error">True to throw on error.</param>
@@ -79,7 +148,7 @@ public static NtResult<NtAtom> Add(string name, bool throw_on_error)
         }
 
         /// <summary>
-        /// Add an atom name
+        /// Add a global atom name
         /// </summary>
         /// <param name="name">The name to add</param>
         /// <returns>A reference to the atom</returns>
@@ -89,19 +158,19 @@ public static NtAtom Add(string name)
         }
 
         /// <summary>
-        /// Find an atom by name.
+        /// Find a global atom by name.
         /// </summary>
         /// <param name="name">The name of the atom.</param>
         /// <param name="throw_on_error">True to throw on error.</param>
         /// <returns>The found atom.</returns>
         public static NtResult<NtAtom> Find(string name, bool throw_on_error)
         {
             return NtSystemCalls.NtFindAtom(name + "\0", (name.Length + 1) * 2, 
-                out ushort atom).CreateResult(throw_on_error, () => new NtAtom(atom));
+                out ushort atom).CreateResult(throw_on_error, () => new NtAtom(atom, true));
         }
 
         /// <summary>
-        /// Find an atom by name.
+        /// Find a global atom by name.
         /// </summary>
         /// <param name="name">The name of the atom.</param>
         /// <returns>The found atom.</returns>
@@ -111,25 +180,37 @@ public static NtAtom Find(string name)
         }
 
         /// <summary>
-        /// Query if the atom exists.
+        /// Query if a global atom exists.
         /// </summary>
         /// <param name="atom">The atom to check.</param>
         /// <returns>True if the atom exists.</returns>
         public static bool Exists(ushort atom)
         {
-            return new NtAtom(atom).GetName(false).IsSuccess;
+            return Exists(atom, true);
         }
 
         /// <summary>
-        /// Open an atom by number.
+        /// Query if the atom exists.
+        /// </summary>
+        /// <param name="atom">The atom to check.</param>
+        /// <param name="global">Specify true to check for a global atom, otherwise gets a user atom.</param>
+        /// <returns>True if the atom exists.</returns>
+        public static bool Exists(ushort atom, bool global)
+        {
+            return new NtAtom(atom, global).GetName(false).IsSuccess;
+        }
+
+        /// <summary>
+        /// Open a global atom by number.
         /// </summary>
         /// <param name="atom">The atom to open.</param>
         /// <param name="check_exists">True to check atom exists.</param>
+        /// <param name="global">True to open a global atom, otherwise a user atom.</param>
         /// <param name="throw_on_error">True to throw on error.</param>
         /// <returns>The atom object.</returns>
-        public static NtResult<NtAtom> Open(ushort atom, bool check_exists, bool throw_on_error)
+        public static NtResult<NtAtom> Open(ushort atom, bool check_exists, bool global, bool throw_on_error)
         {
-            NtAtom ret = new NtAtom(atom);
+            NtAtom ret = new NtAtom(atom, global);
             if (check_exists)
             {
                 return ret.GetName(false).Status.CreateResult(throw_on_error, () => ret);
@@ -138,7 +219,19 @@ public static NtResult<NtAtom> Open(ushort atom, bool check_exists, bool throw_o
         }
 
         /// <summary>
-        /// Open an atom by number.
+        /// Open a global atom by number.
+        /// </summary>
+        /// <param name="atom">The atom to open.</param>
+        /// <param name="check_exists">True to check atom exists.</param>
+        /// <param name="throw_on_error">True to throw on error.</param>
+        /// <returns>The atom object.</returns>
+        public static NtResult<NtAtom> Open(ushort atom, bool check_exists, bool throw_on_error)
+        {
+            return Open(atom, check_exists, true, throw_on_error);
+        }
+
+        /// <summary>
+        /// Open a global atom by number.
         /// </summary>
         /// <param name="atom">The atom to open.</param>
         /// <param name="check_exists">True to check atom exists.</param>
@@ -149,7 +242,7 @@ public static NtAtom Open(ushort atom, bool check_exists)
         }
 
         /// <summary>
-        /// Open an atom by number.
+        /// Open a global atom by number.
         /// </summary>
         /// <param name="atom">The atom to open.</param>
         /// <returns>The atom object.</returns>
@@ -162,37 +255,24 @@ public static NtAtom Open(ushort atom)
         /// Enumerate all atoms.
         /// </summary>
         /// <returns>An enumeration of all atoms on the system.</returns>
+        public static IEnumerable<NtAtom> GetAtoms(bool global)
+        {
+            return global ? GetGlobalAtoms() : GetUserAtoms();
+        }
+
+        /// <summary>
+        /// Enumerate all global atoms.
+        /// </summary>
+        /// <returns>An enumeration of all atoms on the system.</returns>
         public static IEnumerable<NtAtom> GetAtoms()
         {
-            int size = 1024;
-            while (size < 5 * 1024 * 1024)
-            {
-                using (var buffer = new SafeStructureInOutBuffer<AtomTableInformation>(size, true))
-                {
-                    NtStatus status = NtSystemCalls.NtQueryInformationAtom(0,
-                        AtomInformationClass.AtomTableInformation, buffer, buffer.Length, out int return_length);
-                    if (status.IsSuccess())
-                    {
-                        AtomTableInformation table = buffer.Result;
-                        IntPtr data = buffer.Data.DangerousGetHandle();
-                        ushort[] atoms = new ushort[table.NumberOfAtoms];
-                        buffer.Data.ReadArray(0, atoms, 0, atoms.Length);
-                        return atoms.Select(a => new NtAtom(a));
-                    }
-                    else if (status != NtStatus.STATUS_INFO_LENGTH_MISMATCH)
-                    {
-                        throw new NtException(status);
-                    }
-                    size *= 2;
-                }
-            }
-            return new NtAtom[0];
+            return GetAtoms(true);
         }
         #endregion
 
         #region Public Methods
         /// <summary>
-        /// Delete an atom.
+        /// Delete a global atom.
         /// </summary>
         /// <param name="throw_on_error">True to throw on error.</param>
         /// <returns>The NT status code.</returns>
@@ -202,7 +282,7 @@ public NtStatus Delete(bool throw_on_error)
         }
 
         /// <summary>
-        /// Delete an atom.
+        /// Delete a global atom.
         /// </summary>
         public void Delete()
         {
@@ -216,12 +296,7 @@ public void Delete()
         /// <returns>The name of the atom.</returns>
         public NtResult<string> GetName(bool throw_on_error)
         {
-            using (var buffer = new SafeStructureInOutBuffer<AtomBasicInformation>(2048, false))
-            {
-                return NtSystemCalls.NtQueryInformationAtom(Atom, AtomInformationClass.AtomBasicInformation,
-                     buffer, buffer.Length, out int return_length)
-                     .CreateResult(throw_on_error, () => buffer.Data.ReadUnicodeString(buffer.Result.NameLength / 2));
-            }
+            return Global ? GetGlobalName(throw_on_error) : GetUserName(throw_on_error);
         }
 
         #endregion
@@ -239,6 +314,11 @@ public NtResult<string> GetName(bool throw_on_error)
         /// <returns>The name of the atom</returns>
         public string Name => GetName(true).Result;
 
+        /// <summary>
+        /// If true indicates this is a global atom, otherwise it's a user atom.
+        /// </summary>
+        public bool Global { get; }
+
         #endregion
     }
 }

NtApiDotNet/NtAtomNative.cs

@@ -69,6 +69,9 @@ public static extern NtStatus NtQueryInformationAtom(
             int AtomInformationLength,
             out int ReturnLength
         );
+
+        [DllImport("win32u.dll", SetLastError = true)]
+        public static extern int NtUserGetAtomName(ushort Atom,  UnicodeStringAllocated Name);
     }
 #pragma warning restore 1591
 }

NtObjectManager/NtObjectManager.psm1

@@ -7185,13 +7185,15 @@ function Get-NtTypeAccess {
 
 <#
 .SYNOPSIS
-Get an ATOM objects.
+Get an ATOM object.
 .DESCRIPTION
 This cmdlet gets all ATOM objects or by name or atom.
 .PARAMETER Atom
 Specify the ATOM to get.
 .PARAMETER Name
 Specify the name of the ATOM to get.
+.PARAMETER User
+Specify to get a user atom rather than a global.
 .INPUTS
 None
 .OUTPUTS
@@ -7203,12 +7205,15 @@ function Get-NtAtom {
         [Parameter(Mandatory, ParameterSetName = "FromAtom")]
         [uint16]$Atom,
         [Parameter(Mandatory, Position = 0, ParameterSetName = "FromName")]
-        [string]$Name
+        [string]$Name,
+        [Parameter(ParameterSetName = "All")]
+        [Parameter(ParameterSetName = "FromAtom")]
+        [switch]$User
     )
 
     switch ($PSCmdlet.ParameterSetName) {
-        "All" { [NtApiDotNet.NtAtom]::GetAtoms() | Write-Output }
-        "FromAtom" { [NtApiDotNet.NtAtom]::Open($Atom) | Write-Output }
+        "All" { [NtApiDotNet.NtAtom]::GetAtoms(!$User) | Write-Output }
+        "FromAtom" { [NtApiDotNet.NtAtom]::Open($Atom, $true, !$User, $true).Result | Write-Output }
         "FromName" { [NtApiDotNet.NtAtom]::Find($Name) | Write-Output }
     }
 }