Allow to execute specific function

Signed-off-by: Julien Fabre <[email protected]>

internal/timecraft/module.go

@@ -15,6 +15,10 @@ type ModuleSpec struct {
 	// Path is the path of the WebAssembly module.
 	Path string
 
+	// Name of the exported function to call in the WebAssembly module.
+	// If empty, the _start function will be executed.
+	Function string
+
 	// Args are command-line arguments to pass to the module.
 	Args []string
 

internal/timecraft/process.go

@@ -132,7 +132,7 @@ func (pm *ProcessManager) Start(moduleSpec ModuleSpec, logSpec *LogSpec, parentI
 	if err != nil {
 		return ProcessID{}, fmt.Errorf("could not read wasm file '%s': %w", wasmPath, err)
 	}
-
+	function := moduleSpec.Function
 	wasmModule, err := pm.runtime.CompileModule(pm.ctx, wasmCode)
 	if err != nil {
 		return ProcessID{}, err
@@ -233,6 +233,9 @@ func (pm *ProcessManager) Start(moduleSpec ModuleSpec, logSpec *LogSpec, parentI
 		}, object.Tag{
 			Name:  "timecraft.module.name",
 			Value: wasmModule.Name(),
+		}, object.Tag{
+			Name:  "timecraft.module.function",
+			Value: function,
 		})
 		if err != nil {
 			return ProcessID{}, err
@@ -388,7 +391,7 @@ func (pm *ProcessManager) Start(moduleSpec ModuleSpec, logSpec *LogSpec, parentI
 
 	// Run the module in the background, and tidy up once complete.
 	pm.group.Go(func() error {
-		err := runModule(ctx, pm.runtime, wasmModule)
+		err := runModule(ctx, pm.runtime, wasmModule, function)
 		cancel(err)
 
 		pm.mu.Lock()

internal/timecraft/replay.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/stealthrocket/timecraft/format"
 	"github.com/stealthrocket/timecraft/internal/stream"
 	"github.com/stealthrocket/timecraft/internal/timemachine"
 	"github.com/stealthrocket/timecraft/internal/timemachine/wasicall"
@@ -55,7 +56,7 @@ func (r *Replay) SetTrace(w io.Writer) {
 
 // Replay replays process execution.
 func (r *Replay) Replay(ctx context.Context) error {
-	moduleCode, err := r.ModuleCode(ctx)
+	moduleCode, function, err := r.ModuleCode(ctx)
 	if err != nil {
 		return err
 	}
@@ -66,28 +67,39 @@ func (r *Replay) Replay(ctx context.Context) error {
 	}
 	defer records.Close()
 
-	return r.ReplayRecords(ctx, moduleCode, records)
+	return r.ReplayRecords(ctx, function, moduleCode, records)
 }
 
 // ModuleCode reads the module's WebAssembly code.
-func (r *Replay) ModuleCode(ctx context.Context) ([]byte, error) {
+func (r *Replay) ModuleCode(ctx context.Context) ([]byte, string, error) {
 	manifest, err := r.registry.LookupLogManifest(ctx, r.processID)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
+	fn := lookupFunction(manifest)
 	process, err := r.registry.LookupProcess(ctx, manifest.Process.Digest)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	processConfig, err := r.registry.LookupConfig(ctx, process.Config.Digest)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 	module, err := r.registry.LookupModule(ctx, processConfig.Modules[0].Digest)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
-	return module.Code, nil
+	return module.Code, fn, nil
+}
+
+func lookupFunction(m *format.Manifest) string {
+	for key, value := range m.Process.Annotations {
+		// TODO: create well known tags list
+		if key == "timecraft.module.function" {
+			return value
+		}
+	}
+	return ""
 }
 
 // RecordReader constructs a reader for the process replay log.
@@ -108,7 +120,7 @@ func (r *Replay) RecordReader(ctx context.Context) (records stream.ReadCloser[ti
 
 // ReplayRecordsModule replays process execution using the specified records on
 // a pre-compiled module.
-func (r *Replay) ReplayRecordsModule(ctx context.Context, compiledModule wazero.CompiledModule, records stream.Reader[timemachine.Record]) error {
+func (r *Replay) ReplayRecordsModule(ctx context.Context, function string, compiledModule wazero.CompiledModule, records stream.Reader[timemachine.Record]) error {
 	replay := wasicall.NewReplay(records)
 	defer replay.Close(ctx)
 
@@ -128,17 +140,17 @@ func (r *Replay) ReplayRecordsModule(ctx context.Context, compiledModule wazero.
 	hostModuleInstance := wazergo.MustInstantiate(ctx, r.runtime, hostModule, wasi_snapshot_preview1.WithWASI(system))
 	ctx = wazergo.WithModuleInstance(ctx, hostModuleInstance)
 
-	return runModule(ctx, r.runtime, compiledModule)
+	return runModule(ctx, r.runtime, compiledModule, function)
 }
 
 // ReplayRecords replays process execution using the specified records.
-func (r *Replay) ReplayRecords(ctx context.Context, moduleCode []byte, records stream.Reader[timemachine.Record]) error {
+func (r *Replay) ReplayRecords(ctx context.Context, function string, moduleCode []byte, records stream.Reader[timemachine.Record]) error {
 	compiledModule, err := r.runtime.CompileModule(ctx, moduleCode)
 	if err != nil {
 		return err
 	}
 	defer compiledModule.Close(ctx)
-	return r.ReplayRecordsModule(ctx, compiledModule, records)
+	return r.ReplayRecordsModule(ctx, function, compiledModule, records)
 }
 
 type recordReadCloser struct {

internal/timecraft/run.go

@@ -8,15 +8,21 @@ import (
 	"github.com/tetratelabs/wazero/sys"
 )
 
-func runModule(ctx context.Context, runtime wazero.Runtime, compiledModule wazero.CompiledModule) error {
+const defaultFunction = "_start"
+
+func runModule(ctx context.Context, runtime wazero.Runtime, compiledModule wazero.CompiledModule, fn string) error {
+	if fn == "" {
+		fn = defaultFunction
+	}
+
 	module, err := runtime.InstantiateModule(ctx, compiledModule, wazero.NewModuleConfig().
 		WithStartFunctions())
 	if err != nil {
 		return err
 	}
 	defer module.Close(ctx)
 
-	_, err = module.ExportedFunction("_start").Call(ctx)
+	_, err = module.ExportedFunction(fn).Call(ctx)
 	switch err {
 	case context.Canceled, context.DeadlineExceeded:
 		err = nil

profile.go

@@ -115,7 +115,7 @@ func profile(ctx context.Context, args []string) error {
 		startTime = human.Time(processStartTime)
 	}
 
-	moduleCode, err := replay.ModuleCode(ctx)
+	moduleCode, function, err := replay.ModuleCode(ctx)
 	if err != nil {
 		return err
 	}
@@ -152,7 +152,7 @@ func profile(ctx context.Context, args []string) error {
 		return err
 	}
 
-	if err := replay.ReplayRecordsModule(ctx, compiledModule, records); err != nil {
+	if err := replay.ReplayRecordsModule(ctx, function, compiledModule, records); err != nil {
 		return err
 	}
 

run.go

@@ -28,6 +28,7 @@ Options:
    -D, --dial addr                Expose a socket connected to the specified address
        --dir dir                  Expose a directory to the guest module
    -e, --env name=value           Pass an environment variable to the guest module
+   -f, --function function        Exported function to call in the guest module (_start if empty)
        --fly-blind                Disable recording of the guest module execution
    -h, --help                     Show this usage information
    -L, --listen addr              Expose a socket listening on the specified address