Merge pull request #50 from dylanhitt/AttachCtx

Attach ctx

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)
 }
 

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
+		}
+	}
+}

frida/session.go

@@ -3,6 +3,7 @@ package frida
 //#include <frida-core.h>
 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)
 }