Golang lua.CreateLuaSandbox函數代碼示例

// Determines the script type and creates interpreter
func (this *SandboxFilter) Init(config interface{}) (err error) {
	if this.sb != nil {
		return nil // no-op already initialized
	}
	this.sbc = config.(*SandboxConfig)
	this.sbc.ScriptFilename = pipeline.PrependShareDir(this.sbc.ScriptFilename)

	data_dir := pipeline.PrependBaseDir(DATA_DIR)
	if !fileExists(data_dir) {
		err = os.MkdirAll(data_dir, 0700)
		if err != nil {
			return
		}
	}

	switch this.sbc.ScriptType {
	case "lua":
		this.sb, err = lua.CreateLuaSandbox(this.sbc)
		if err != nil {
			return
		}
	default:
		return fmt.Errorf("unsupported script type: %s", this.sbc.ScriptType)
	}

	this.preservationFile = filepath.Join(data_dir, this.name+DATA_EXT)
	if this.sbc.PreserveData && fileExists(this.preservationFile) {
		err = this.sb.Init(this.preservationFile, "filter")
	} else {
		err = this.sb.Init("", "filter")
	}

	return
}

func TestReadMessage(t *testing.T) {
	var sbc SandboxConfig
	captures := map[string]string{"exists": "found"}
	sbc.ScriptFilename = "./testsupport/read_message.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	msg := getTestMessage()
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	r := sb.ProcessMessage(msg, captures)
	if r != 0 {
		t.Errorf("ProcessMessage should return 0, received %d", r)
	}
	r = sb.TimerEvent(time.Now().UnixNano())
	if r != 0 {
		t.Errorf("read_message should return nil in timer_event")
	}
	sb.Destroy("")
}

