Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
projectdiscovery
GitHub Repository: projectdiscovery/nuclei
Path: blob/dev/pkg/js/gojs/set.go
2070 views
1
package gojs
2
3
import (
4
"context"
5
"reflect"
6
7
"github.com/Mzack9999/goja"
8
"github.com/projectdiscovery/utils/errkit"
9
)
10
11
var (
12
ErrInvalidFuncOpts = errkit.New("invalid function options")
13
ErrNilRuntime = errkit.New("runtime is nil")
14
)
15
16
type FuncOpts struct {
17
Name string
18
Signatures []string
19
Description string
20
FuncDecl interface{}
21
}
22
23
// valid checks if the function options are valid
24
func (f *FuncOpts) valid() bool {
25
return f.Name != "" && f.FuncDecl != nil && len(f.Signatures) > 0 && f.Description != ""
26
}
27
28
// wrapWithContext wraps a Go function with context injection
29
// nolint
30
func wrapWithContext(runtime *goja.Runtime, fn interface{}) interface{} {
31
fnType := reflect.TypeOf(fn)
32
if fnType.Kind() != reflect.Func {
33
return fn
34
}
35
36
// Only wrap if first parameter is context.Context
37
if fnType.NumIn() == 0 || fnType.In(0) != reflect.TypeFor[context.Context]() {
38
return fn // Return original function unchanged if it doesn't have context.Context as first arg
39
}
40
41
// Create input and output type slices
42
inTypes := make([]reflect.Type, fnType.NumIn())
43
for i := 0; i < fnType.NumIn(); i++ {
44
inTypes[i] = fnType.In(i)
45
}
46
outTypes := make([]reflect.Type, fnType.NumOut())
47
for i := 0; i < fnType.NumOut(); i++ {
48
outTypes[i] = fnType.Out(i)
49
}
50
51
// Create a new function with same signature
52
newFnType := reflect.FuncOf(inTypes, outTypes, fnType.IsVariadic())
53
newFn := reflect.MakeFunc(newFnType, func(args []reflect.Value) []reflect.Value {
54
// Get context from runtime
55
var ctx context.Context
56
if ctxVal := runtime.Get("context"); ctxVal != nil {
57
if ctxObj, ok := ctxVal.Export().(context.Context); ok {
58
ctx = ctxObj
59
}
60
}
61
if ctx == nil {
62
ctx = context.Background()
63
}
64
65
// Add execution ID to context if available
66
if execID := runtime.Get("executionId"); execID != nil {
67
ctx = context.WithValue(ctx, "executionId", execID.String())
68
}
69
70
// Replace first argument (context) with our context
71
args[0] = reflect.ValueOf(ctx)
72
73
// Call original function with modified arguments
74
return reflect.ValueOf(fn).Call(args)
75
})
76
77
return newFn.Interface()
78
}
79
80
// RegisterFunc registers a function with given name, signatures and description
81
func RegisterFuncWithSignature(runtime *goja.Runtime, opts FuncOpts) error {
82
if runtime == nil {
83
return ErrNilRuntime
84
}
85
if !opts.valid() {
86
return errkit.Newf("invalid function options: name: %s, signatures: %v, description: %s", opts.Name, opts.Signatures, opts.Description)
87
}
88
89
// Wrap the function with context injection
90
// wrappedFn := wrapWithContext(runtime, opts.FuncDecl)
91
return runtime.Set(opts.Name, opts.FuncDecl /* wrappedFn */)
92
}
93
94