script.go 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338
  1. package tengo
  2. import (
  3. "context"
  4. "fmt"
  5. "path/filepath"
  6. "sync"
  7. "github.com/d5/tengo/v2/parser"
  8. )
  9. // Script can simplify compilation and execution of embedded scripts.
  10. type Script struct {
  11. variables map[string]*Variable
  12. modules ModuleGetter
  13. input []byte
  14. maxAllocs int64
  15. maxConstObjects int
  16. enableFileImport bool
  17. importDir string
  18. }
  19. // NewScript creates a Script instance with an input script.
  20. func NewScript(input []byte) *Script {
  21. return &Script{
  22. variables: make(map[string]*Variable),
  23. input: input,
  24. maxAllocs: -1,
  25. maxConstObjects: -1,
  26. }
  27. }
  28. // Add adds a new variable or updates an existing variable to the script.
  29. func (s *Script) Add(name string, value interface{}) error {
  30. obj, err := FromInterface(value)
  31. if err != nil {
  32. return err
  33. }
  34. s.variables[name] = &Variable{
  35. name: name,
  36. value: obj,
  37. }
  38. return nil
  39. }
  40. // Remove removes (undefines) an existing variable for the script. It returns
  41. // false if the variable name is not defined.
  42. func (s *Script) Remove(name string) bool {
  43. if _, ok := s.variables[name]; !ok {
  44. return false
  45. }
  46. delete(s.variables, name)
  47. return true
  48. }
  49. // SetImports sets import modules.
  50. func (s *Script) SetImports(modules ModuleGetter) {
  51. s.modules = modules
  52. }
  53. // SetImportDir sets the initial import directory for script files.
  54. func (s *Script) SetImportDir(dir string) error {
  55. dir, err := filepath.Abs(dir)
  56. if err != nil {
  57. return err
  58. }
  59. s.importDir = dir
  60. return nil
  61. }
  62. // SetMaxAllocs sets the maximum number of objects allocations during the run
  63. // time. Compiled script will return ErrObjectAllocLimit error if it
  64. // exceeds this limit.
  65. func (s *Script) SetMaxAllocs(n int64) {
  66. s.maxAllocs = n
  67. }
  68. // SetMaxConstObjects sets the maximum number of objects in the compiled
  69. // constants.
  70. func (s *Script) SetMaxConstObjects(n int) {
  71. s.maxConstObjects = n
  72. }
  73. // EnableFileImport enables or disables module loading from local files. Local
  74. // file modules are disabled by default.
  75. func (s *Script) EnableFileImport(enable bool) {
  76. s.enableFileImport = enable
  77. }
  78. // Compile compiles the script with all the defined variables, and, returns
  79. // Compiled object.
  80. func (s *Script) Compile() (*Compiled, error) {
  81. symbolTable, globals, err := s.prepCompile()
  82. if err != nil {
  83. return nil, err
  84. }
  85. fileSet := parser.NewFileSet()
  86. srcFile := fileSet.AddFile("(main)", -1, len(s.input))
  87. p := parser.NewParser(srcFile, s.input, nil)
  88. file, err := p.ParseFile()
  89. if err != nil {
  90. return nil, err
  91. }
  92. c := NewCompiler(srcFile, symbolTable, nil, s.modules, nil)
  93. c.EnableFileImport(s.enableFileImport)
  94. c.SetImportDir(s.importDir)
  95. if err := c.Compile(file); err != nil {
  96. return nil, err
  97. }
  98. // reduce globals size
  99. globals = globals[:symbolTable.MaxSymbols()+1]
  100. // global symbol names to indexes
  101. globalIndexes := make(map[string]int, len(globals))
  102. for _, name := range symbolTable.Names() {
  103. symbol, _, _ := symbolTable.Resolve(name, false)
  104. if symbol.Scope == ScopeGlobal {
  105. globalIndexes[name] = symbol.Index
  106. }
  107. }
  108. // remove duplicates from constants
  109. bytecode := c.Bytecode()
  110. bytecode.RemoveDuplicates()
  111. // check the constant objects limit
  112. if s.maxConstObjects >= 0 {
  113. cnt := bytecode.CountObjects()
  114. if cnt > s.maxConstObjects {
  115. return nil, fmt.Errorf("exceeding constant objects limit: %d", cnt)
  116. }
  117. }
  118. return &Compiled{
  119. globalIndexes: globalIndexes,
  120. bytecode: bytecode,
  121. globals: globals,
  122. maxAllocs: s.maxAllocs,
  123. }, nil
  124. }
  125. // Run compiles and runs the scripts. Use returned compiled object to access
  126. // global variables.
  127. func (s *Script) Run() (compiled *Compiled, err error) {
  128. compiled, err = s.Compile()
  129. if err != nil {
  130. return
  131. }
  132. err = compiled.Run()
  133. return
  134. }
  135. // RunContext is like Run but includes a context.
  136. func (s *Script) RunContext(
  137. ctx context.Context,
  138. ) (compiled *Compiled, err error) {
  139. compiled, err = s.Compile()
  140. if err != nil {
  141. return
  142. }
  143. err = compiled.RunContext(ctx)
  144. return
  145. }
  146. func (s *Script) prepCompile() (
  147. symbolTable *SymbolTable,
  148. globals []Object,
  149. err error,
  150. ) {
  151. var names []string
  152. for name := range s.variables {
  153. names = append(names, name)
  154. }
  155. symbolTable = NewSymbolTable()
  156. for idx, fn := range builtinFuncs {
  157. symbolTable.DefineBuiltin(idx, fn.Name)
  158. }
  159. globals = make([]Object, GlobalsSize)
  160. for idx, name := range names {
  161. symbol := symbolTable.Define(name)
  162. if symbol.Index != idx {
  163. panic(fmt.Errorf("wrong symbol index: %d != %d",
  164. idx, symbol.Index))
  165. }
  166. globals[symbol.Index] = s.variables[name].value
  167. }
  168. return
  169. }
  170. // Compiled is a compiled instance of the user script. Use Script.Compile() to
  171. // create Compiled object.
  172. type Compiled struct {
  173. globalIndexes map[string]int // global symbol name to index
  174. bytecode *Bytecode
  175. globals []Object
  176. maxAllocs int64
  177. lock sync.RWMutex
  178. }
  179. // Run executes the compiled script in the virtual machine.
  180. func (c *Compiled) Run() error {
  181. c.lock.Lock()
  182. defer c.lock.Unlock()
  183. v := NewVM(c.bytecode, c.globals, c.maxAllocs)
  184. return v.Run()
  185. }
  186. // RunContext is like Run but includes a context.
  187. func (c *Compiled) RunContext(ctx context.Context) (err error) {
  188. c.lock.Lock()
  189. defer c.lock.Unlock()
  190. v := NewVM(c.bytecode, c.globals, c.maxAllocs)
  191. ch := make(chan error, 1)
  192. go func() {
  193. defer func() {
  194. if r := recover(); r != nil {
  195. switch e := r.(type) {
  196. case string:
  197. ch <- fmt.Errorf(e)
  198. case error:
  199. ch <- e
  200. default:
  201. ch <- fmt.Errorf("unknown panic: %v", e)
  202. }
  203. }
  204. }()
  205. ch <- v.Run()
  206. }()
  207. select {
  208. case <-ctx.Done():
  209. v.Abort()
  210. <-ch
  211. err = ctx.Err()
  212. case err = <-ch:
  213. }
  214. return
  215. }
  216. // Clone creates a new copy of Compiled. Cloned copies are safe for concurrent
  217. // use by multiple goroutines.
  218. func (c *Compiled) Clone() *Compiled {
  219. c.lock.Lock()
  220. defer c.lock.Unlock()
  221. clone := &Compiled{
  222. globalIndexes: c.globalIndexes,
  223. bytecode: c.bytecode,
  224. globals: make([]Object, len(c.globals)),
  225. maxAllocs: c.maxAllocs,
  226. }
  227. // copy global objects
  228. for idx, g := range c.globals {
  229. if g != nil {
  230. clone.globals[idx] = g
  231. }
  232. }
  233. return clone
  234. }
  235. // IsDefined returns true if the variable name is defined (has value) before or
  236. // after the execution.
  237. func (c *Compiled) IsDefined(name string) bool {
  238. c.lock.RLock()
  239. defer c.lock.RUnlock()
  240. idx, ok := c.globalIndexes[name]
  241. if !ok {
  242. return false
  243. }
  244. v := c.globals[idx]
  245. if v == nil {
  246. return false
  247. }
  248. return v != UndefinedValue
  249. }
  250. // Get returns a variable identified by the name.
  251. func (c *Compiled) Get(name string) *Variable {
  252. c.lock.RLock()
  253. defer c.lock.RUnlock()
  254. value := UndefinedValue
  255. if idx, ok := c.globalIndexes[name]; ok {
  256. value = c.globals[idx]
  257. if value == nil {
  258. value = UndefinedValue
  259. }
  260. }
  261. return &Variable{
  262. name: name,
  263. value: value,
  264. }
  265. }
  266. // GetAll returns all the variables that are defined by the compiled script.
  267. func (c *Compiled) GetAll() []*Variable {
  268. c.lock.RLock()
  269. defer c.lock.RUnlock()
  270. var vars []*Variable
  271. for name, idx := range c.globalIndexes {
  272. value := c.globals[idx]
  273. if value == nil {
  274. value = UndefinedValue
  275. }
  276. vars = append(vars, &Variable{
  277. name: name,
  278. value: value,
  279. })
  280. }
  281. return vars
  282. }
  283. // Set replaces the value of a global variable identified by the name. An error
  284. // will be returned if the name was not defined during compilation.
  285. func (c *Compiled) Set(name string, value interface{}) error {
  286. c.lock.Lock()
  287. defer c.lock.Unlock()
  288. obj, err := FromInterface(value)
  289. if err != nil {
  290. return err
  291. }
  292. idx, ok := c.globalIndexes[name]
  293. if !ok {
  294. return fmt.Errorf("'%s' is not defined", name)
  295. }
  296. c.globals[idx] = obj
  297. return nil
  298. }