func TestInit(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/hello_world.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	if STATUS_UNKNOWN != sb.Status() {
		t.Errorf("status should be %d, received %d",
			STATUS_UNKNOWN, sb.Status())
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	b := sb.Usage(TYPE_MEMORY, STAT_CURRENT)
	if b == 0 {
		t.Errorf("current memory should be >0, using %d", b)
	}
	b = sb.Usage(TYPE_MEMORY, STAT_MAXIMUM)
	if b == 0 {
		t.Errorf("maximum memory should be >0, using %d", b)
	}
	b = sb.Usage(TYPE_MEMORY, STAT_LIMIT)
	if b != sbc.MemoryLimit {
		t.Errorf("memory limit should be %d, using %d", sbc.MemoryLimit, b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_CURRENT)
	if b != 9 {
		t.Errorf("current instructions should be 9, using %d", b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_MAXIMUM)
	if b != 9 {
		t.Errorf("maximum instructions should be 9, using %d", b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_LIMIT)
	if b != sbc.InstructionLimit {
		t.Errorf("instruction limit should be %d, using %d", sbc.InstructionLimit, b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_CURRENT)
	if b != 12 {
		t.Errorf("current output should be 12, using %d", b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_MAXIMUM)
	if b != 12 {
		t.Errorf("maximum output should be 12, using %d", b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_LIMIT)
	if b != sbc.OutputLimit {
		t.Errorf("output limit should be %d, using %d", sbc.OutputLimit, b)
	}
	if STATUS_RUNNING != sb.Status() {
		t.Errorf("status should be %d, received %d",
			STATUS_RUNNING, sb.Status())
	}
	sb.Destroy("")
}

// Determines the script type and creates interpreter sandbox.
func (this *SandboxFilter) Init(config interface{}) (err error) {
	if this.sb != nil {
		return nil // no-op already initialized
	}
	this.sbc = config.(*sandbox.SandboxConfig)
	this.sbc.ScriptFilename = GetHekaConfigDir(this.sbc.ScriptFilename)

	switch this.sbc.ScriptType {
	case "lua":
		this.sb, err = lua.CreateLuaSandbox(this.sbc)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("unsupported script type: %s", this.sbc.ScriptType)
	}

	this.preservationFile = filepath.Join(filepath.Dir(this.sbc.ScriptFilename), this.name+".data")
	if this.sbc.PreserveData && fileExists(this.preservationFile) {
		err = this.sb.Init(this.preservationFile)
	} else {
		err = this.sb.Init("")
	}

	return err
}

func TestMissingTimeEvent(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/hello_world.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	r := sb.TimerEvent(time.Now().UnixNano())
	if r == 0 {
		t.Errorf("TimerEvent() expected: 1, received: %d", r)
	}
	if STATUS_TERMINATED != sb.Status() {
		t.Errorf("status should be %d, received %d",
			STATUS_TERMINATED, sb.Status())
	}
	r = sb.TimerEvent(time.Now().UnixNano()) // try to use the terminated plugin
	if r == 0 {
		t.Errorf("TimerEvent() expected: 1, received: %d", r)
	}
	sb.Destroy("")
}

func (s *SandboxEncoder) Init(config interface{}) (err error) {
	conf := config.(*SandboxEncoderConfig)
	s.sbc = &sandbox.SandboxConfig{
		ScriptType:       conf.ScriptType,
		ScriptFilename:   conf.ScriptFilename,
		ModuleDirectory:  conf.ModuleDirectory,
		PreserveData:     conf.PreserveData,
		MemoryLimit:      conf.MemoryLimit,
		InstructionLimit: conf.InstructionLimit,
		OutputLimit:      conf.OutputLimit,
		Profile:          conf.Profile,
		Config:           conf.Config,
	}
	s.sbc.ScriptFilename = pipeline.PrependShareDir(s.sbc.ScriptFilename)
	s.sampleDenominator = pipeline.Globals().SampleDenominator

	s.tz = time.UTC
	if tz, ok := s.sbc.Config["tz"]; ok {
		if s.tz, err = time.LoadLocation(tz.(string)); err != nil {
			return
		}
	}

	dataDir := pipeline.PrependBaseDir(sandbox.DATA_DIR)
	if !fileExists(dataDir) {
		if err = os.MkdirAll(dataDir, 0700); err != nil {
			return
		}
	}

	switch s.sbc.ScriptType {
	case "lua":
		s.sb, err = lua.CreateLuaSandbox(s.sbc)
	default:
		return fmt.Errorf("Unsupported script type: %s", s.sbc.ScriptType)
	}

	if err != nil {
		return fmt.Errorf("Sandbox creation failed: '%s'", err)
	}

	s.preservationFile = filepath.Join(dataDir, s.name+sandbox.DATA_EXT)
	if s.sbc.PreserveData && fileExists(s.preservationFile) {
		err = s.sb.Init(s.preservationFile, "encoder")
	} else {
		err = s.sb.Init("", "encoder")
	}
	if err != nil {
		return fmt.Errorf("Sandbox initialization failed: %s", err)
	}

	s.sb.InjectMessage(func(payload, payload_type, payload_name string) int {
		s.injected = true
		s.output = []byte(payload)
		return 0
	})
	s.sample = true
	s.cEncoder = client.NewProtobufEncoder(nil)
	return
}

func TestPreserve(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/serialize.lua"
	sbc.MemoryLimit = 64000
	sbc.InstructionLimit = 1000
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	output := "/tmp/serialize.lua.data"
	saved := "./testsupport/serialize.lua.data"
	err = sb.Destroy("/tmp/serialize.lua.data")
	if err != nil {
		t.Errorf("%s", err)
	} else {
		o, err := ioutil.ReadFile(output)
		if err != nil {
			t.Errorf("%s", err)
		}
		s, err := ioutil.ReadFile(saved)
		if err != nil {
			t.Errorf("%s", err)
		}
		if 0 != bytes.Compare(o, s) {
			t.Errorf("The preserved data does not match")
		}
	}
}

func TestMissingProcessMessage(t *testing.T) {
	var sbc SandboxConfig
	var captures map[string]string
	sbc.ScriptFilename = "./testsupport/hello_world.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	msg := getTestMessage()
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	r := sb.ProcessMessage(msg, captures)
	if r == 0 {
		t.Errorf("ProcessMessage() expected: 1, received: %d", r)
	}
	s := "process_message() function was not found"
	if sb.LastError() != s {
		t.Errorf("LastError() should be \"%s\", received: \"%s\"", s, sb.LastError())
	}
	if STATUS_TERMINATED != sb.Status() {
		t.Errorf("status should be %d, received %d",
			STATUS_TERMINATED, sb.Status())
	}
	r = sb.ProcessMessage(msg, captures) // try to use the terminated plugin
	if r == 0 {
		t.Errorf("ProcessMessage() expected: 1, received: %d", r)
	}
	sb.Destroy("")
}

func TestRestore(t *testing.T) {
	var sbc SandboxConfig
	var captures map[string]string
	sbc.ScriptFilename = "./testsupport/simple_count.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	msg := getTestMessage()
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("./testsupport/simple_count.lua.data")
	if err != nil {
		t.Errorf("%s", err)
	}
	sb.InjectMessage(func(p, pt, pn string) int {
		if p != "11" {
			t.Errorf("State was not restored")
		}
		return 0
	})
	r := sb.ProcessMessage(msg, captures)
	if r != 0 {
		t.Errorf("ProcessMessage should return 0, received %d", r)
	}
	sb.Destroy("")
}

func TestPreserveFailure(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/serialize_failure.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	output := "/tmp/serialize_failure.lua.data"
	err = sb.Destroy(output)
	if err == nil {
		t.Errorf("The key of type 'function' should have failed")
	} else {
		expect := "Destroy() serialize_data cannot preserve type 'function'"
		if err.Error() != expect {
			t.Errorf("expected '%s' got '%s'", expect, err)
		}
	}
	_, err = os.Stat(output)
	if err == nil {
		t.Errorf("The output file should be removed on failure")
	}
}

func TestFailedMessageInjection(t *testing.T) {
	var sbc SandboxConfig
	var captures map[string]string
	sbc.ScriptFilename = "./testsupport/loop.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	msg := getTestMessage()
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	err = sb.Init("")
	if err != nil {
		t.Errorf("%s", err)
	}
	sb.InjectMessage(func(p, pt, pn string) int {
		return 1
	})
	r := sb.ProcessMessage(msg, captures)
	if r != 1 {
		t.Errorf("ProcessMessage should return 1, received %d", r)
	}
	if STATUS_TERMINATED != sb.Status() {
		t.Errorf("status should be %d, received %d",
			STATUS_TERMINATED, sb.Status())
	}
	s := sb.LastError()
	errMsg := "process_message() inject_message() exceeded MaxMsgLoops"
	if s != errMsg {
		t.Errorf("error should be \"%s\", received \"%s\"", errMsg, s)
	}
	sb.Destroy("")
}

func TestAPIErrors(t *testing.T) {
	msg := getTestMessage()
	tests := []string{
		"inject_message() incorrect number of args",
		"output() no arg",
		"out of memory",
		"out of instructions",
		"operation on a nil",
		"invalid return",
		"no return",
		"read_message() incorrect number of args",
		"read_message() incorrect field name type",
		"read_message() negative field index",
		"read_message() negative array index",
		"output limit exceeded"}

	msgs := []string{
		"process_message() ./testsupport/errors.lua:11: inject_message() takes a maximum of 2 arguments",
		"process_message() ./testsupport/errors.lua:13: output() must have at least one argument",
		"process_message() not enough memory",
		"process_message() instruction_limit exceeded",
		"process_message() ./testsupport/errors.lua:22: attempt to perform arithmetic on global 'x' (a nil value)",
		"process_message() must return a single numeric value",
		"process_message() must return a single numeric value",
		"process_message() ./testsupport/errors.lua:28: read_message() incorrect number of arguments",
		"process_message() ./testsupport/errors.lua:30: bad argument #1 to 'read_message' (string expected, got nil)",
		"process_message() ./testsupport/errors.lua:32: bad argument #2 to 'read_message' (field index must be >= 0)",
		"process_message() ./testsupport/errors.lua:34: bad argument #3 to 'read_message' (array index must be >= 0)",
		"process_message() ./testsupport/errors.lua:37: output_limit exceeded"}

	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/errors.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 128
	for i, v := range tests {
		sb, err := lua.CreateLuaSandbox(&sbc)
		if err != nil {
			t.Errorf("%s", err)
		}
		err = sb.Init("")
		if err != nil {
			t.Errorf("%s", err)
		}
		msg.SetPayload(v)
		r := sb.ProcessMessage(msg)
		if r != 1 || STATUS_TERMINATED != sb.Status() {
			t.Errorf("test: %s status should be %d, received %d",
				v, STATUS_TERMINATED, sb.Status())
		}
		s := sb.LastError()
		if s != msgs[i] {
			t.Errorf("test: %s error should be \"%s\", received \"%s\"",
				v, msgs[i], s)
		}
		sb.Destroy("")
	}
}

func TestCreation(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/hello_world.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	sbc.OutputLimit = 1024
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err != nil {
		t.Errorf("%s", err)
	}
	b := sb.Usage(TYPE_MEMORY, STAT_CURRENT)
	if b != 0 {
		t.Errorf("current memory should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_MEMORY, STAT_MAXIMUM)
	if b != 0 {
		t.Errorf("maximum memory should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_MEMORY, STAT_LIMIT)
	if b != sbc.MemoryLimit {
		t.Errorf("memory limit should be %d, using %d", sbc.MemoryLimit, b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_CURRENT)
	if b != 0 {
		t.Errorf("current instructions should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_MAXIMUM)
	if b != 0 {
		t.Errorf("maximum instructions should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_INSTRUCTIONS, STAT_LIMIT)
	if b != sbc.InstructionLimit {
		t.Errorf("instruction limit should be %d, using %d", sbc.InstructionLimit, b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_CURRENT)
	if b != 0 {
		t.Errorf("current output should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_MAXIMUM)
	if b != 0 {
		t.Errorf("maximum output should be 0, using %d", b)
	}
	b = sb.Usage(TYPE_OUTPUT, STAT_LIMIT)
	if b != sbc.OutputLimit {
		t.Errorf("output limit should be %d, using %d", sbc.OutputLimit, b)
	}
	b = sb.Usage(TYPE_OUTPUT, 99)
	if b != 0 {
		t.Errorf("invalid index should return 0, received %d", b)
	}
	if sb.LastError() != "" {
		t.Errorf("LastError() should be empty, received: %s", sb.LastError())
	}
	sb.Destroy("")
}

func TestCreationTooManyInstructions(t *testing.T) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/hello_world.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000001
	sb, err := lua.CreateLuaSandbox(&sbc)
	if err == nil {
		t.Errorf("Sandbox creation should have failed on InstructionLimit")
		sb.Destroy("")
	}
}

func BenchmarkSandboxCreateInitDestroyPreserve(b *testing.B) {
	var sbc SandboxConfig
	sbc.ScriptFilename = "./testsupport/serialize.lua"
	sbc.MemoryLimit = 32767
	sbc.InstructionLimit = 1000
	for i := 0; i < b.N; i++ {
		sb, _ := lua.CreateLuaSandbox(&sbc)
		sb.Init("")
		sb.Destroy("/tmp/serialize.lua.data")
	}
}