From 9b028d1bff7e0c0b2d77ea8a4704913968749177 Mon Sep 17 00:00:00 2001
From: "Joel Sallow (/u/ta11ow)" <32407840+vexx32@users.noreply.github.com>
Date: Sun, 24 Nov 2019 10:14:34 -0500
Subject: [PATCH] :bug: Fix issue with Katas' verification script (#328)

CheckRestrictedLanguage() is far too stringent.
It prevents usage of basic array indexing and member access.

Switch to detecting command names in the AST.
This also lets users define sub-functions if they want.

Only functions defined in the test function itself will be allowed.
---
 .../Koans/Katas/ProcessingStrings.Koans.ps1   | 25 ++++++++++++++++---
 .../Koans/Katas/SortingCharacters.Koans.ps1   | 23 ++++++++++++++---
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/PSKoans/Koans/Katas/ProcessingStrings.Koans.ps1 b/PSKoans/Koans/Katas/ProcessingStrings.Koans.ps1
index 545eca3bd..02744b3ac 100644
--- a/PSKoans/Koans/Katas/ProcessingStrings.Koans.ps1
+++ b/PSKoans/Koans/Katas/ProcessingStrings.Koans.ps1
@@ -1,4 +1,6 @@
 using module PSKoans
+using namespace System.Management.Automation.Language
+using namespace System.Collections.Generic
 [Koan(Position = 150)]
 param()
 <#
@@ -51,10 +53,25 @@ Describe "The Stock Challenge" {
         #>
 
         $Verification = {
-            [string[]]$AllowedCommands = @()
-            [string[]]$AllowedVariables = @("*")
-            $Function = Get-Item -Path 'Function:Get-GreatestVarianceDate'
-            $Function.ScriptBlock.CheckRestrictedLanguage($AllowedCommands, $AllowedVariables, $false)
+            $Functions = [Hashset[string]]::new([StringComparer]::OrdinalIgnoreCase)
+            $Ast = (Get-Command 'Get-GreatestVarianceDate' -CommandType Function).ScriptBlock.Ast
+            $Ast.FindAll(
+                {
+                    param($node)
+                    if ($node -is [CommandAst] -and ($name = $node.GetCommandName()) -and !$Functions.Contains($name)) {
+                        throw 'Usage of external cmdlets and functions is not permitted.'
+                    }
+
+                    if ($node -is [FunctionDefinitionAst]) {
+                        $Functions.Add($node.Name) > $null
+                        return
+                    }
+
+                    if ($node.Left -is [VariableExpressionAst] -and $node.Left.VariablePath.DriveName -eq 'Function') {
+                        $Functions.Add($node.Left.VariablePath.UserPath -replace '^function:') > $null
+                    }
+                }, $true
+            )
         }
 
         function Get-GreatestVarianceDate {
diff --git a/PSKoans/Koans/Katas/SortingCharacters.Koans.ps1 b/PSKoans/Koans/Katas/SortingCharacters.Koans.ps1
index 46aa55dba..7adcdead9 100644
--- a/PSKoans/Koans/Katas/SortingCharacters.Koans.ps1
+++ b/PSKoans/Koans/Katas/SortingCharacters.Koans.ps1
@@ -16,10 +16,25 @@ param()
 Describe 'Kata - Sorting Characters' {
     BeforeAll {
         $Verification = {
-            [string[]]$AllowedCommands = @()
-            [string[]]$AllowedVariables = @("*")
-            $Function = Get-Item -Path 'Function:Get-SortedString'
-            $Function.ScriptBlock.CheckRestrictedLanguage($AllowedCommands, $AllowedVariables, $false)
+            $Functions = [Hashset[string]]::new([StringComparer]::OrdinalIgnoreCase)
+            $Ast = (Get-Command 'Get-GreatestVarianceDate' -CommandType Function).ScriptBlock.Ast
+            $Ast.FindAll(
+                {
+                    param($node)
+                    if ($node -is [CommandAst] -and ($name = $node.GetCommandName()) -and !$Functions.Contains($name)) {
+                        throw 'Usage of external cmdlets and functions is not permitted.'
+                    }
+
+                    if ($node -is [FunctionDefinitionAst]) {
+                        $Functions.Add($node.Name) > $null
+                        return
+                    }
+
+                    if ($node.Left -is [VariableExpressionAst] -and $node.Left.VariablePath.DriveName -eq 'Function') {
+                        $Functions.Add($node.Left.VariablePath.UserPath -replace '^function:') > $null
+                    }
+                }, $true
+            )
         }
 
         function Get-SortedString {