new plugins system

Overlord-Client/cmd/agent/plugins/loader_windows.go

@@ -5,7 +5,7 @@ package plugins
 import (
 	"errors"
 	"fmt"
-	"sync"
+	"runtime"
 	"syscall"
 	"unsafe"
 )
@@ -15,59 +15,96 @@ func loadNativePlugin(data []byte) (NativePlugin, error) {
 		return nil, errors.New("empty plugin binary")
 	}
 
-	mm, err := LoadMemoryModule(data)
-	if err != nil {
-		return nil, fmt.Errorf("pe load: %w", err)
+	type initResult struct {
+		dp  *dllPlugin
+		err error
 	}
+	ch := make(chan initResult, 1)
 
-	if err := mm.CallEntryPoint(dllProcessAttach); err != nil {
-		mm.Free()
-		return nil, fmt.Errorf("DllMain init: %w", err)
-	}
+	go func() {
+		runtime.LockOSThread()
 
-	onLoad, err := mm.GetExport("PluginOnLoad")
-	if err != nil {
-		mm.Free()
-		return nil, err
-	}
-	onEvent, err := mm.GetExport("PluginOnEvent")
-	if err != nil {
-		mm.Free()
-		return nil, err
-	}
-	onUnload, err := mm.GetExport("PluginOnUnload")
-	if err != nil {
-		mm.Free()
-		return nil, err
-	}
+		mm, err := LoadMemoryModule(data)
+		if err != nil {
+			ch <- initResult{err: fmt.Errorf("pe load: %w", err)}
+			runtime.UnlockOSThread()
+			return
+		}
 
-	setCallback, _ := mm.GetExport("PluginSetCallback")
+		if err := mm.CallEntryPoint(dllProcessAttach); err != nil {
+			mm.Free()
+			ch <- initResult{err: fmt.Errorf("DllMain init: %w", err)}
+			runtime.UnlockOSThread()
+			return
+		}
 
-	runtime := "go"
-	if getRuntimeAddr, err := mm.GetExport("PluginGetRuntime"); err == nil {
-		ret, _, _ := syscall.SyscallN(getRuntimeAddr)
-		if ret != 0 {
-			var buf [32]byte
-			for i := range buf {
-				b := *(*byte)(unsafe.Pointer(ret + uintptr(i)))
-				if b == 0 {
-					runtime = string(buf[:i])
-					break
+		onLoad, err := mm.GetExport("PluginOnLoad")
+		if err != nil {
+			mm.Free()
+			ch <- initResult{err: err}
+			runtime.UnlockOSThread()
+			return
+		}
+		onEvent, err := mm.GetExport("PluginOnEvent")
+		if err != nil {
+			mm.Free()
+			ch <- initResult{err: err}
+			runtime.UnlockOSThread()
+			return
+		}
+		onUnload, err := mm.GetExport("PluginOnUnload")
+		if err != nil {
+			mm.Free()
+			ch <- initResult{err: err}
+			runtime.UnlockOSThread()
+			return
+		}
+
+		setCallback, _ := mm.GetExport("PluginSetCallback")
+
+		rt := "go"
+		if getRuntimeAddr, err := mm.GetExport("PluginGetRuntime"); err == nil {
+			ret, _, _ := syscall.SyscallN(getRuntimeAddr)
+			if ret != 0 {
+				var buf [32]byte
+				for i := range buf {
+					b := *(*byte)(unsafe.Pointer(ret + uintptr(i)))
+					if b == 0 {
+						rt = string(buf[:i])
+						break
+					}
+					buf[i] = b
 				}
-				buf[i] = b
 			}
 		}
-	}
 
-	dp := &dllPlugin{
-		mem:             mm,
-		onLoadAddr:      onLoad,
-		onEventAddr:     onEvent,
-		onUnloadAddr:    onUnload,
-		setCallbackAddr: setCallback,
-		pluginRuntime:   runtime,
+		dp := &dllPlugin{
+			mem:             mm,
+			onLoadAddr:      onLoad,
+			onEventAddr:     onEvent,
+			onUnloadAddr:    onUnload,
+			setCallbackAddr: setCallback,
+			pluginRuntime:   rt,
+			workCh:          make(chan pluginWork),
+		}
+
+		ch <- initResult{dp: dp}
+
+		dp.workerLoop()
+
+		runtime.UnlockOSThread()
+	}()
+
+	res := <-ch
+	if res.err != nil {
+		return nil, res.err
 	}
-	return dp, nil
+	return res.dp, nil
+}
+
+type pluginWork struct {
+	fn   func() error
+	done chan error
 }
 
 type dllPlugin struct {
@@ -78,76 +115,94 @@ type dllPlugin struct {
 	setCallbackAddr uintptr
 	callbackHandle  uintptr // prevent GC of the callback closure
 	pluginRuntime   string  // "go", "c", "cpp", etc.
-	mu              sync.Mutex
+	workCh          chan pluginWork
+}
+
+func (p *dllPlugin) workerLoop() {
+	for w := range p.workCh {
+		w.done <- w.fn()
+	}
+}
+
+func (p *dllPlugin) runOnWorker(fn func() error) error {
+	done := make(chan error, 1)
+	p.workCh <- pluginWork{fn: fn, done: done}
+	return <-done
 }
 
 func (p *dllPlugin) Load(send func(string, []byte), hostInfo []byte) error {
-	p.mu.Lock()
-	defer p.mu.Unlock()
+	return p.runOnWorker(func() error {
+		// Create a stdcall callback the DLL can invoke to send events to the host.
+		cb := syscall.NewCallback(func(eventPtr, eventLen, payloadPtr, payloadLen uintptr) uintptr {
+			event := make([]byte, eventLen)
+			if eventLen > 0 {
+				copy(event, unsafe.Slice((*byte)(unsafe.Pointer(eventPtr)), eventLen))
+			}
+			payload := make([]byte, payloadLen)
+			if payloadLen > 0 {
+				copy(payload, unsafe.Slice((*byte)(unsafe.Pointer(payloadPtr)), payloadLen))
+			}
+			send(string(event), payload)
+			return 0
+		})
+		p.callbackHandle = cb
 
-	// Create a stdcall callback the DLL can invoke to send events to the host.
-	cb := syscall.NewCallback(func(eventPtr, eventLen, payloadPtr, payloadLen uintptr) uintptr {
-		event := make([]byte, eventLen)
-		if eventLen > 0 {
-			copy(event, unsafe.Slice((*byte)(unsafe.Pointer(eventPtr)), eventLen))
+		if p.setCallbackAddr != 0 {
+			syscall.SyscallN(p.setCallbackAddr, cb)
 		}
-		payload := make([]byte, payloadLen)
-		if payloadLen > 0 {
-			copy(payload, unsafe.Slice((*byte)(unsafe.Pointer(payloadPtr)), payloadLen))
+
+		var infoPtr uintptr
+		infoLen := uintptr(len(hostInfo))
+		if len(hostInfo) > 0 {
+			infoPtr = uintptr(unsafe.Pointer(&hostInfo[0]))
 		}
-		send(string(event), payload)
-		return 0
+		ret, _, _ := syscall.SyscallN(p.onLoadAddr, infoPtr, infoLen, cb)
+		if int32(ret) != 0 {
+			return errors.New("PluginOnLoad returned non-zero")
+		}
+		return nil
 	})
-	p.callbackHandle = cb
-
-	if p.setCallbackAddr != 0 {
-		syscall.SyscallN(p.setCallbackAddr, cb)
-	}
-
-	var infoPtr uintptr
-	infoLen := uintptr(len(hostInfo))
-	if len(hostInfo) > 0 {
-		infoPtr = uintptr(unsafe.Pointer(&hostInfo[0]))
-	}
-	ret, _, _ := syscall.SyscallN(p.onLoadAddr, infoPtr, infoLen, cb)
-	if int32(ret) != 0 {
-		return errors.New("PluginOnLoad returned non-zero")
-	}
-	return nil
 }
 
 func (p *dllPlugin) Event(event string, payload []byte) error {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-
 	eventBytes := []byte(event)
-	var eventPtr, payloadPtr uintptr
-	eventLen := uintptr(len(eventBytes))
-	payloadLen := uintptr(len(payload))
-	if len(eventBytes) > 0 {
-		eventPtr = uintptr(unsafe.Pointer(&eventBytes[0]))
-	}
-	if len(payload) > 0 {
-		payloadPtr = uintptr(unsafe.Pointer(&payload[0]))
-	}
+	payloadCopy := make([]byte, len(payload))
+	copy(payloadCopy, payload)
 
-	ret, _, _ := syscall.SyscallN(p.onEventAddr, eventPtr, eventLen, payloadPtr, payloadLen)
-	if int32(ret) != 0 {
-		return errors.New("PluginOnEvent returned non-zero")
-	}
+	go func() {
+		runtime.LockOSThread()
+		defer runtime.UnlockOSThread()
+
+		cleanup := p.mem.SetupThreadTLS()
+		defer cleanup()
+
+		var eventPtr, payloadPtr uintptr
+		eventLen := uintptr(len(eventBytes))
+		payloadLen := uintptr(len(payloadCopy))
+		if len(eventBytes) > 0 {
+			eventPtr = uintptr(unsafe.Pointer(&eventBytes[0]))
+		}
+		if len(payloadCopy) > 0 {
+			payloadPtr = uintptr(unsafe.Pointer(&payloadCopy[0]))
+		}
+
+		syscall.SyscallN(p.onEventAddr, eventPtr, eventLen, payloadPtr, payloadLen)
+	}()
 	return nil
 }
 
 func (p *dllPlugin) Unload() {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	if p.onUnloadAddr != 0 {
-		syscall.SyscallN(p.onUnloadAddr)
-	}
+	_ = p.runOnWorker(func() error {
+		if p.onUnloadAddr != 0 {
+			syscall.SyscallN(p.onUnloadAddr)
+		}
+		return nil
+	})
 }
 
 func (p *dllPlugin) Close() error {
 	p.Unload()
+	close(p.workCh)
 	if p.pluginRuntime != "go" {
 		if p.mem != nil {
 			p.mem.Free()