From 2a08834b9d49f8e7f2d6ab7e1fc3d7fcbab5f434 Mon Sep 17 00:00:00 2001 From: KnicKnic Date: Tue, 9 Jul 2019 08:35:18 -0700 Subject: [PATCH] add support for calling back into powershell from a callback --- .gitignore | 2 +- examples/cmd/main.go | 2 +- native-powershell | 2 +- pkg/powershell/callbackexample_test.go | 16 +++++---------- pkg/powershell/chelpers.go | 2 +- pkg/powershell/context.go | 6 +++--- pkg/powershell/higherops_test.go | 6 +++--- pkg/powershell/hostcommand.go | 2 +- pkg/powershell/hostcommand_test.go | 2 +- pkg/powershell/powershell.go | 19 +++++++++++++++++- pkg/powershell/runspace.go | 27 ++++++++++++++++++++------ scripts/update_bin.ps1 | 5 ++++- 12 files changed, 60 insertions(+), 31 deletions(-) diff --git a/.gitignore b/.gitignore index 032265a..79cb12b 100644 --- a/.gitignore +++ b/.gitignore @@ -7,5 +7,5 @@ c/** **.a host.h __debug_bin - +**debug.test diff --git a/examples/cmd/main.go b/examples/cmd/main.go index 6539400..18b3ba1 100644 --- a/examples/cmd/main.go +++ b/examples/cmd/main.go @@ -18,7 +18,7 @@ func (logger fmtPrintLogger) Write(arg string) { type callbackTest struct{} -func (c callbackTest) Callback(str string, input []powershell.Object, results powershell.CallbackResultsWriter) { +func (c callbackTest) Callback(_ powershell.Runspace, str string, input []powershell.Object, results powershell.CallbackResultsWriter) { fmt.Println("\tIn callback:", str) results.WriteString(str) for i, object := range input { diff --git a/native-powershell b/native-powershell index 29c3535..79d1e4c 160000 --- a/native-powershell +++ b/native-powershell @@ -1 +1 @@ -Subproject commit 29c3535b5a88f08604c1c1a2c700a664593ffa07 +Subproject commit 79d1e4c9f9211252ea82f80608d76d7713b06f2a diff --git a/pkg/powershell/callbackexample_test.go b/pkg/powershell/callbackexample_test.go index 2d44bb4..1a7f3ad 100644 --- a/pkg/powershell/callbackexample_test.go +++ b/pkg/powershell/callbackexample_test.go @@ -6,12 +6,10 @@ import ( ) // the callback we want to add -type callbackAdd10MultipleRunspace struct { - // to be able to execute powershell statements inside our callback we need a new runspace - newRunspace Runspace +type callbackAdd10Nested struct { } -func (callback *callbackAdd10MultipleRunspace) Callback(str string, input []Object, results CallbackResultsWriter) { +func (callback callbackAdd10Nested) Callback(runspace Runspace, str string, input []Object, results CallbackResultsWriter) { switch str { // check if we are processing the "add 10" message case "add 10": @@ -27,7 +25,7 @@ func (callback *callbackAdd10MultipleRunspace) Callback(str string, input []Obje // or write them back as a powershell integer // convert object into a powershell integer - execResults := callback.newRunspace.ExecScript(`[int]$args[0]`, true, nil, fmt.Sprint(num)) + execResults := runspace.ExecScript(`[int]$args[0]`, true, nil, fmt.Sprint(num)) // we need to close our execResults.Object[0] for us after it has been processed // however we do not know when that is, so tell the results to auto do it @@ -36,12 +34,8 @@ func (callback *callbackAdd10MultipleRunspace) Callback(str string, input []Obje } } } -func Example_powershellCallbackMultipleRunspace() { - // create a separate callback runspace - // you can share objects between runspaces, but you cannot execute 2 statements in the same runspace - // at the same time! - callback := &callbackAdd10MultipleRunspace{CreateRunspaceSimple()} - defer callback.newRunspace.Close() +func Example_powershellCallbackNested() { + callback := callbackAdd10Nested{} // create a runspace (where you run your powershell statements in) runspace := CreateRunspace(nil, callback) diff --git a/pkg/powershell/chelpers.go b/pkg/powershell/chelpers.go index 8c15bcc..2e6c376 100644 --- a/pkg/powershell/chelpers.go +++ b/pkg/powershell/chelpers.go @@ -59,7 +59,7 @@ func commandWchart(context uint64, cMessage *C.wchar_t, input *C.PowerShellObjec inputArr[i] = makePowerShellObjectIndexed(input, i) } message := makeString(cMessage) - contextInterface.Callback.Callback(message, inputArr, &resultsWriter) + contextInterface.Callback.Callback(contextInterface.recreateRunspace(), message, inputArr, &resultsWriter) } resultsWriter.filloutResults(ret) } diff --git a/pkg/powershell/context.go b/pkg/powershell/context.go index 1aff869..986db68 100644 --- a/pkg/powershell/context.go +++ b/pkg/powershell/context.go @@ -13,18 +13,18 @@ import ( var contextCache sync.Map var contextLookupKey uint64 -func storeRunspaceContext(context runspaceContext) uint64 { +func storeRunspaceContext(context *runspaceContext) uint64 { contextLookup := atomic.AddUint64(&contextLookupKey, 1) contextCache.Store(contextLookup, context) return contextLookup } -func getRunspaceContext(key uint64) runspaceContext { +func getRunspaceContext(key uint64) *runspaceContext { contextInterface, ok := contextCache.Load(key) if !ok { panic(fmt.Sprint("failed to load context key:", key)) } - return contextInterface.(runspaceContext) + return contextInterface.(*runspaceContext) } func deleteRunspaceContextLookup(key uint64) { contextCache.Delete(key) diff --git a/pkg/powershell/higherops_test.go b/pkg/powershell/higherops_test.go index 244d760..663709a 100644 --- a/pkg/powershell/higherops_test.go +++ b/pkg/powershell/higherops_test.go @@ -37,7 +37,7 @@ type callbackTest struct { lines *string } -func (c callbackTest) Callback(str string, input []Object, results CallbackResultsWriter) { +func (c callbackTest) Callback(_ Runspace, str string, input []Object, results CallbackResultsWriter) { record.Println(" In callback:", str) results.WriteString(str) for i, object := range input { @@ -228,7 +228,7 @@ func TestClocalScope(t *testing.T) { type callbackAddRef struct{} -func (c callbackAddRef) Callback(str string, input []Object, results CallbackResultsWriter) { +func (c callbackAddRef) Callback(_ Runspace, str string, input []Object, results CallbackResultsWriter) { results.WriteString(str) for _, object := range input { results.Write(object.AddRef(), true) @@ -254,7 +254,7 @@ type callbackAddRefSave struct { objects []Object } -func (c *callbackAddRefSave) Callback(str string, input []Object, results CallbackResultsWriter) { +func (c *callbackAddRefSave) Callback(runspace Runspace, str string, input []Object, results CallbackResultsWriter) { for _, object := range input { c.objects = append(c.objects, object.AddRef()) } diff --git a/pkg/powershell/hostcommand.go b/pkg/powershell/hostcommand.go index 83387d9..480133f 100644 --- a/pkg/powershell/hostcommand.go +++ b/pkg/powershell/hostcommand.go @@ -20,7 +20,7 @@ type CallbackResultsWriter interface { // CallbackHolder callback function pointer for Send-HostCommand callbacks type CallbackHolder interface { - Callback(str string, input []Object, results CallbackResultsWriter) + Callback(runspace Runspace, message string, input []Object, results CallbackResultsWriter) } // callbackResultsWriter is the internal implementation of CallbackResultsWriter diff --git a/pkg/powershell/hostcommand_test.go b/pkg/powershell/hostcommand_test.go index bfe2550..90e08a0 100644 --- a/pkg/powershell/hostcommand_test.go +++ b/pkg/powershell/hostcommand_test.go @@ -9,7 +9,7 @@ import ( type callbackAdd10 struct { } -func (callbackAdd10) Callback(str string, input []Object, results CallbackResultsWriter) { +func (callbackAdd10) Callback(runspace Runspace, str string, input []Object, results CallbackResultsWriter) { switch str { // check if we are processing the "add 10" message case "add 10": diff --git a/pkg/powershell/powershell.go b/pkg/powershell/powershell.go index d9f1cc3..1cc0eaa 100644 --- a/pkg/powershell/powershell.go +++ b/pkg/powershell/powershell.go @@ -21,6 +21,7 @@ import "C" // psCommand represents a powershell command, must call Close type psCommand struct { handle C.PowershellHandle + context * runspaceContext } // InvokeResults the results of an Invoke on a psCommand @@ -31,7 +32,17 @@ type InvokeResults struct { // createCommand using a runspace, still need to create a command in the powershell command func (runspace Runspace) createCommand() psCommand { - return psCommand{C.CreatePowershell(runspace.handle)} + currentlyInvoking := runspace.context.invoking + if(len(currentlyInvoking)!= 0){ + currentCommand:= currentlyInvoking[len(currentlyInvoking) -1] + return currentCommand.createNested() + } + return psCommand{C.CreatePowershell(runspace.handle), runspace.context} +} + +// createNested a nested powershell command +func (command psCommand) createNested() psCommand{ + return psCommand{C.CreatePowershellNested(command.handle), command.context} } // Close and free a psCommand @@ -100,6 +111,10 @@ func (command psCommand) AddParameter(paramName string, object Object) { _ = C.AddParameterObject(command.handle, (*C.wchar_t)(ptrName), object.handle) } +func (command psCommand) completeInvoke() { + command.context.invoking = command.context.invoking[:len(command.context.invoking)-1] +} + // Invoke the powershell command // // If wanting to call another powershell command do not reuse after Invoke, create another psCommand object and use that one @@ -109,6 +124,8 @@ func (command psCommand) Invoke() InvokeResults { var objects *C.PowerShellObject var count C.uint + command.context.invoking = append(command.context.invoking, command) + defer command.completeInvoke() exception := C.InvokeCommand(command.handle, &objects, &count) return makeInvokeResults(objects, count, exception) } diff --git a/pkg/powershell/runspace.go b/pkg/powershell/runspace.go index 8f09d4c..1faa37d 100644 --- a/pkg/powershell/runspace.go +++ b/pkg/powershell/runspace.go @@ -22,12 +22,22 @@ func init() { type runspaceContext struct { Log logger.Full Callback CallbackHolder + invoking []psCommand // in order list of psCommands we are currently invoking + + // runspaceContext should contain all the datamembers to reconstrut runspace + handle C.RunspaceHandle + contextLookup uint64 +} + +// recreateRunspace will give you a runspace from it's context +func (context *runspaceContext) recreateRunspace() Runspace{ + return Runspace{context.handle, context, context.contextLookup} } // Runspace a context handle for a runspace, use .Close() to free type Runspace struct { handle C.RunspaceHandle - context runspaceContext + context *runspaceContext contextLookup uint64 } @@ -42,9 +52,14 @@ func CreateRunspaceSimple() Runspace { // // You must call Close when done with this object func CreateRunspace(loggerCallback logger.Simple, callback CallbackHolder) Runspace { - context := runspaceContext{logger.MakeLoggerFull(loggerCallback), callback} - contextLookup := storeRunspaceContext(context) - + context := &runspaceContext{Log: logger.MakeLoggerFull(loggerCallback), + Callback: callback, + invoking: nil, + handle: C.ulonglong(0), + contextLookup: 0, + } + context.contextLookup = storeRunspaceContext(context) + var useLogger C.char = 1 if loggerCallback == nil { useLogger = 0 @@ -53,8 +68,8 @@ func CreateRunspace(loggerCallback logger.Simple, callback CallbackHolder) Runsp if callback == nil { useCommand = 0 } - runspace := C.CreateRunspaceHelper(C.ulonglong(contextLookup), useLogger, useCommand) - return Runspace{runspace, context, contextLookup} + context.handle = C.CreateRunspaceHelper(C.ulonglong(context.contextLookup), useLogger, useCommand) + return context.recreateRunspace() } // Close and free a Runspace diff --git a/scripts/update_bin.ps1 b/scripts/update_bin.ps1 index 2017d18..491ce9d 100644 --- a/scripts/update_bin.ps1 +++ b/scripts/update_bin.ps1 @@ -2,4 +2,7 @@ cp .\native-powershell\host.h .\native-powershell\native-powershell-bin\ cp .\native-powershell\x64\Release\psh_host.dll .\native-powershell\native-powershell-bin\ cp .\native-powershell\x64\Release\psh_host.pdb .\native-powershell\native-powershell-bin\ -cp .\native-powershell\x64\Release\psh_host.lib .\native-powershell\native-powershell-bin\ \ No newline at end of file +cp .\native-powershell\x64\Release\psh_host.lib .\native-powershell\native-powershell-bin\ + +cp .\native-powershell\native-powershell-bin\psh_host.dll .\ +cp .\native-powershell\native-powershell-bin\psh_host.dll .\pkg\powershell \ No newline at end of file