diff --git a/frida/device.go b/frida/device.go index 3891992..82b7572 100644 --- a/frida/device.go +++ b/frida/device.go @@ -20,6 +20,7 @@ type DeviceInt interface { Manager() *DeviceManager IsLost() bool Params(opts ...OptFunc) (map[string]any, error) + ParamsWithContext(ctx context.Context) (map[string]any, error) FrontmostApplication(scope Scope) (*Application, error) EnumerateApplications(identifier string, scope Scope, opts ...OptFunc) ([]*Application, error) ProcessByPID(pid int, scope Scope) (*Process, error) @@ -35,7 +36,8 @@ type DeviceInt interface { Input(pid int, data []byte) error Resume(pid int) error Kill(pid int) error - Attach(val any, opts *SessionOptions) (*Session, error) + Attach(val any, sessionOpts *SessionOptions, opts ...OptFunc) (*Session, error) + AttachWithContext(ctx context.Context, val any, opts *SessionOptions) (*Session, error) InjectLibraryFile(target any, path, entrypoint, data string) (uint, error) InjectLibraryBlob(target any, byteData []byte, entrypoint, data string) (uint, error) OpenChannel(address string) (*IOStream, error) @@ -113,36 +115,20 @@ func (d *Device) IsLost() bool { return false } -// ParamsCtx runs Params but with context. +// ParamsWithContext runs Params but with context. // This function will properly handle cancelling the frida operation. // It is advised to use this rather than handling Cancellable yourself. -func (d *Device) ParamsCtx(ctx context.Context) (map[string]any, error) { - paramC := make(chan map[string]any, 1) - errC := make(chan error, 1) - - c := NewCancellable() - go func() { +func (d *Device) ParamsWithContext(ctx context.Context) (map[string]any, error) { + rawParams, err := handleWithContext(ctx, func(c *Cancellable, doneC chan any, errC chan error) { params, err := d.Params(WithCancel(c)) if err != nil { errC <- err return } - paramC <- params - }() - - for { - select { - case <-ctx.Done(): - c.Cancel() - return nil, ErrContextCancelled - case params := <-paramC: - c.Unref() - return params, nil - case err := <-errC: - c.Unref() - return nil, err - } - } + doneC <- params + }) + params, _ := rawParams.(map[string]any) + return params, err } // Params returns system parameters of the device @@ -486,10 +472,31 @@ func (d *Device) Kill(pid int) error { return handleGError(err) } +// AttachWithContext runs Attach but with context. +// This function will properly handle cancelling the frida operation. +// It is advised to use this rather than handling Cancellable yourself. +func (d *Device) AttachWithContext(ctx context.Context, val any, sessionOpts *SessionOptions) (*Session, error) { + rawSession, err := handleWithContext(ctx, func(c *Cancellable, doneC chan any, errC chan error) { + session, err := d.Attach(val, sessionOpts, WithCancel(c)) + if err != nil { + errC <- err + return + } + doneC <- session + }) + session, _ := rawSession.(*Session) + return session, err +} + // Attach will attach on specified process name or PID. // You can pass the nil as SessionOptions or you can create it if you want // the session to persist for specific timeout. -func (d *Device) Attach(val any, opts *SessionOptions) (*Session, error) { +func (d *Device) Attach(val any, sessionOpts *SessionOptions, opts ...OptFunc) (*Session, error) { + o := setupOptions(opts) + return d.attach(val, sessionOpts, o) +} + +func (d *Device) attach(val any, sessionOpts *SessionOptions, opts options) (*Session, error) { if d.device == nil { return nil, errors.New("could not attach for nil device") } @@ -508,13 +515,13 @@ func (d *Device) Attach(val any, opts *SessionOptions) (*Session, error) { } var opt *C.FridaSessionOptions = nil - if opts != nil { - opt = opts.opts + if sessionOpts != nil { + opt = sessionOpts.opts defer clean(unsafe.Pointer(opt), unrefFrida) } var err *C.GError - s := C.frida_device_attach_sync(d.device, C.guint(pid), opt, nil, &err) + s := C.frida_device_attach_sync(d.device, C.guint(pid), opt, opts.cancellable, &err) return &Session{s}, handleGError(err) } diff --git a/frida/misc.go b/frida/misc.go index a00b71c..424df1b 100644 --- a/frida/misc.go +++ b/frida/misc.go @@ -4,6 +4,7 @@ package frida import "C" import ( + "context" "encoding/json" "fmt" "unsafe" @@ -116,3 +117,25 @@ func ScriptMessageToMessage(message string) (*Message, error) { } return &m, nil } + +func handleWithContext(ctx context.Context, f func(c *Cancellable, done chan any, errC chan error)) (any, error) { + doneC := make(chan any, 1) + errC := make(chan error, 1) + + c := NewCancellable() + go f(c, doneC, errC) + + for { + select { + case <-ctx.Done(): + c.Cancel() + return nil, ErrContextCancelled + case done := <-doneC: + c.Unref() + return done, nil + case err := <-errC: + c.Unref() + return nil, err + } + } +} diff --git a/frida/session.go b/frida/session.go index 49c5895..10acecb 100644 --- a/frida/session.go +++ b/frida/session.go @@ -3,6 +3,7 @@ package frida //#include import "C" import ( + "context" "runtime" "unsafe" ) @@ -18,10 +19,26 @@ func (s *Session) IsDetached() bool { return int(detached) == 1 } +// DetachWithContext runs Detach but with context. +// This function will properly handle cancelling the frida operation. +// It is advised to use this rather than handling Cancellable yourself. +func (s *Session) DetachWithContext(ctx context.Context) error { + _, err := handleWithContext(ctx, func(c *Cancellable, done chan any, errC chan error) { + errC <- s.Detach(WithCancel(c)) + }) + return err +} + // Detach detaches the current session. -func (s *Session) Detach() error { +func (s *Session) Detach(opts ...OptFunc) error { + o := setupOptions(opts) + return s.detach(o) +} + +// detach +func (s *Session) detach(opts options) error { var err *C.GError - C.frida_session_detach_sync(s.s, nil, &err) + C.frida_session_detach_sync(s.s, opts.cancellable, &err) return handleGError(err) }