package elf import ( "flag" "fmt" "log" "os" "os/signal" "sync" "sync/atomic" "syscall" "go.starlark.net/starlark" ) type threadPool struct { pool []*starlark.Thread mutex *sync.Mutex } var ( scriptFile = flag.String("mapping-script", "map.starlark", "Starlark script used to do mapping logic") threadPoolSize = flag.Int("script-thread-pool", 128, "Thread pool size for starlark execution engine") builtinFunc = make(starlark.StringDict) requiredFunc = map[string]struct{}{ "getMapping": struct{}{}, "getNodes": struct{}{}, "getSuffix": struct{}{}, } parsedFunc *starlark.StringDict parseThread = &starlark.Thread{ Name: "parseThread", } reloadSignal = make(chan os.Signal) pool *threadPool counter uint64 ) func newThreadPool(size int) *threadPool { ret := &threadPool{ pool: make([]*starlark.Thread, size), mutex: &sync.Mutex{}, } for i := range ret.pool { counter += 1 ret.pool[i] = &starlark.Thread{ Name: fmt.Sprintf("ElfThread-%d", counter), } } return ret } func (pool *threadPool) Get() (ret *starlark.Thread) { pool.mutex.Lock() defer pool.mutex.Unlock() if len(pool.pool) == 0 { return &starlark.Thread{ Name: fmt.Sprintf("ElfThread-%d", atomic.AddUint64(&counter, 1)), } } ret = pool.pool[len(pool.pool)-1] pool.pool = pool.pool[:len(pool.pool)-1] return } func (pool *threadPool) Put(thread *starlark.Thread) { pool.mutex.Lock() defer pool.mutex.Unlock() if len(pool.pool) == *threadPoolSize { return } pool.pool = append(pool.pool, thread) } func RegisterFunc(name string, impl func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error)) { if builtinFunc.Has(name) { panic(fmt.Errorf("Function %s has already been declared as: %s", name, builtinFunc[name].String())) } builtinFunc[name] = starlark.NewBuiltin(name, impl) log.Println("Registered function", name, "to starlark") } func RequireFunc(name string) { requiredFunc[name] = struct{}{} } func Initialize() { globals, err := starlark.ExecFile( parseThread, *scriptFile, nil, builtinFunc) if err != nil { panic(fmt.Errorf("Unable to parse starlark file: %s", err.Error())) } parsedFunc = &globals for name := range requiredFunc { if !parsedFunc.Has(name) { panic(fmt.Errorf("Required function %s not provided", name)) } if _, ok := globals[name].(starlark.Callable); !ok { panic(fmt.Errorf( "Variable %s is not callable as required function", name)) } } signal.Notify(reloadSignal, syscall.SIGHUP) go func() { for { <-reloadSignal reload() } }() pool = newThreadPool(*threadPoolSize) } func reload() { globals, err := starlark.ExecFile( parseThread, *scriptFile, nil, builtinFunc) if err != nil { log.Println("Reload failed: Unable to parse starlark file:", err.Error()) return } for name := range requiredFunc { if !globals.Has(name) { log.Println( "Reload failed: Required function", name, "not provided") return } if _, ok := globals[name].(starlark.Callable); !ok { log.Println( "Variable", name, "is not callable as required function") return } } log.Println("Reload success, new rules applied :)") parsedFunc = &globals } func GetMapping(ip string) string { thread := pool.Get() ret, err := starlark.Call(thread, (*parsedFunc)["getMapping"], starlark.Tuple{starlark.String(ip)}, nil) if err != nil { log.Println("Starlark execute error:", err.Error()) return "" } if r, ok := ret.(starlark.String); ok { return string(r) } log.Println("Script returned unexpected result:", ret.String()) return "" } func GetNodes() (res []string) { thread := pool.Get() ret, err := starlark.Call(thread, (*parsedFunc)["getNodes"], nil, nil) if err != nil { log.Println("Starlark execute error:", err.Error()) return nil } if r, ok := ret.(*starlark.List); ok { res = make([]string, r.Len()) for i := 0; i < r.Len(); i++ { v := r.Index(i) if s, ok := v.(starlark.String); ok { res[i] = string(s) } else { log.Println("Script returned unexpected result:", ret.String()) } } } else { log.Println("Script returned unexpected result:", ret.String()) } return } func GetSuffix(ip string) string { thread := pool.Get() ret, err := starlark.Call(thread, (*parsedFunc)["getSuffix"], starlark.Tuple{starlark.String(ip)}, nil) if err != nil { log.Println("Starlark execute error:", err.Error()) return "" } if r, ok := ret.(starlark.String); ok { return string(r) } log.Println("Script returned unexpected result:", ret.String()) return "" }