diff --git a/session/script_builtin_runtime.go b/session/script_builtin_runtime.go index ffcd14b4..04d02d56 100644 --- a/session/script_builtin_runtime.go +++ b/session/script_builtin_runtime.go @@ -2,8 +2,10 @@ package session import ( "encoding/json" + "fmt" "io/ioutil" "os" + "sync" "github.com/bettercap/bettercap/v2/js" "github.com/evilsocket/islazy/fs" @@ -14,6 +16,8 @@ import ( // see https://github.com/robertkrimen/otto/issues/213 var jsRuntime = otto.New() +var jsListeners = sync.Map{} + func jsRunFunc(call otto.FunctionCall) otto.Value { argv := call.ArgumentList argc := len(argv) @@ -57,40 +61,94 @@ func jsOnEventFunc(call otto.FunctionCall) otto.Value { cb = argv[1] } + listenerKey := fmt.Sprintf("%s:%s", filterExpr, cb.String()) + + if _, found := jsListeners.Load(listenerKey); found { + return js.ReportError("listener already exists") + } + + // add to listeners + closeChan := make(chan bool) + jsListeners.Store(listenerKey, closeChan) + // start a go routine for this event listener - go func(expr string, cb otto.Value) { + go func(expr string, cb otto.Value, closeChan chan bool) { listener := I.Events.Listen() defer I.Events.Unlisten(listener) + defer close(closeChan) - for event := range listener { - if expr == "" || event.Tag == expr { - // some objects don't do well with js, so convert them to a generic map - // before passing them to the callback - var opaque interface{} - if raw, err := json.Marshal(event); err != nil { - I.Events.Log(log.ERROR, "error serializing event %s: %v", event.Tag, err) - } else if err = json.Unmarshal(raw, &opaque); err != nil { - I.Events.Log(log.ERROR, "error serializing event %s: %v", event.Tag, err) + for { + select { + case event := <-listener: + if expr == "" || event.Tag == expr { + // some objects don't do well with js, so convert them to a generic map + // before passing them to the callback + var opaque interface{} + if raw, err := json.Marshal(event); err != nil { + I.Events.Log(log.ERROR, "error serializing event %s: %v", event.Tag, err) + } else if err = json.Unmarshal(raw, &opaque); err != nil { + I.Events.Log(log.ERROR, "error serializing event %s: %v", event.Tag, err) + } + + // lock vm if ready and available + locked := false + if I.script != nil { + I.script.Lock() + locked = true + } + + if _, err := cb.Call(otto.NullValue(), opaque); err != nil { + I.Events.Log(log.ERROR, "error dispatching event %s: %v", event.Tag, err) + } + + // unlock vm if ready and available + if locked { + I.script.Unlock() + } } - // lock vm if ready and available - locked := false - if I.script != nil { - I.script.Lock() - locked = true - } - - if _, err := cb.Call(otto.NullValue(), opaque); err != nil { - I.Events.Log(log.ERROR, "error dispatching event %s: %v", event.Tag, err) - } - - // unlock vm if ready and available - if locked { - I.script.Unlock() - } + case <-closeChan: + return } + } - }(filterExpr, cb) + }(filterExpr, cb, closeChan) + + return js.NullValue +} + +func jsRemoveEventListenerFunc(call otto.FunctionCall) otto.Value { + argv := call.ArgumentList + argc := len(argv) + cb := otto.NullValue() + filterExpr := "" + + // just one argument, a function to receive all events + if argc == 1 { + if argv[0].IsFunction() == false { + return js.ReportError("the single argument must be a function") + } + cb = argv[0] + } else { + if argc != 2 { + return js.ReportError("expected two arguments (event_name, callback), got %d", argc) + } else if argv[0].IsString() == false { + return js.ReportError("first argument must be a string") + } else if argv[1].IsFunction() == false { + return js.ReportError("second argument must be a function") + } + + filterExpr = argv[0].String() + cb = argv[1] + } + + listenerKey := fmt.Sprintf("%s:%s", filterExpr, cb.String()) + if closer, found := jsListeners.Load(listenerKey); found { + closer.(chan bool) <- true + jsListeners.Delete(listenerKey) + } else { + return js.ReportError("listener not found") + } return js.NullValue } diff --git a/session/session.go b/session/session.go index 4d0d0ef6..983ef1a2 100644 --- a/session/session.go +++ b/session/session.go @@ -330,6 +330,7 @@ func (s *Session) Start() error { plugin.Defines["saveJSON"] = jsSaveJSONFunc plugin.Defines["saveToFile"] = jsSaveToFileFunc plugin.Defines["onEvent"] = jsOnEventFunc + plugin.Defines["removeEventListener"] = jsRemoveEventListenerFunc plugin.Defines["session"] = s // load the script here so the session and its internal objects are ready