Release v0.1.0: now with starlark script support.
You may send SIGHUP signal to the process to make it reread the script. It will use older loaded version if parsing of new file failed.
This commit is contained in:
154
mapping/elf/core.go
Normal file
154
mapping/elf/core.go
Normal file
@ -0,0 +1,154 @@
|
||||
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{}{},
|
||||
}
|
||||
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 ""
|
||||
}
|
Reference in New Issue
Block a user