github.com/undoio/delve@v1.9.0/pkg/terminal/starbind/starlark.go (about)

     1  package starbind
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"runtime"
     9  	"strings"
    10  	"sync"
    11  
    12  	"go.starlark.net/resolve"
    13  	"go.starlark.net/starlark"
    14  
    15  	"github.com/undoio/delve/service"
    16  	"github.com/undoio/delve/service/api"
    17  )
    18  
    19  //go:generate go run ../../../_scripts/gen-starlark-bindings.go go ./starlark_mapping.go
    20  //go:generate go run ../../../_scripts/gen-starlark-bindings.go doc ../../../Documentation/cli/starlark.md
    21  
    22  const (
    23  	dlvCommandBuiltinName        = "dlv_command"
    24  	readFileBuiltinName          = "read_file"
    25  	writeFileBuiltinName         = "write_file"
    26  	commandPrefix                = "command_"
    27  	dlvContextName               = "dlv_context"
    28  	curScopeBuiltinName          = "cur_scope"
    29  	defaultLoadConfigBuiltinName = "default_load_config"
    30  )
    31  
    32  func init() {
    33  	resolve.AllowNestedDef = true
    34  	resolve.AllowLambda = true
    35  	resolve.AllowFloat = true
    36  	resolve.AllowSet = true
    37  	resolve.AllowBitwise = true
    38  	resolve.AllowRecursion = true
    39  	resolve.AllowGlobalReassign = true
    40  }
    41  
    42  // Context is the context in which starlark scripts are evaluated.
    43  // It contains methods to call API functions, command line commands, etc.
    44  type Context interface {
    45  	Client() service.Client
    46  	RegisterCommand(name, helpMsg string, cmdfn func(args string) error)
    47  	CallCommand(cmdstr string) error
    48  	Scope() api.EvalScope
    49  	LoadConfig() api.LoadConfig
    50  }
    51  
    52  // Env is the environment used to evaluate starlark scripts.
    53  type Env struct {
    54  	env       starlark.StringDict
    55  	contextMu sync.Mutex
    56  	thread    *starlark.Thread
    57  	cancelfn  context.CancelFunc
    58  
    59  	ctx Context
    60  	out EchoWriter
    61  }
    62  
    63  // New creates a new starlark binding environment.
    64  func New(ctx Context, out EchoWriter) *Env {
    65  	env := &Env{}
    66  
    67  	env.ctx = ctx
    68  	env.out = out
    69  
    70  	env.env = env.starlarkPredeclare()
    71  	env.env[dlvCommandBuiltinName] = starlark.NewBuiltin(dlvCommandBuiltinName, func(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    72  		if err := isCancelled(thread); err != nil {
    73  			return starlark.None, err
    74  		}
    75  		argstrs := make([]string, len(args))
    76  		for i := range args {
    77  			a, ok := args[i].(starlark.String)
    78  			if !ok {
    79  				return nil, fmt.Errorf("argument of dlv_command is not a string")
    80  			}
    81  			argstrs[i] = string(a)
    82  		}
    83  		err := env.ctx.CallCommand(strings.Join(argstrs, " "))
    84  		if err != nil && strings.Contains(err.Error(), " has exited with status ") {
    85  			return env.interfaceToStarlarkValue(err), nil
    86  		}
    87  		return starlark.None, decorateError(thread, err)
    88  	})
    89  	env.env[readFileBuiltinName] = starlark.NewBuiltin(readFileBuiltinName, func(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    90  		if len(args) != 1 {
    91  			return nil, decorateError(thread, fmt.Errorf("wrong number of arguments"))
    92  		}
    93  		path, ok := args[0].(starlark.String)
    94  		if !ok {
    95  			return nil, decorateError(thread, fmt.Errorf("argument of read_file was not a string"))
    96  		}
    97  		buf, err := ioutil.ReadFile(string(path))
    98  		if err != nil {
    99  			return nil, decorateError(thread, err)
   100  		}
   101  		return starlark.String(string(buf)), nil
   102  	})
   103  	env.env[writeFileBuiltinName] = starlark.NewBuiltin(writeFileBuiltinName, func(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   104  		if len(args) != 2 {
   105  			return nil, decorateError(thread, fmt.Errorf("wrong number of arguments"))
   106  		}
   107  		path, ok := args[0].(starlark.String)
   108  		if !ok {
   109  			return nil, decorateError(thread, fmt.Errorf("first argument of write_file was not a string"))
   110  		}
   111  		err := ioutil.WriteFile(string(path), []byte(args[1].String()), 0640)
   112  		return starlark.None, decorateError(thread, err)
   113  	})
   114  	env.env[curScopeBuiltinName] = starlark.NewBuiltin(curScopeBuiltinName, func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   115  		return env.interfaceToStarlarkValue(env.ctx.Scope()), nil
   116  	})
   117  	env.env[defaultLoadConfigBuiltinName] = starlark.NewBuiltin(defaultLoadConfigBuiltinName, func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
   118  		return env.interfaceToStarlarkValue(env.ctx.LoadConfig()), nil
   119  	})
   120  	return env
   121  }
   122  
   123  // Redirect redirects starlark output to out.
   124  func (env *Env) Redirect(out EchoWriter) {
   125  	env.out = out
   126  	if env.thread != nil {
   127  		env.thread.Print = env.printFunc()
   128  	}
   129  }
   130  
   131  func (env *Env) printFunc() func(_ *starlark.Thread, msg string) {
   132  	return func(_ *starlark.Thread, msg string) { fmt.Fprintln(env.out, msg) }
   133  }
   134  
   135  // Execute executes a script. Path is the name of the file to execute and
   136  // source is the source code to execute.
   137  // Source can be either a []byte, a string or a io.Reader. If source is nil
   138  // Execute will execute the file specified by 'path'.
   139  // After the file is executed if a function named mainFnName exists it will be called, passing args to it.
   140  func (env *Env) Execute(path string, source interface{}, mainFnName string, args []interface{}) (starlark.Value, error) {
   141  	defer func() {
   142  		err := recover()
   143  		if err == nil {
   144  			return
   145  		}
   146  		fmt.Fprintf(env.out, "panic executing starlark script: %v\n", err)
   147  		for i := 0; ; i++ {
   148  			pc, file, line, ok := runtime.Caller(i)
   149  			if !ok {
   150  				break
   151  			}
   152  			fname := "<unknown>"
   153  			fn := runtime.FuncForPC(pc)
   154  			if fn != nil {
   155  				fname = fn.Name()
   156  			}
   157  			fmt.Fprintf(env.out, "%s\n\tin %s:%d\n", fname, file, line)
   158  		}
   159  	}()
   160  
   161  	thread := env.newThread()
   162  	globals, err := starlark.ExecFile(thread, path, source, env.env)
   163  	if err != nil {
   164  		return starlark.None, err
   165  	}
   166  
   167  	err = env.exportGlobals(globals)
   168  	if err != nil {
   169  		return starlark.None, err
   170  	}
   171  
   172  	return env.callMain(thread, globals, mainFnName, args)
   173  }
   174  
   175  // exportGlobals saves globals with a name starting with a capital letter
   176  // into the environment and creates commands from globals with a name
   177  // starting with "command_"
   178  func (env *Env) exportGlobals(globals starlark.StringDict) error {
   179  	for name, val := range globals {
   180  		switch {
   181  		case strings.HasPrefix(name, commandPrefix):
   182  			err := env.createCommand(name, val)
   183  			if err != nil {
   184  				return err
   185  			}
   186  		case name[0] >= 'A' && name[0] <= 'Z':
   187  			env.env[name] = val
   188  		}
   189  	}
   190  	return nil
   191  }
   192  
   193  // Cancel cancels the execution of a currently running script or function.
   194  func (env *Env) Cancel() {
   195  	if env == nil {
   196  		return
   197  	}
   198  	env.contextMu.Lock()
   199  	if env.cancelfn != nil {
   200  		env.cancelfn()
   201  		env.cancelfn = nil
   202  	}
   203  	if env.thread != nil {
   204  		env.thread.Cancel("user interrupt")
   205  	}
   206  	env.contextMu.Unlock()
   207  }
   208  
   209  func (env *Env) newThread() *starlark.Thread {
   210  	thread := &starlark.Thread{
   211  		Print: env.printFunc(),
   212  	}
   213  	env.contextMu.Lock()
   214  	var ctx context.Context
   215  	ctx, env.cancelfn = context.WithCancel(context.Background())
   216  	env.thread = thread
   217  	env.contextMu.Unlock()
   218  	thread.SetLocal(dlvContextName, ctx)
   219  	return thread
   220  }
   221  
   222  func (env *Env) createCommand(name string, val starlark.Value) error {
   223  	fnval, ok := val.(*starlark.Function)
   224  	if !ok {
   225  		return nil
   226  	}
   227  
   228  	name = name[len(commandPrefix):]
   229  
   230  	helpMsg := fnval.Doc()
   231  	if helpMsg == "" {
   232  		helpMsg = "user defined"
   233  	}
   234  
   235  	if fnval.NumParams() == 1 {
   236  		if p0, _ := fnval.Param(0); p0 == "args" {
   237  			env.ctx.RegisterCommand(name, helpMsg, func(args string) error {
   238  				_, err := starlark.Call(env.newThread(), fnval, starlark.Tuple{starlark.String(args)}, nil)
   239  				return err
   240  			})
   241  			return nil
   242  		}
   243  	}
   244  
   245  	env.ctx.RegisterCommand(name, helpMsg, func(args string) error {
   246  		thread := env.newThread()
   247  		argval, err := starlark.Eval(thread, "<input>", "("+args+")", env.env)
   248  		if err != nil {
   249  			return err
   250  		}
   251  		argtuple, ok := argval.(starlark.Tuple)
   252  		if !ok {
   253  			argtuple = starlark.Tuple{argval}
   254  		}
   255  		_, err = starlark.Call(thread, fnval, argtuple, nil)
   256  		return err
   257  	})
   258  	return nil
   259  }
   260  
   261  // callMain calls the main function in globals, if one was defined.
   262  func (env *Env) callMain(thread *starlark.Thread, globals starlark.StringDict, mainFnName string, args []interface{}) (starlark.Value, error) {
   263  	if mainFnName == "" {
   264  		return starlark.None, nil
   265  	}
   266  	mainval := globals[mainFnName]
   267  	if mainval == nil {
   268  		return starlark.None, nil
   269  	}
   270  	mainfn, ok := mainval.(*starlark.Function)
   271  	if !ok {
   272  		return starlark.None, fmt.Errorf("%s is not a function", mainFnName)
   273  	}
   274  	if mainfn.NumParams() != len(args) {
   275  		return starlark.None, fmt.Errorf("wrong number of arguments for %s", mainFnName)
   276  	}
   277  	argtuple := make(starlark.Tuple, len(args))
   278  	for i := range args {
   279  		argtuple[i] = env.interfaceToStarlarkValue(args[i])
   280  	}
   281  	return starlark.Call(thread, mainfn, argtuple, nil)
   282  }
   283  
   284  func isCancelled(thread *starlark.Thread) error {
   285  	if ctx, ok := thread.Local(dlvContextName).(context.Context); ok {
   286  		select {
   287  		case <-ctx.Done():
   288  			return ctx.Err()
   289  		default:
   290  		}
   291  	}
   292  	return nil
   293  }
   294  
   295  func decorateError(thread *starlark.Thread, err error) error {
   296  	if err == nil {
   297  		return nil
   298  	}
   299  	pos := thread.CallFrame(1).Pos
   300  	if pos.Col > 0 {
   301  		return fmt.Errorf("%s:%d:%d: %v", pos.Filename(), pos.Line, pos.Col, err)
   302  	}
   303  	return fmt.Errorf("%s:%d: %v", pos.Filename(), pos.Line, err)
   304  }
   305  
   306  type EchoWriter interface {
   307  	io.Writer
   308  	Echo(string)
   309  	Flush()
   310  }