<commit_before>package driver

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutils "github.com/hashicorp/nomad/client/testutil"
)

func TestExecDriver_Fingerprint(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:      "foo",
		Driver:    "exec",
		Resources: structs.DefaultResources(),
	}
	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)
	node := &structs.Node{
		Attributes: map[string]string{
			"unique.cgroup.mountpoint": "/sys/fs/cgroup",
		},
	}
	apply, err := d.Fingerprint(&config.Config{}, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !apply {
		t.Fatalf("should apply")
	}
	if node.Attributes["driver.exec"] == "" {
		t.Fatalf("missing driver")
	}
}

func TestExecDriver_StartOpen_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"5"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Attempt to open
	handle2, err := d.Open(ctx.ExecCtx, handle.ID())
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle2 == nil {
		t.Fatalf("missing handle")
	}

	handle.Kill()
	handle2.Kill()
}
func TestExecDriver_Start_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"2"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestExecDriver_Start_Wait_AllocDir(t *testing.T) {
	ctestutils.ExecCompatible(t)

	exp := []byte{'w', 'i', 'n'}
	file := "output.txt"
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/bash",
			"args": []string{
				"-c",
				fmt.Sprintf(`sleep 1; echo -n %s > ${%s}/%s`, string(exp), env.AllocDir, file),
			},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}

	// Check that data was written to the shared alloc directory.
	outputFile := filepath.Join(ctx.AllocDir.SharedDir, file)
	act, err := ioutil.ReadFile(outputFile)
	if err != nil {
		t.Fatalf("Couldn't read expected output: %v", err)
	}

	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("Command outputted %v; want %v", act, exp)
	}
}

func TestExecDriver_Start_Kill_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"100"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources:   basicResources,
		KillTimeout: 10 * time.Second,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	go func() {
		time.Sleep(100 * time.Millisecond)
		err := handle.Kill()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if res.Successful() {
			t.Fatal("should err")
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*10) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestExecDriverUser(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		User:   "alice",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"100"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources:   basicResources,
		KillTimeout: 10 * time.Second,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err == nil {
		handle.Kill()
		t.Fatalf("Should've failed")
	}
	msg := "user alice"
	if !strings.Contains(err.Error(), msg) {
		t.Fatalf("Expecting '%v' in '%v'", msg, err)
	}
}

// TestExecDriver_HandlerExec ensures the exec driver's handle properly executes commands inside the chroot.
func TestExecDriver_HandlerExec(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"9000"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Exec a command that should work
	out, code, err := handle.Exec(context.TODO(), "/usr/bin/stat", []string{"/alloc"})
	if err != nil {
		t.Fatalf("error exec'ing stat: %v", err)
	}
	if code != 0 {
		t.Fatalf("expected `stat /alloc` to succeed but exit code was: %d", code)
	}
	if expected := 100; len(out) < expected {
		t.Fatalf("expected at least %d bytes of output but found %d:\n%s", expected, len(out), out)
	}

	// Exec a command that should fail
	out, code, err = handle.Exec(context.TODO(), "/usr/bin/stat", []string{"lkjhdsaflkjshowaisxmcvnlia"})
	if err != nil {
		t.Fatalf("error exec'ing stat: %v", err)
	}
	if code == 0 {
		t.Fatalf("expected `stat` to fail but exit code was: %d", code)
	}
	if expected := "No such file or directory"; !bytes.Contains(out, []byte(expected)) {
		t.Fatalf("expected output to contain %q but found: %q", expected, out)
	}

	if err := handle.Kill(); err != nil {
		t.Fatalf("error killing exec handle: %v", err)
	}
}
<commit_msg>Test env+cgroups for exec driver checks<commit_after>package driver

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutils "github.com/hashicorp/nomad/client/testutil"
)

func TestExecDriver_Fingerprint(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:      "foo",
		Driver:    "exec",
		Resources: structs.DefaultResources(),
	}
	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)
	node := &structs.Node{
		Attributes: map[string]string{
			"unique.cgroup.mountpoint": "/sys/fs/cgroup",
		},
	}
	apply, err := d.Fingerprint(&config.Config{}, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !apply {
		t.Fatalf("should apply")
	}
	if node.Attributes["driver.exec"] == "" {
		t.Fatalf("missing driver")
	}
}

func TestExecDriver_StartOpen_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"5"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Attempt to open
	handle2, err := d.Open(ctx.ExecCtx, handle.ID())
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle2 == nil {
		t.Fatalf("missing handle")
	}

	handle.Kill()
	handle2.Kill()
}
func TestExecDriver_Start_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"2"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Update should be a no-op
	err = handle.Update(task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestExecDriver_Start_Wait_AllocDir(t *testing.T) {
	ctestutils.ExecCompatible(t)

	exp := []byte{'w', 'i', 'n'}
	file := "output.txt"
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/bash",
			"args": []string{
				"-c",
				fmt.Sprintf(`sleep 1; echo -n %s > ${%s}/%s`, string(exp), env.AllocDir, file),
			},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if !res.Successful() {
			t.Fatalf("err: %v", res)
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*5) * time.Second):
		t.Fatalf("timeout")
	}

	// Check that data was written to the shared alloc directory.
	outputFile := filepath.Join(ctx.AllocDir.SharedDir, file)
	act, err := ioutil.ReadFile(outputFile)
	if err != nil {
		t.Fatalf("Couldn't read expected output: %v", err)
	}

	if !reflect.DeepEqual(act, exp) {
		t.Fatalf("Command outputted %v; want %v", act, exp)
	}
}

func TestExecDriver_Start_Kill_Wait(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"100"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources:   basicResources,
		KillTimeout: 10 * time.Second,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	go func() {
		time.Sleep(100 * time.Millisecond)
		err := handle.Kill()
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}()

	// Task should terminate quickly
	select {
	case res := <-handle.WaitCh():
		if res.Successful() {
			t.Fatal("should err")
		}
	case <-time.After(time.Duration(testutil.TestMultiplier()*10) * time.Second):
		t.Fatalf("timeout")
	}
}

func TestExecDriverUser(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		User:   "alice",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"100"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources:   basicResources,
		KillTimeout: 10 * time.Second,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err == nil {
		handle.Kill()
		t.Fatalf("Should've failed")
	}
	msg := "user alice"
	if !strings.Contains(err.Error(), msg) {
		t.Fatalf("Expecting '%v' in '%v'", msg, err)
	}
}

// TestExecDriver_HandlerExec ensures the exec driver's handle properly
// executes commands inside the container.
func TestExecDriver_HandlerExec(t *testing.T) {
	ctestutils.ExecCompatible(t)
	task := &structs.Task{
		Name:   "sleep",
		Driver: "exec",
		Config: map[string]interface{}{
			"command": "/bin/sleep",
			"args":    []string{"9000"},
		},
		LogConfig: &structs.LogConfig{
			MaxFiles:      10,
			MaxFileSizeMB: 10,
		},
		Resources: basicResources,
	}

	ctx := testDriverContexts(t, task)
	defer ctx.AllocDir.Destroy()
	d := NewExecDriver(ctx.DriverCtx)

	if _, err := d.Prestart(ctx.ExecCtx, task); err != nil {
		t.Fatalf("prestart err: %v", err)
	}
	handle, err := d.Start(ctx.ExecCtx, task)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if handle == nil {
		t.Fatalf("missing handle")
	}

	// Exec a command that should work and dump the environment
	out, code, err := handle.Exec(context.Background(), "/bin/sh", []string{"-c", "env | grep NOMAD"})
	if err != nil {
		t.Fatalf("error exec'ing stat: %v", err)
	}
	if code != 0 {
		t.Fatalf("expected `stat /alloc` to succeed but exit code was: %d", code)
	}

	// Assert exec'd commands are run in a task-like environment
	scriptEnv := make(map[string]string)
	for _, line := range strings.Split(string(out), "\n") {
		if line == "" {
			continue
		}
		parts := strings.SplitN(string(line), "=", 2)
		if len(parts) != 2 {
			t.Fatalf("Invalid env var: %q", line)
		}
		scriptEnv[parts[0]] = parts[1]
	}
	if v, ok := scriptEnv["NOMAD_SECRETS_DIR"]; !ok || v != "/secrets" {
		t.Errorf("Expected NOMAD_SECRETS_DIR=/secrets but found=%t value=%q", ok, v)
	}
	if v, ok := scriptEnv["NOMAD_ALLOC_ID"]; !ok || v != ctx.DriverCtx.allocID {
		t.Errorf("Expected NOMAD_SECRETS_DIR=%q but found=%t value=%q", ok, v)
	}

	// Assert cgroup membership
	out, code, err = handle.Exec(context.Background(), "/bin/cat", []string{"/proc/self/cgroup"})
	if err != nil {
		t.Fatalf("error exec'ing cat /proc/self/cgroup: %v", err)
	}
	if code != 0 {
		t.Fatalf("expected `cat /proc/self/cgroup` to succeed but exit code was: %d", code)
	}
	found := false
	for _, line := range strings.Split(string(out), "\n") {
		// Every cgroup entry should be /nomad/$ALLOC_ID
		if line == "" {
			continue
		}
		if !strings.Contains(line, ":/nomad/") {
			t.Errorf("Not a member of the alloc's cgroup: expected=...:/nomad/... -- found=%q", line)
			continue
		}
		found = true
	}
	if !found {
		t.Errorf("exec'd command isn't in the task's cgroup")
	}

	// Exec a command that should fail
	out, code, err = handle.Exec(context.Background(), "/usr/bin/stat", []string{"lkjhdsaflkjshowaisxmcvnlia"})
	if err != nil {
		t.Fatalf("error exec'ing stat: %v", err)
	}
	if code == 0 {
		t.Fatalf("expected `stat` to fail but exit code was: %d", code)
	}
	if expected := "No such file or directory"; !bytes.Contains(out, []byte(expected)) {
		t.Fatalf("expected output to contain %q but found: %q", expected, out)
	}

	if err := handle.Kill(); err != nil {
		t.Fatalf("error killing exec handle: %v", err)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"log"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutil "github.com/hashicorp/nomad/client/testutil"
)

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

type MockTaskStateUpdater struct {
	Count       int
	Name        []string
	Status      []string
	Description []string
}

func (m *MockTaskStateUpdater) Update(name, status, desc string) {
	m.Count += 1
	m.Name = append(m.Name, name)
	m.Status = append(m.Status, status)
	m.Description = append(m.Description, desc)
}

func testTaskRunner() (*MockTaskStateUpdater, *TaskRunner) {
	logger := testLogger()
	conf := DefaultConfig()
	conf.StateDir = "/tmp"
	upd := &MockTaskStateUpdater{}
	ctx := driver.NewExecContext()
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	tr := NewTaskRunner(logger, conf, upd.Update, ctx, alloc.ID, task)
	return upd, tr
}

func testTaskRunnerPostOffer() (*MockTaskStateUpdater, *TaskRunner) {
	state, runner := testTaskRunner()
	// testTaskRunner gives us a job in pre-offer state. We need it in post-
	// offer state some some other things are initialized. We'll init them here.
	runner.task.Resources.Networks[0].ReservedPorts[0] = 80
	return state, runner
}

func TestTaskRunner_SimpleRun(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunnerPostOffer()
	go tr.Run()
	defer tr.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Name[0] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[0] != "task started" {
		t.Fatalf("bad: %#v", upd.Description)
	}

	if upd.Name[1] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[1] != "task completed" {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Destroy(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()

	// Begin the tear down
	go func() {
		time.Sleep(100 * time.Millisecond)
		tr.Destroy()
	}()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if !strings.Contains(upd.Description[1], "task failed") {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Update(t *testing.T) {
	ctestutil.ExecCompatible(t)
	_, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Update the task definition
	newTask := new(structs.Task)
	*newTask = *tr.task
	newTask.Driver = "foobar"
	tr.Update(newTask)

	// Wait for update to take place
	testutil.WaitForResult(func() (bool, error) {
		return tr.task == newTask, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

/*
TODO: This test is disabled til a follow-up api changes the restore state interface.
The driver/executor interface will be changed from Open to Cleanup, in which
clean-up tears down previous allocs.

func TestTaskRunner_SaveRestoreState(t *testing.T) {
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Snapshot state
	time.Sleep(200 * time.Millisecond)
	err := tr.SaveState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a new task runner
	tr2 := NewTaskRunner(tr.logger, tr.config, upd.Update,
		tr.ctx, tr.allocID, &structs.Task{Name: tr.task.Name})
	err = tr2.RestoreState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	go tr2.Run()
	defer tr2.Destroy()

	// Destroy and wait
	tr2.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}
}
*/
<commit_msg>Initialize reserved ports<commit_after>package client

import (
	"log"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"

	ctestutil "github.com/hashicorp/nomad/client/testutil"
)

func testLogger() *log.Logger {
	return log.New(os.Stderr, "", log.LstdFlags)
}

type MockTaskStateUpdater struct {
	Count       int
	Name        []string
	Status      []string
	Description []string
}

func (m *MockTaskStateUpdater) Update(name, status, desc string) {
	m.Count += 1
	m.Name = append(m.Name, name)
	m.Status = append(m.Status, status)
	m.Description = append(m.Description, desc)
}

func testTaskRunner() (*MockTaskStateUpdater, *TaskRunner) {
	logger := testLogger()
	conf := DefaultConfig()
	conf.StateDir = "/tmp"
	upd := &MockTaskStateUpdater{}
	ctx := driver.NewExecContext()
	alloc := mock.Alloc()
	task := alloc.Job.TaskGroups[0].Tasks[0]
	tr := NewTaskRunner(logger, conf, upd.Update, ctx, alloc.ID, task)
	return upd, tr
}

func testTaskRunnerPostOffer() (*MockTaskStateUpdater, *TaskRunner) {
	state, runner := testTaskRunner()
	// testTaskRunner gives us a job in pre-offer state. We need it in post-
	// offer state some some other things are initialized. We'll init them here.
	runner.task.Resources.Networks[0].ReservedPorts = []int{80}
	return state, runner
}

func TestTaskRunner_SimpleRun(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunnerPostOffer()
	go tr.Run()
	defer tr.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Name[0] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[0] != "task started" {
		t.Fatalf("bad: %#v", upd.Description)
	}

	if upd.Name[1] != tr.task.Name {
		t.Fatalf("bad: %#v", upd.Name)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Description[1] != "task completed" {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Destroy(t *testing.T) {
	ctestutil.ExecCompatible(t)
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()

	// Begin the tear down
	go func() {
		time.Sleep(100 * time.Millisecond)
		tr.Destroy()
	}()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}

	if upd.Count != 2 {
		t.Fatalf("should have 2 updates: %#v", upd)
	}
	if upd.Status[0] != structs.AllocClientStatusRunning {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if upd.Status[1] != structs.AllocClientStatusDead {
		t.Fatalf("bad: %#v", upd.Status)
	}
	if !strings.Contains(upd.Description[1], "task failed") {
		t.Fatalf("bad: %#v", upd.Description)
	}
}

func TestTaskRunner_Update(t *testing.T) {
	ctestutil.ExecCompatible(t)
	_, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Update the task definition
	newTask := new(structs.Task)
	*newTask = *tr.task
	newTask.Driver = "foobar"
	tr.Update(newTask)

	// Wait for update to take place
	testutil.WaitForResult(func() (bool, error) {
		return tr.task == newTask, nil
	}, func(err error) {
		t.Fatalf("err: %v", err)
	})
}

/*
TODO: This test is disabled til a follow-up api changes the restore state interface.
The driver/executor interface will be changed from Open to Cleanup, in which
clean-up tears down previous allocs.

func TestTaskRunner_SaveRestoreState(t *testing.T) {
	upd, tr := testTaskRunner()

	// Change command to ensure we run for a bit
	tr.task.Config["command"] = "/bin/sleep"
	tr.task.Config["args"] = "10"
	go tr.Run()
	defer tr.Destroy()

	// Snapshot state
	time.Sleep(200 * time.Millisecond)
	err := tr.SaveState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a new task runner
	tr2 := NewTaskRunner(tr.logger, tr.config, upd.Update,
		tr.ctx, tr.allocID, &structs.Task{Name: tr.task.Name})
	err = tr2.RestoreState()
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	go tr2.Run()
	defer tr2.Destroy()

	// Destroy and wait
	tr2.Destroy()

	select {
	case <-tr.WaitCh():
	case <-time.After(2 * time.Second):
		t.Fatalf("timeout")
	}
}
*/
<|endoftext|>


<commit_before>package log

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
)

type (
	Level  uint32
	Logger bool
)

const (
	Trace = iota
	Debug
	Verbose
	Info
	Off
)

var lock sync.Mutex

var v Level = Info

func SetV(level Level) {
	atomic.StoreUint32((*uint32)(&v), uint32(level))
}

func V(level Level) Logger {
	if atomic.LoadUint32((*uint32)(&v)) > uint32(level) {
		return Logger(false)
	}
	return Logger(true)
}

func (l Logger) log(v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprint(os.Stderr, v...)
	}
}

func (l Logger) unsafeLog(v ...interface{}) {
	if l {
		fmt.Fprint(os.Stderr, v...)
	}
}

func (l Logger) logf(format string, v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprintf(os.Stderr, format, v...)
	}
}

func (l Logger) unsafeLogf(format string, v ...interface{}) {
	if l {
		fmt.Fprintf(os.Stderr, format, v...)
	}
}

func (l Logger) logln(v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprintln(os.Stderr, v...)
	}
}

func (l Logger) unsafeLogln(v ...interface{}) {
	if l {
		fmt.Fprintln(os.Stderr, v...)
	}
}

func (l Logger) Lock() {
	lock.Lock()
}

func (l Logger) Unlock() {
	lock.Unlock()
}

func (l Logger) Print(v ...interface{}) {
	l.log(v...)
}

func (l Logger) UnsafePrint(v ...interface{}) {
	l.unsafeLog(v...)
}

func (l Logger) Printf(format string, v ...interface{}) {
	l.logf(format, v...)
}

func (l Logger) UnsafePrintf(format string, v ...interface{}) {
	l.unsafeLogf(format, v...)
}

func (l Logger) Println(v ...interface{}) {
	l.logln(v...)
}

func (l Logger) UnsafePrintln(v ...interface{}) {
	l.unsafeLogln(v...)
}

func (l Logger) Fatal(v ...interface{}) {
	l.log(v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatal(v ...interface{}) {
	l.unsafeLog(v...)
	os.Exit(1)
}

func (l Logger) Fatalf(format string, v ...interface{}) {
	l.logf(format, v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatalf(format string, v ...interface{}) {
	l.unsafeLogf(format, v...)
	os.Exit(1)
}

func (l Logger) Fatalln(v ...interface{}) {
	l.logln(v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatalln(v ...interface{}) {
	l.unsafeLogln(v...)
	os.Exit(1)
}

func (l Logger) Run(msg string) {
	l.logf("[RUN]      %v\n", msg)
}

func (l Logger) UnsafeRun(msg string) {
	l.unsafeLogf("[RUN]      %v\n", msg)
}

func (l Logger) Skip(msg string) {
	l.logf("[SKIP]     %v\n", msg)
}

func (l Logger) UnsafeSkip(msg string) {
	l.unsafeLogf("[SKIP]     %v\n", msg)
}

func (l Logger) Warn(msg string) {
	l.logf("[WARN]     %v\n", msg)
}

func (l Logger) UnsafeWarn(msg string) {
	l.unsafeLogf("[WARN]     %v\n", msg)
}

func (l Logger) Go(msg string) {
	l.logf("[GO]       %v\n", msg)
}

func (l Logger) UnsafeGo(msg string) {
	l.unsafeLogf("[GO]       %v\n", msg)
}

func (l Logger) Ok(msg string) {
	l.logf("[OK]       %v\n", msg)
}

func (l Logger) UnsafeOk(msg string) {
	l.unsafeLogf("[OK]       %v\n", msg)
}

func (l Logger) Fail(msg string) {
	l.logf("[FAIL]     %v\n", msg)
}

func (l Logger) UnsafeFail(msg string) {
	l.unsafeLogf("[FAIL]     %v\n", msg)
}

func (l Logger) Rollback(msg string) {
	l.logf("[ROLLBACK] %v\n", msg)
}

func (l Logger) UnsafeRollback(msg string) {
	l.unsafeLogf("[ROLLBACK] %v\n", msg)
}

func (l Logger) NewLine(msg string) {
	l.logf("           %v\n", msg)
}

func (l Logger) UnsafeNewLine(msg string) {
	l.unsafeLogf("           %v\n", msg)
}

func (l Logger) Stderr(stderr *bytes.Buffer) {
	if stderr != nil && stderr.Len() != 0 {
		lock.Lock()
		defer lock.Unlock()
		l.UnsafePrintln("<<<<< stderr")
		l.UnsafePrint(stderr)
		l.UnsafePrintln(">>>>> stderr")
	}
}

func (l Logger) UnsafeStderr(stderr *bytes.Buffer) {
	if stderr != nil && stderr.Len() != 0 {
		l.UnsafePrintln("<<<<< stderr")
		l.UnsafePrint(stderr)
		l.UnsafePrintln(">>>>> stderr")
	}
}

func (l Logger) Die(msg string, err error) {
	lock.Lock()
	defer lock.Unlock()
	l.UnsafeFail(msg)
	l.UnsafeFatalln("\nError:", err)
}

func (l Logger) UnsafeDie(msg string, err error) {
	l.UnsafeFail(msg)
	l.UnsafeFatalln("\nError:", err)
}

func (l Logger) FailWithDetails(msg string, details *bytes.Buffer) {
	if msg != "" {
		l.Fail(msg)
	}
	l.Stderr(details)
}

func Print(v ...interface{}) {
	V(Info).Print(v...)
}

func UnsafePrint(v ...interface{}) {
	V(Info).UnsafePrint(v...)
}

func Printf(format string, v ...interface{}) {
	V(Info).Printf(format, v...)
}

func UnsafePrintf(format string, v ...interface{}) {
	V(Info).UnsafePrintf(format, v...)
}

func Println(v ...interface{}) {
	V(Info).Println(v...)
}

func UnsafePrintln(v ...interface{}) {
	V(Info).UnsafePrintln(v...)
}

func Fatal(v ...interface{}) {
	V(Info).Fatal(v...)
}

func UnsafeFatal(v ...interface{}) {
	V(Info).UnsafeFatal(v...)
}

func Fatalf(format string, v ...interface{}) {
	V(Info).Fatalf(format, v...)
}

func UnsafeFatalf(format string, v ...interface{}) {
	V(Info).UnsafeFatalf(format, v...)
}

func Fatalln(v ...interface{}) {
	V(Info).Fatalln(v...)
}

func UnsafeFatalln(v ...interface{}) {
	V(Info).UnsafeFatalln(v...)
}

func Run(msg string) {
	V(Info).Run(msg)
}

func UnsafeRun(msg string) {
	V(Info).UnsafeRun(msg)
}

func Skip(msg string) {
	V(Info).Run(msg)
}

func UnsafeSkip(msg string) {
	V(Info).UnsafeSkip(msg)
}

func Go(msg string) {
	V(Info).Go(msg)
}

func UnsafeGo(msg string) {
	V(Info).UnsafeGo(msg)
}

func Ok(msg string) {
	V(Info).Ok(msg)
}

func UnsafeOk(msg string) {
	V(Info).UnsafeOk(msg)
}

func Warn(msg string) {
	V(Info).Warn(msg)
}

func UnsafeWarn(msg string) {
	V(Info).UnsafeWarn(msg)
}

func Fail(msg string) {
	V(Info).Fail(msg)
}

func UnsafeFail(msg string) {
	V(Info).UnsafeFail(msg)
}

func Rollback(msg string) {
	V(Info).Rollback(msg)
}

func UnsafeRollback(msg string) {
	V(Info).UnsafeRollback(msg)
}

func NewLine(msg string) {
	V(Info).NewLine(msg)
}

func UnsafeNewLine(msg string) {
	V(Info).UnsafeNewLine(msg)
}

func FailWithDetails(msg string, details *bytes.Buffer) {
	V(Info).FailWithDetails(msg, details)
}

func Stderr(stderr *bytes.Buffer) {
	V(Info).Stderr(stderr)
}

func UnsafeStderr(stderr *bytes.Buffer) {
	V(Info).UnsafeStderr(stderr)
}

func Die(msg string, err error) {
	V(Info).Die(msg, err)
}

func UnsafeDie(msg string, err error) {
	V(Info).UnsafeDie(msg, err)
}

var levelToStringMap = map[Level]string{
	Trace:   "trace",
	Debug:   "debug",
	Verbose: "verbose",
	Info:    "info",
	Off:     "off",
}

func LevelToString(level Level) (string, bool) {
	v, ok := levelToStringMap[level]
	return v, ok
}

func MustLevelToString(level Level) string {
	v, ok := LevelToString(level)
	if !ok {
		panic(fmt.Errorf("invalid log level: %v", level))
	}
	return v
}

var stringToLevelMap = map[string]Level{
	"trace":   Trace,
	"debug":   Debug,
	"verbose": Verbose,
	"info":    Info,
	"off":     Off,
}

func StringToLevel(levelString string) (Level, bool) {
	v, ok := stringToLevelMap[levelString]
	return v, ok
}

func MustStringToLevel(levelString string) Level {
	level, ok := StringToLevel(levelString)
	if !ok {
		panic(errors.New("invalid log level string: " + levelString))
	}
	return level
}

func LevelStrings() []string {
	levels := make([]string, 0, len(stringToLevelMap))
	for k := range stringToLevelMap {
		levels = append(levels, k)
	}
	return levels
}
<commit_msg>Added `Log` function.<commit_after>package log

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
)

type (
	Level  uint32
	Logger bool
)

const (
	Trace = iota
	Debug
	Verbose
	Info
	Off
)

var lock sync.Mutex

var v Level = Info

func SetV(level Level) {
	atomic.StoreUint32((*uint32)(&v), uint32(level))
}

func V(level Level) Logger {
	if atomic.LoadUint32((*uint32)(&v)) > uint32(level) {
		return Logger(false)
	}
	return Logger(true)
}

func (l Logger) log(v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprint(os.Stderr, v...)
	}
}

func (l Logger) unsafeLog(v ...interface{}) {
	if l {
		fmt.Fprint(os.Stderr, v...)
	}
}

func (l Logger) logf(format string, v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprintf(os.Stderr, format, v...)
	}
}

func (l Logger) unsafeLogf(format string, v ...interface{}) {
	if l {
		fmt.Fprintf(os.Stderr, format, v...)
	}
}

func (l Logger) logln(v ...interface{}) {
	if l {
		lock.Lock()
		defer lock.Unlock()
		fmt.Fprintln(os.Stderr, v...)
	}
}

func (l Logger) unsafeLogln(v ...interface{}) {
	if l {
		fmt.Fprintln(os.Stderr, v...)
	}
}

func (l Logger) Lock() {
	lock.Lock()
}

func (l Logger) Unlock() {
	lock.Unlock()
}

func (l Logger) Print(v ...interface{}) {
	l.log(v...)
}

func (l Logger) UnsafePrint(v ...interface{}) {
	l.unsafeLog(v...)
}

func (l Logger) Printf(format string, v ...interface{}) {
	l.logf(format, v...)
}

func (l Logger) UnsafePrintf(format string, v ...interface{}) {
	l.unsafeLogf(format, v...)
}

func (l Logger) Println(v ...interface{}) {
	l.logln(v...)
}

func (l Logger) UnsafePrintln(v ...interface{}) {
	l.unsafeLogln(v...)
}

func (l Logger) Fatal(v ...interface{}) {
	l.log(v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatal(v ...interface{}) {
	l.unsafeLog(v...)
	os.Exit(1)
}

func (l Logger) Fatalf(format string, v ...interface{}) {
	l.logf(format, v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatalf(format string, v ...interface{}) {
	l.unsafeLogf(format, v...)
	os.Exit(1)
}

func (l Logger) Fatalln(v ...interface{}) {
	l.logln(v...)
	os.Exit(1)
}

func (l Logger) UnsafeFatalln(v ...interface{}) {
	l.unsafeLogln(v...)
	os.Exit(1)
}

func (l Logger) Run(msg string) {
	l.logf("[RUN]      %v\n", msg)
}

func (l Logger) UnsafeRun(msg string) {
	l.unsafeLogf("[RUN]      %v\n", msg)
}

func (l Logger) Skip(msg string) {
	l.logf("[SKIP]     %v\n", msg)
}

func (l Logger) UnsafeSkip(msg string) {
	l.unsafeLogf("[SKIP]     %v\n", msg)
}

func (l Logger) Warn(msg string) {
	l.logf("[WARN]     %v\n", msg)
}

func (l Logger) UnsafeWarn(msg string) {
	l.unsafeLogf("[WARN]     %v\n", msg)
}

func (l Logger) Go(msg string) {
	l.logf("[GO]       %v\n", msg)
}

func (l Logger) UnsafeGo(msg string) {
	l.unsafeLogf("[GO]       %v\n", msg)
}

func (l Logger) Log(msg string) {
	l.logf("[LOG]      %v\n", msg)
}

func (l Logger) UnsafeLog(msg string) {
	l.unsafeLogf("[LOG]      %v\n", msg)
}

func (l Logger) Ok(msg string) {
	l.logf("[OK]       %v\n", msg)
}

func (l Logger) UnsafeOk(msg string) {
	l.unsafeLogf("[OK]       %v\n", msg)
}

func (l Logger) Fail(msg string) {
	l.logf("[FAIL]     %v\n", msg)
}

func (l Logger) UnsafeFail(msg string) {
	l.unsafeLogf("[FAIL]     %v\n", msg)
}

func (l Logger) Rollback(msg string) {
	l.logf("[ROLLBACK] %v\n", msg)
}

func (l Logger) UnsafeRollback(msg string) {
	l.unsafeLogf("[ROLLBACK] %v\n", msg)
}

func (l Logger) NewLine(msg string) {
	l.logf("           %v\n", msg)
}

func (l Logger) UnsafeNewLine(msg string) {
	l.unsafeLogf("           %v\n", msg)
}

func (l Logger) Stderr(stderr *bytes.Buffer) {
	if stderr != nil && stderr.Len() != 0 {
		lock.Lock()
		defer lock.Unlock()
		l.UnsafePrintln("<<<<< stderr")
		l.UnsafePrint(stderr)
		l.UnsafePrintln(">>>>> stderr")
	}
}

func (l Logger) UnsafeStderr(stderr *bytes.Buffer) {
	if stderr != nil && stderr.Len() != 0 {
		l.UnsafePrintln("<<<<< stderr")
		l.UnsafePrint(stderr)
		l.UnsafePrintln(">>>>> stderr")
	}
}

func (l Logger) Die(msg string, err error) {
	lock.Lock()
	defer lock.Unlock()
	l.UnsafeFail(msg)
	l.UnsafeFatalln("\nError:", err)
}

func (l Logger) UnsafeDie(msg string, err error) {
	l.UnsafeFail(msg)
	l.UnsafeFatalln("\nError:", err)
}

func (l Logger) FailWithDetails(msg string, details *bytes.Buffer) {
	if msg != "" {
		l.Fail(msg)
	}
	l.Stderr(details)
}

func Print(v ...interface{}) {
	V(Info).Print(v...)
}

func UnsafePrint(v ...interface{}) {
	V(Info).UnsafePrint(v...)
}

func Printf(format string, v ...interface{}) {
	V(Info).Printf(format, v...)
}

func UnsafePrintf(format string, v ...interface{}) {
	V(Info).UnsafePrintf(format, v...)
}

func Println(v ...interface{}) {
	V(Info).Println(v...)
}

func UnsafePrintln(v ...interface{}) {
	V(Info).UnsafePrintln(v...)
}

func Fatal(v ...interface{}) {
	V(Info).Fatal(v...)
}

func UnsafeFatal(v ...interface{}) {
	V(Info).UnsafeFatal(v...)
}

func Fatalf(format string, v ...interface{}) {
	V(Info).Fatalf(format, v...)
}

func UnsafeFatalf(format string, v ...interface{}) {
	V(Info).UnsafeFatalf(format, v...)
}

func Fatalln(v ...interface{}) {
	V(Info).Fatalln(v...)
}

func UnsafeFatalln(v ...interface{}) {
	V(Info).UnsafeFatalln(v...)
}

func Run(msg string) {
	V(Info).Run(msg)
}

func UnsafeRun(msg string) {
	V(Info).UnsafeRun(msg)
}

func Log(msg string) {
	V(Info).Log(msg)
}

func UnsafeLog(msg string) {
	V(Info).UnsafeLog(msg)
}

func Skip(msg string) {
	V(Info).Run(msg)
}

func UnsafeSkip(msg string) {
	V(Info).UnsafeSkip(msg)
}

func Go(msg string) {
	V(Info).Go(msg)
}

func UnsafeGo(msg string) {
	V(Info).UnsafeGo(msg)
}

func Ok(msg string) {
	V(Info).Ok(msg)
}

func UnsafeOk(msg string) {
	V(Info).UnsafeOk(msg)
}

func Warn(msg string) {
	V(Info).Warn(msg)
}

func UnsafeWarn(msg string) {
	V(Info).UnsafeWarn(msg)
}

func Fail(msg string) {
	V(Info).Fail(msg)
}

func UnsafeFail(msg string) {
	V(Info).UnsafeFail(msg)
}

func Rollback(msg string) {
	V(Info).Rollback(msg)
}

func UnsafeRollback(msg string) {
	V(Info).UnsafeRollback(msg)
}

func NewLine(msg string) {
	V(Info).NewLine(msg)
}

func UnsafeNewLine(msg string) {
	V(Info).UnsafeNewLine(msg)
}

func FailWithDetails(msg string, details *bytes.Buffer) {
	V(Info).FailWithDetails(msg, details)
}

func Stderr(stderr *bytes.Buffer) {
	V(Info).Stderr(stderr)
}

func UnsafeStderr(stderr *bytes.Buffer) {
	V(Info).UnsafeStderr(stderr)
}

func Die(msg string, err error) {
	V(Info).Die(msg, err)
}

func UnsafeDie(msg string, err error) {
	V(Info).UnsafeDie(msg, err)
}

var levelToStringMap = map[Level]string{
	Trace:   "trace",
	Debug:   "debug",
	Verbose: "verbose",
	Info:    "info",
	Off:     "off",
}

func LevelToString(level Level) (string, bool) {
	v, ok := levelToStringMap[level]
	return v, ok
}

func MustLevelToString(level Level) string {
	v, ok := LevelToString(level)
	if !ok {
		panic(fmt.Errorf("invalid log level: %v", level))
	}
	return v
}

var stringToLevelMap = map[string]Level{
	"trace":   Trace,
	"debug":   Debug,
	"verbose": Verbose,
	"info":    Info,
	"off":     Off,
}

func StringToLevel(levelString string) (Level, bool) {
	v, ok := stringToLevelMap[levelString]
	return v, ok
}

func MustStringToLevel(levelString string) Level {
	level, ok := StringToLevel(levelString)
	if !ok {
		panic(errors.New("invalid log level string: " + levelString))
	}
	return level
}

func LevelStrings() []string {
	levels := make([]string, 0, len(stringToLevelMap))
	for k := range stringToLevelMap {
		levels = append(levels, k)
	}
	return levels
}
<|endoftext|>


<commit_before>// Copyright 2016 Koichi Shiraishi. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package nvim

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/juju/errors"
	"github.com/neovim-go/vim"
)

var (
	// ErrorColor highlight error message use Identifier syntax color.
	ErrorColor = "Identifier"
	// ProgressColor highlight progress message use Identifier syntax color.
	ProgressColor = "Identifier"
	// SuccessColor highlight success message use Identifier syntax color.
	SuccessColor = "Function"
)

// Echo provide the vim 'echo' command.
func Echo(v *vim.Vim, format string, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echo '" + fmt.Sprintf(format, a...) + "'")
}

// EchoRaw provide the raw output vim 'echo' command.
func EchoRaw(v *vim.Vim, a string) error {
	v.Command("redraw")
	return v.Command("echo \"" + a + "\"")
}

// Echomsg provide the vim 'echomsg' command.
func Echomsg(v *vim.Vim, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echomsg '" + strings.TrimSpace(fmt.Sprintln(a...)) + "'")
}

// Echoerr provide the vim 'echoerr' command.
func Echoerr(v *vim.Vim, format string, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echoerr '" + fmt.Sprintf(format, a...) + "'")
}

// ErrorWrap splits the errors.Annotate's cause and error messages,
// and provide the vim 'echo' message with 'echohl' highlighting to cause text.
func ErrorWrap(v *vim.Vim, err error) error {
	v.Command("redraw")
	er := strings.SplitAfterN(fmt.Sprintf("%s", err), ": ", 2)
	if strings.Contains(er[1], `"`) {
		er[1] = strings.Replace(er[1], `"`, "", -1)
	}
	if os.Getenv("NVIM_GO_DEBUG") != "" {
		log.Printf("Error stack\n%s", errors.ErrorStack(err))
	}
	return v.Command("echo \"" + er[0] + "\" | echohl " + ErrorColor + " | echon \"" + er[1] + "\" | echohl None")
}

// EchohlErr provide the vim 'echo' command with the 'echohl' highlighting prefix text.
func EchohlErr(v *vim.Vim, prefix string, a ...interface{}) error {
	v.Command("redraw")
	if prefix != "" {
		prefix += ": "
	}
	er := fmt.Sprintf("%s", a...)
	return v.Command("echo '" + prefix + "' | echohl " + ErrorColor + " | echon \"" + er + "\" | echohl None")
}

// EchohlBefore provide the vim 'echo' command with the 'echohl' highlighting prefix text.
func EchohlBefore(v *vim.Vim, prefix string, highlight string, format string, a ...interface{}) error {
	v.Command("redraw")
	suffix := "\" | echohl None | echon \""
	if prefix != "" {
		suffix += ": "
	}
	return v.Command("echohl " + highlight + " | echo \"" + prefix + suffix + fmt.Sprintf(format, a...) + "\" | echohl None")
}

// EchohlAfter provide the vim 'echo' command with the 'echohl' highlighting message text.
func EchohlAfter(v *vim.Vim, prefix string, highlight string, format string, a ...interface{}) error {
	v.Command("redraw")
	if prefix != "" {
		prefix += ": "
	}
	return v.Command("echo \"" + prefix + "\" | echohl " + highlight + " | echon \"" + fmt.Sprintf(format, a...) + "\" | echohl None")
}

// EchoProgress displays a command progress message to echo area.
func EchoProgress(v *vim.Vim, prefix, format string, a ...interface{}) error {
	v.Command("redraw")
	msg := fmt.Sprintf(format, a...)
	return v.Command(fmt.Sprintf("echo \"%s: \" | echohl %s | echon \"%s ...\" | echohl None", prefix, ProgressColor, msg))
}

// EchoSuccess displays the success of the command to echo area.
func EchoSuccess(v *vim.Vim, prefix string, msg string) error {
	v.Command("redraw")
	if msg != "" {
		msg = " | " + msg
	}
	return v.Command(fmt.Sprintf("echo \"%s: \" | echohl %s | echon 'SUCCESS' | echohl None | echon '%s' | echohl None", prefix, SuccessColor, msg))
}

// ReportError output of the accumulated errors report.
// TODO(zchee): research vim.ReportError behavior
// Why it does not immediately display error?
func ReportError(v *vim.Vim, format string, a ...interface{}) error {
	return v.ReportError(fmt.Sprintf(format, a...))
}

// ClearMsg cleanups the echo area.
func ClearMsg(v *vim.Vim) error {
	return v.Command("echon")
}
<commit_msg>nvim: Fix ErrorWrap to return immediately if err is nil<commit_after>// Copyright 2016 Koichi Shiraishi. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package nvim

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/juju/errors"
	"github.com/neovim-go/vim"
)

var (
	// ErrorColor highlight error message use Identifier syntax color.
	ErrorColor = "Identifier"
	// ProgressColor highlight progress message use Identifier syntax color.
	ProgressColor = "Identifier"
	// SuccessColor highlight success message use Identifier syntax color.
	SuccessColor = "Function"
)

// Echo provide the vim 'echo' command.
func Echo(v *vim.Vim, format string, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echo '" + fmt.Sprintf(format, a...) + "'")
}

// EchoRaw provide the raw output vim 'echo' command.
func EchoRaw(v *vim.Vim, a string) error {
	v.Command("redraw")
	return v.Command("echo \"" + a + "\"")
}

// Echomsg provide the vim 'echomsg' command.
func Echomsg(v *vim.Vim, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echomsg '" + strings.TrimSpace(fmt.Sprintln(a...)) + "'")
}

// Echoerr provide the vim 'echoerr' command.
func Echoerr(v *vim.Vim, format string, a ...interface{}) error {
	v.Command("redraw")
	return v.Command("echoerr '" + fmt.Sprintf(format, a...) + "'")
}

// ErrorWrap splits the errors.Annotate's cause and error messages,
// and provide the vim 'echo' message with 'echohl' highlighting to cause text.
func ErrorWrap(v *vim.Vim, err error) error {
	if err == nil {
		return nil
	}
	v.Command("redraw")
	er := strings.SplitAfterN(fmt.Sprintf("%s", err), ": ", 2)
	if strings.Contains(er[1], `"`) {
		er[1] = strings.Replace(er[1], `"`, "", -1)
	}
	if os.Getenv("NVIM_GO_DEBUG") != "" {
		log.Printf("Error stack\n%s", errors.ErrorStack(err))
	}
	return v.Command("echo \"" + er[0] + "\" | echohl " + ErrorColor + " | echon \"" + er[1] + "\" | echohl None")
}

// EchohlErr provide the vim 'echo' command with the 'echohl' highlighting prefix text.
func EchohlErr(v *vim.Vim, prefix string, a ...interface{}) error {
	v.Command("redraw")
	if prefix != "" {
		prefix += ": "
	}
	er := fmt.Sprintf("%s", a...)
	return v.Command("echo '" + prefix + "' | echohl " + ErrorColor + " | echon \"" + er + "\" | echohl None")
}

// EchohlBefore provide the vim 'echo' command with the 'echohl' highlighting prefix text.
func EchohlBefore(v *vim.Vim, prefix string, highlight string, format string, a ...interface{}) error {
	v.Command("redraw")
	suffix := "\" | echohl None | echon \""
	if prefix != "" {
		suffix += ": "
	}
	return v.Command("echohl " + highlight + " | echo \"" + prefix + suffix + fmt.Sprintf(format, a...) + "\" | echohl None")
}

// EchohlAfter provide the vim 'echo' command with the 'echohl' highlighting message text.
func EchohlAfter(v *vim.Vim, prefix string, highlight string, format string, a ...interface{}) error {
	v.Command("redraw")
	if prefix != "" {
		prefix += ": "
	}
	return v.Command("echo \"" + prefix + "\" | echohl " + highlight + " | echon \"" + fmt.Sprintf(format, a...) + "\" | echohl None")
}

// EchoProgress displays a command progress message to echo area.
func EchoProgress(v *vim.Vim, prefix, format string, a ...interface{}) error {
	v.Command("redraw")
	msg := fmt.Sprintf(format, a...)
	return v.Command(fmt.Sprintf("echo \"%s: \" | echohl %s | echon \"%s ...\" | echohl None", prefix, ProgressColor, msg))
}

// EchoSuccess displays the success of the command to echo area.
func EchoSuccess(v *vim.Vim, prefix string, msg string) error {
	v.Command("redraw")
	if msg != "" {
		msg = " | " + msg
	}
	return v.Command(fmt.Sprintf("echo \"%s: \" | echohl %s | echon 'SUCCESS' | echohl None | echon '%s' | echohl None", prefix, SuccessColor, msg))
}

// ReportError output of the accumulated errors report.
// TODO(zchee): research vim.ReportError behavior
// Why it does not immediately display error?
func ReportError(v *vim.Vim, format string, a ...interface{}) error {
	return v.ReportError(fmt.Sprintf(format, a...))
}

// ClearMsg cleanups the echo area.
func ClearMsg(v *vim.Vim) error {
	return v.Command("echon")
}
<|endoftext|>


<commit_before>// santiago - webhook dispatching service
// https://github.com/topfreegames/santiago
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package api_test

import (
	"encoding/json"
	"net/http"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/satori/go.uuid"
	. "github.com/topfreegames/santiago/testing"
)

var _ = Describe("Status Handler", func() {
	var logger *MockLogger

	BeforeEach(func() {
		logger = NewMockLogger()
	})

	It("Should respond with number of items in queue and errors", func() {
		a, err := GetDefaultTestApp(logger)
		a.Queue = uuid.NewV4().String()

		Expect(err).NotTo(HaveOccurred())
		res := Get(a, "/status")

		Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))

		var obj map[string]interface{}
		err = json.Unmarshal([]byte(res.Body().Raw()), &obj)
		Expect(obj["errors"]).To(BeEquivalentTo(0.0))
		Expect(obj["messagesInQueue"]).To(BeEquivalentTo(0))
	})

	It("Should respond with number of items in queue when queue has items", func() {
		a, err := GetDefaultTestApp(logger)
		a.Queue = uuid.NewV4().String()

		for i := 0; i < 10; i++ {
			a.Client.LPush(a.Queue, "{\"x\":1}")
		}

		Expect(err).NotTo(HaveOccurred())
		res := Get(a, "/status")

		Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))

		var obj map[string]interface{}
		err = json.Unmarshal([]byte(res.Body().Raw()), &obj)
		Expect(obj["errors"]).To(BeEquivalentTo(0.0))
		Expect(obj["messagesInQueue"]).To(BeEquivalentTo(10))
	})

	Measure("it should get status", func(b Benchmarker) {
		app, err := GetDefaultTestApp(logger)
		Expect(err).NotTo(HaveOccurred())

		runtime := b.Time("runtime", func() {
			res := Get(app, "/status")
			Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))
			Expect(res.Body().Raw()).NotTo(BeEmpty())
		})

		Expect(runtime.Seconds()).Should(BeNumerically("<", 0.01), "Status shouldn't take too long.")
	}, 200)

})
<commit_msg>Status should take at most 20ms<commit_after>// santiago - webhook dispatching service
// https://github.com/topfreegames/santiago
// Licensed under the MIT license:
// http://www.opensource.org/licenses/mit-license
// Copyright © 2016 Top Free Games <backend@tfgco.com>

package api_test

import (
	"encoding/json"
	"net/http"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/satori/go.uuid"
	. "github.com/topfreegames/santiago/testing"
)

var _ = Describe("Status Handler", func() {
	var logger *MockLogger

	BeforeEach(func() {
		logger = NewMockLogger()
	})

	It("Should respond with number of items in queue and errors", func() {
		a, err := GetDefaultTestApp(logger)
		a.Queue = uuid.NewV4().String()

		Expect(err).NotTo(HaveOccurred())
		res := Get(a, "/status")

		Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))

		var obj map[string]interface{}
		err = json.Unmarshal([]byte(res.Body().Raw()), &obj)
		Expect(obj["errors"]).To(BeEquivalentTo(0.0))
		Expect(obj["messagesInQueue"]).To(BeEquivalentTo(0))
	})

	It("Should respond with number of items in queue when queue has items", func() {
		a, err := GetDefaultTestApp(logger)
		a.Queue = uuid.NewV4().String()

		for i := 0; i < 10; i++ {
			a.Client.LPush(a.Queue, "{\"x\":1}")
		}

		Expect(err).NotTo(HaveOccurred())
		res := Get(a, "/status")

		Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))

		var obj map[string]interface{}
		err = json.Unmarshal([]byte(res.Body().Raw()), &obj)
		Expect(obj["errors"]).To(BeEquivalentTo(0.0))
		Expect(obj["messagesInQueue"]).To(BeEquivalentTo(10))
	})

	Measure("it should get status", func(b Benchmarker) {
		app, err := GetDefaultTestApp(logger)
		Expect(err).NotTo(HaveOccurred())

		runtime := b.Time("runtime", func() {
			res := Get(app, "/status")
			Expect(res.Raw().StatusCode).To(Equal(http.StatusOK))
			Expect(res.Body().Raw()).NotTo(BeEmpty())
		})

		Expect(runtime.Seconds()).Should(BeNumerically("<", 0.02), "Status shouldn't take too long.")
	}, 200)

})
<|endoftext|>


<commit_before>package main

import (
	"os"
	"github.com/ArthurHlt/travis-resource/common"
	"github.com/ArthurHlt/travis-resource/model"
	"encoding/json"
	"errors"
	"github.com/ArthurHlt/travis-resource/travis"
	"strconv"
	"time"
)

func main() {

	var request model.OutRequest
	err := json.NewDecoder(os.Stdin).Decode(&request)
	common.FatalIf("failed to read request", err)
	if request.Source.Repository == "" {
		common.FatalIf("can't get build", errors.New("there is no repository set"))
	}
	travisClient, err := common.MakeTravisClient(request.Source)
	common.FatalIf("failed to create travis client", err)

	var build travis.Build
	repository := request.Source.Repository
	if request.OutParams.Repository != "" {
		repository = request.OutParams.Repository
	}
	buildParam := ""
	if buildParamNumber, ok := request.OutParams.Build.(float64); ok {
		buildParam = strconv.FormatFloat(buildParamNumber, 'f', 0, 64)
	}
	if buildParamString, ok := request.OutParams.Build.(string); ok {
		buildParam = buildParamString
	}
	if buildParam == "latest" || (request.OutParams.Repository != "" && request.OutParams.Build == "" && request.OutParams.Branch == "") {
		build, err = travisClient.Builds.GetFirstFinishedBuild(repository)
		common.FatalIf("can't get build", err)
	} else if buildParam != "" {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, buildParam)
		common.FatalIf("can't get build", err)
	} else if request.OutParams.Branch != "" {
		build, err = travisClient.Builds.GetFirstFinishedBuildWithBranch(repository, request.OutParams.Branch)
		common.FatalIf("can't get build", err)
	} else {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, request.Version.BuildNumber)
		common.FatalIf("can't get build", err)
	}

	travisClient.Builds.Restart(build.Id)
	if !request.OutParams.UnWaitBuild {
		waitBuild(travisClient, repository, build.Number, request)
	}
	build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, build.Number)
	common.FatalIf("can't get build after restart", err)
	response := model.InResponse{common.GetMetadatasFromBuild(build), model.Version{build.Number}}
	json.NewEncoder(os.Stdout).Encode(response)
}
func waitBuild(travisClient *travis.Client, repository, buildNumber string, request model.OutRequest) {
	var build travis.Build
	var err error
	var travisUrl string
	if request.Source.Url != "" {
		travisUrl = request.Source.Url
	} else {
		travisUrl = common.GetTravisUrl(request.Source.Pro)
	}
	travisUrl = common.GetTravisDashboardUrl(travisUrl)
	for {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, buildNumber)
		common.FatalIf("can't get build after restart", err)
		if !stringInSlice(build.State, travis.RUNNING_STATE) {
			break
		}
		time.Sleep(5 * time.Second)
	}
	if build.State != travis.SUCCEEDED_STATE {
		common.FatalIf("Build '" + build.Number + "' in error, see",
			errors.New(travisUrl + repository + "/builds/" + strconv.Itoa(int(build.Id))))
	}
}
func stringInSlice(str string, list []string) bool {
	for _, v := range list {
		if v == str {
			return true
		}
	}
	return false
}<commit_msg>add state in out error message<commit_after>package main

import (
	"os"
	"github.com/ArthurHlt/travis-resource/common"
	"github.com/ArthurHlt/travis-resource/model"
	"encoding/json"
	"errors"
	"github.com/ArthurHlt/travis-resource/travis"
	"strconv"
	"time"
)

func main() {

	var request model.OutRequest
	err := json.NewDecoder(os.Stdin).Decode(&request)
	common.FatalIf("failed to read request", err)
	if request.Source.Repository == "" {
		common.FatalIf("can't get build", errors.New("there is no repository set"))
	}
	travisClient, err := common.MakeTravisClient(request.Source)
	common.FatalIf("failed to create travis client", err)

	var build travis.Build
	repository := request.Source.Repository
	if request.OutParams.Repository != "" {
		repository = request.OutParams.Repository
	}
	buildParam := ""
	if buildParamNumber, ok := request.OutParams.Build.(float64); ok {
		buildParam = strconv.FormatFloat(buildParamNumber, 'f', 0, 64)
	}
	if buildParamString, ok := request.OutParams.Build.(string); ok {
		buildParam = buildParamString
	}
	if buildParam == "latest" || (request.OutParams.Repository != "" && request.OutParams.Build == "" && request.OutParams.Branch == "") {
		build, err = travisClient.Builds.GetFirstFinishedBuild(repository)
		common.FatalIf("can't get build", err)
	} else if buildParam != "" {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, buildParam)
		common.FatalIf("can't get build", err)
	} else if request.OutParams.Branch != "" {
		build, err = travisClient.Builds.GetFirstFinishedBuildWithBranch(repository, request.OutParams.Branch)
		common.FatalIf("can't get build", err)
	} else {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, request.Version.BuildNumber)
		common.FatalIf("can't get build", err)
	}

	travisClient.Builds.Restart(build.Id)
	if !request.OutParams.UnWaitBuild {
		waitBuild(travisClient, repository, build.Number, request)
	}
	build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, build.Number)
	common.FatalIf("can't get build after restart", err)
	response := model.InResponse{common.GetMetadatasFromBuild(build), model.Version{build.Number}}
	json.NewEncoder(os.Stdout).Encode(response)
}
func waitBuild(travisClient *travis.Client, repository, buildNumber string, request model.OutRequest) {
	var build travis.Build
	var err error
	var travisUrl string
	if request.Source.Url != "" {
		travisUrl = request.Source.Url
	} else {
		travisUrl = common.GetTravisUrl(request.Source.Pro)
	}
	travisUrl = common.GetTravisDashboardUrl(travisUrl)
	for {
		build, err = travisClient.Builds.GetFirstBuildFromBuildNumber(repository, buildNumber)
		common.FatalIf("can't get build after restart", err)
		if !stringInSlice(build.State, travis.RUNNING_STATE) {
			break
		}
		time.Sleep(5 * time.Second)
	}
	if build.State != travis.SUCCEEDED_STATE {
		common.FatalIf("Build '" + build.Number + "' in failed",
			errors.New("\n\tstate:" + build.State + "\n\tsee:" + travisUrl + repository + "/builds/" + strconv.Itoa(int(build.Id))))
	}
}
func stringInSlice(str string, list []string) bool {
	for _, v := range list {
		if v == str {
			return true
		}
	}
	return false
}<|endoftext|>


<commit_before>package jsonschema

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"testing"
)

func TestDraft4(t *testing.T) {
	testResources := filepath.Join("JSON-Schema-Test-Suite", "tests", "draft4")
	_, err := os.Stat(testResources)
	if err != nil {
		t.Error("Test suite missing. Run `git submodule update` to download it.")
	}
	err = filepath.Walk(testResources, testFileRunner(t))
	if err != nil {
		t.Error(err.Error())
	}
}

func testFileRunner(t *testing.T) func(string, os.FileInfo, error) error {
	return func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		var testFile []testCase
		err = json.NewDecoder(file).Decode(&testFile)
		if err != nil {
			return err
		}

		for _, description := range testFile {
			schema, err := Parse(bytes.NewReader(description.Schema))
			if err != nil {
				return err
			}
			for _, test := range description.Tests {
				errorList := schema.Validate(test.Data)
				message := failureMessage(errorList, test, description, path)
				if len(message) > 0 {
					t.Error(message)
				}
			}
		}
		return nil
	}
}

func failureMessage(errorList []ValidationError, tst testInstance, cse testCase, path string) string {
	var validated bool
	if len(errorList) == 0 {
		validated = true
	} else if len(errorList) > 0 {
		validated = false
	}

	var failure string
	if validated && !tst.Valid {
		failure = "schema.Validate validated bad data."
	} else if !validated && tst.Valid {
		failure = "schema.Validate failed to validate good data."
	}

	var message string
	if len(failure) > 0 {
		message = fmt.Sprintf(`%s
file: %s
test case description: %s
schema: %s
test instance description: %s
test data: %s
result of schema.Validate:
	expected: %t
	actual: %t
	actual validation errors: %s

`, failure, path, cse.Description, cse.Schema, tst.Description, tst.Data, tst.Valid, validated, errorList)
	}
	return message
}

type testCase struct {
	Description string          `json:"description"`
	Schema      json.RawMessage `json:"schema"`
	Properties  json.RawMessage `json:"properties"`
	Required    json.RawMessage `json:"required"`
	Tests       []testInstance  `json:"tests"`
}

type testInstance struct {
	Description string          `json:"description"`
	Data        json.RawMessage `json:"data"`
	Valid       bool            `json:"valid"`
}
<commit_msg>Log number of successes and failures.<commit_after>package jsonschema

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"testing"
)

func TestDraft4(t *testing.T) {
	testResources := filepath.Join("JSON-Schema-Test-Suite", "tests", "draft4")
	_, err := os.Stat(testResources)
	if err != nil {
		t.Error("Test suite missing. Run `git submodule update` to download it.")
	}
	var failures, successes int
	err = filepath.Walk(testResources, testFileRunner(t, &failures, &successes))
	if err != nil {
		t.Error(err.Error())
	}
	t.Logf("%d failures, %d successes.", failures, successes)
}

func testFileRunner(t *testing.T, failures, successes *int) func(string, os.FileInfo, error) error {
	return func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		var testFile []testCase
		err = json.NewDecoder(file).Decode(&testFile)
		if err != nil {
			return err
		}

		for _, description := range testFile {
			schema, err := Parse(bytes.NewReader(description.Schema))
			if err != nil {
				return err
			}
			for _, test := range description.Tests {
				errorList := schema.Validate(test.Data)
				message := failureMessage(errorList, test, description, path)
				if len(message) > 0 {
					t.Error(message)
					*failures += 1
				} else {
					*successes += 1
				}
			}
		}
		return nil
	}
}

func failureMessage(errorList []ValidationError, tst testInstance, cse testCase, path string) string {
	var validated bool
	if len(errorList) == 0 {
		validated = true
	} else if len(errorList) > 0 {
		validated = false
	}

	var failureName string
	if validated && !tst.Valid {
		failureName = "schema.Validate validated bad data."
	} else if !validated && tst.Valid {
		failureName = "schema.Validate failed to validate good data."
	}

	var message string
	if len(failureName) > 0 {
		message = fmt.Sprintf(`%s
file: %s
test case description: %s
schema: %s
test instance description: %s
test data: %s
result of schema.Validate:
	expected: %t
	actual: %t
	actual validation errors: %s

`, failureName, path, cse.Description, cse.Schema, tst.Description, tst.Data, tst.Valid, validated, errorList)
	}
	return message
}

type testCase struct {
	Description string          `json:"description"`
	Schema      json.RawMessage `json:"schema"`
	Properties  json.RawMessage `json:"properties"`
	Required    json.RawMessage `json:"required"`
	Tests       []testInstance  `json:"tests"`
}

type testInstance struct {
	Description string          `json:"description"`
	Data        json.RawMessage `json:"data"`
	Valid       bool            `json:"valid"`
}
<|endoftext|>


<commit_before>package main

import ()

func main () {
}
<commit_msg>orchestrator: Make it not terminate<commit_after>package main

import ()

func main() {
	for {
	}
}
<|endoftext|>


<commit_before>// Tideland Go REST Server Library - JSON Web Token - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package jwt_test

//--------------------
// IMPORTS
//--------------------

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"

	"github.com/tideland/gorest/jwt"
	"github.com/tideland/gorest/rest"
	"github.com/tideland/gorest/restaudit"
)

//--------------------
// TESTS
//--------------------

// TestDecodeInvalidRequest tests the decoding of requests
// without a header or an invalid one.
func TestDecodeInvalidRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode invalid requests")
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, false)
	err := mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform request without authorization.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	ts.DoRequest(req)
}

// TestDecodeRequest tests the decoding of a token
// in a handler.
func TestDecodeRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestDecodeCachedRequest tests the decoding of a token
// in a handler including usage of the cache.
func TestDecodeCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	claimsOut = jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyRequest tests the verification of a token
// in a handler.
func TestVerifyRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyCachedRequest tests the verification of a token
// in a handler including usage of the cache.
func TestVerifyCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

//--------------------
// HANDLER
//--------------------

// testAsserter instances will handle the assertions in the testHandler.
type testAsserter func(job rest.Job) (bool, error)

func newDecodeAsserter(assert audit.Assertion, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.DecodeCachedFromJob(job, cache)
		} else {
			token, err = jwt.DecodeFromJob(job)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

func newVerifyAsserter(assert audit.Assertion, key jwt.Key, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.VerifyCachedFromJob(job, cache, key)
		} else {
			token, err = jwt.VerifyFromJob(job, key)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

// testHandler is used in the test scenarios.
type testHandler struct {
	id       string
	asserter testAsserter
}

func newTestHandler(id string, asserter testAsserter) rest.ResourceHandler {
	return &testHandler{
		id:       id,
		asserter: asserter,
	}
}

func (th *testHandler) ID() string {
	return th.id
}

func (th *testHandler) Init(env rest.Environment, domain, resource string) error {
	return nil
}

func (th *testHandler) Get(job rest.Job) (bool, error) {
	return th.asserter(job)
}

//--------------------
// HELPERS
//--------------------

// newMultiplexer creates a new multiplexer with a testing context
// and a testing configuration.
func newMultiplexer(assert audit.Assertion) rest.Multiplexer {
	cfgStr := "{etc {basepath /}{default-domain default}{default-resource default}}"
	cfg, err := etc.ReadString(cfgStr)
	assert.Nil(err)
	return rest.NewMultiplexer(context.Background(), cfg)
}

// EOF
<commit_msg>Finalized testing<commit_after>// Tideland Go REST Server Library - JSON Web Token - Unit Tests
//
// Copyright (C) 2016-2017 Frank Mueller / Tideland / Oldenburg / Germany
//
// All rights reserved. Use of this source code is governed
// by the new BSD license.

package jwt_test

//--------------------
// IMPORTS
//--------------------

import (
	"context"
	"net/http"
	"testing"
	"time"

	"github.com/tideland/golib/audit"
	"github.com/tideland/golib/etc"

	"github.com/tideland/gorest/jwt"
	"github.com/tideland/gorest/rest"
	"github.com/tideland/gorest/restaudit"
)

//--------------------
// TESTS
//--------------------

// TestDecodeInvalidRequest tests the decoding of requests
// without a header or an invalid one.
func TestDecodeInvalidRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode invalid requests")
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newHeaderAsserter(assert, ".* request contains no authorization header")
	err := mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform request without authorization.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	resp := ts.DoRequest(req)
	ok := ""
	resp.AssertUnmarshalledBody(&ok)
	assert.Equal(ok, "OK")
}

// TestDecodeRequest tests the decoding of a token
// in a handler.
func TestDecodeRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestDecodeCachedRequest tests the decoding of a token
// in a handler including usage of the cache.
func TestDecodeCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing decode a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newDecodeAsserter(assert, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	claimsOut = jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyRequest tests the verification of a token
// in a handler.
func TestVerifyRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, false)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

// TestVerifyCachedRequest tests the verification of a token
// in a handler including usage of the cache.
func TestVerifyCachedRequest(t *testing.T) {
	assert := audit.NewTestingAssertion(t, true)
	assert.Logf("testing verify a request token using a cache")
	key := []byte("secret")
	claimsIn := initClaims()
	jwtIn, err := jwt.Encode(claimsIn, key, jwt.HS512)
	assert.Nil(err)
	// Setup the test server.
	mux := newMultiplexer(assert)
	ts := restaudit.StartServer(mux, assert)
	defer ts.Close()
	asserter := newVerifyAsserter(assert, key, true)
	err = mux.Register("test", "jwt", newTestHandler("jwt", asserter))
	assert.Nil(err)
	// Perform first test request.
	req := restaudit.NewRequest("GET", "/test/jwt/1234567890")
	req.AddHeader(restaudit.HeaderAccept, restaudit.ApplicationJSON)
	req.SetRequestProcessor(func(req *http.Request) *http.Request {
		return jwt.AddToRequest(req, jwtIn)
	})
	resp := ts.DoRequest(req)
	claimsOut := jwt.Claims{}
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
	// Perform second test request.
	resp = ts.DoRequest(req)
	resp.AssertUnmarshalledBody(&claimsOut)
	assert.Equal(claimsOut, claimsIn)
}

//--------------------
// HANDLER
//--------------------

// testAsserter instances will handle the assertions in the testHandler.
type testAsserter func(job rest.Job) (bool, error)

func newHeaderAsserter(assert audit.Assertion, pattern string) testAsserter {
	return func(job rest.Job) (bool, error) {
		token, err := jwt.DecodeFromJob(job)
		assert.Nil(token)
		assert.ErrorMatch(err, pattern)
		job.JSON(true).Write(rest.StatusOK, "OK")
		return true, nil
	}
}

func newDecodeAsserter(assert audit.Assertion, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.DecodeCachedFromJob(job, cache)
		} else {
			token, err = jwt.DecodeFromJob(job)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

func newVerifyAsserter(assert audit.Assertion, key jwt.Key, cached bool) testAsserter {
	var cache jwt.Cache
	if cached {
		cache = jwt.NewCache(time.Minute, time.Minute, time.Minute, 10)
	}
	return func(job rest.Job) (bool, error) {
		var token jwt.JWT
		var err error
		if cached {
			token, err = jwt.VerifyCachedFromJob(job, cache, key)
		} else {
			token, err = jwt.VerifyFromJob(job, key)
		}
		assert.Nil(err)
		assert.True(token.IsValid(time.Minute))
		subject, ok := token.Claims().Subject()
		assert.True(ok)
		assert.Equal(subject, job.ResourceID())
		job.JSON(true).Write(rest.StatusOK, token.Claims())
		return true, nil
	}
}

// testHandler is used in the test scenarios.
type testHandler struct {
	id       string
	asserter testAsserter
}

func newTestHandler(id string, asserter testAsserter) rest.ResourceHandler {
	return &testHandler{
		id:       id,
		asserter: asserter,
	}
}

func (th *testHandler) ID() string {
	return th.id
}

func (th *testHandler) Init(env rest.Environment, domain, resource string) error {
	return nil
}

func (th *testHandler) Get(job rest.Job) (bool, error) {
	return th.asserter(job)
}

//--------------------
// HELPERS
//--------------------

// newMultiplexer creates a new multiplexer with a testing context
// and a testing configuration.
func newMultiplexer(assert audit.Assertion) rest.Multiplexer {
	cfgStr := "{etc {basepath /}{default-domain default}{default-resource default}}"
	cfg, err := etc.ReadString(cfgStr)
	assert.Nil(err)
	return rest.NewMultiplexer(context.Background(), cfg)
}

// EOF
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]

// Package dispatch implements alert dispatching for the agent as a module.
// Persistent modules which generate alerts will have these alerts forwarded
// to this module if the dispatch module is active. The dispatch module can then
// forward the alerts on based on it's configuration.
package dispatch /* import "mig.ninja/mig/modules/dispatch" */

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime"

	"mig.ninja/mig"
	"mig.ninja/mig/modules"
)

type module struct {
}

// NewRun returns a new instance of a modules.Runner for this module.
func (m *module) NewRun() modules.Runner {
	return new(run)
}

func init() {
	modules.Register("dispatch", new(module))
}

type run struct {
	Parameters Parameters
	Results    modules.Result
}

func buildResults(e elements, r *modules.Result) (buf []byte, err error) {
	r.Success = true
	r.Elements = e
	r.FoundAnything = true
	buf, err = json.Marshal(r)
	return
}

var logChan chan string
var alertChan chan string
var handlerErrChan chan error
var configChan chan modules.ConfigParams

// We keep a copy of the agent environment, hostname, and tags here so we can
// use them in the dispatched output
var env mig.AgentEnv
var tags map[string]string
var agtHostname string

// messageBuf is a queue used to store incoming messages, and is drained by
// runDispatch
var messageBuf chan string

// Dispatch record describes the formatting of JSON data submitted from the dispatch
// module.
type DispatchRecord struct {
	Hostname string            `json:"hostname"`
	Env      mig.AgentEnv      `json:"environment"`
	Tags     map[string]string `json:"tags"`
	Details  interface{}       `json:"details"`
}

func (d *DispatchRecord) fromString(msg string) error {
	d.Hostname = agtHostname
	d.Env = env
	d.Tags = tags
	err := json.Unmarshal([]byte(msg), &d.Details)
	return err
}

func moduleMain() {
	var cfg config

	incfg := <-configChan
	agtHostname = incfg.Hostname
	tags = incfg.Tags
	buf, err := json.Marshal(incfg.Env)
	if err != nil {
		handlerErrChan <- err
		return
	}
	err = json.Unmarshal(buf, &env)
	if err != nil {
		handlerErrChan <- err
		return
	}
	buf, err = json.Marshal(incfg.Config)
	if err != nil {
		handlerErrChan <- err
		return
	}
	err = json.Unmarshal(buf, &cfg)
	if err != nil {
		handlerErrChan <- err
		return
	}
	logChan <- "module received configuration"
	if cfg.Dispatch.ChannelSize == 0 {
		cfg.Dispatch.ChannelSize = 1024
		logChan <- "channel size not specified, defaulting to 1024"
	}

	messageBuf = make(chan string, cfg.Dispatch.ChannelSize)

	// Register the dispatch function, which will be called when the module
	// recieves an alert message from the agent.
	modules.RegisterDispatchFunction(dispatchIn)

	// Start the dispatch dequeue function.
	err = runDispatch(cfg)
	if err != nil {
		handlerErrChan <- err
		return
	}
}

func dispatchIn(msg string) {
	select {
	case messageBuf <- msg:
	default:
		// If we can't queue the message it is just dropped
	}
}

func runDispatch(cfg config) error {
	var (
		httpClient http.Client
		dr         DispatchRecord
	)

	for {
		msg := <-messageBuf
		dr.fromString(msg)
		buf, err := json.Marshal(dr)
		if err != nil {
			logChan <- fmt.Sprintf("create dispatch record: %v", err)
			continue
		}
		b := bytes.NewBuffer(buf)
		resp, err := httpClient.Post(cfg.Dispatch.HTTPURL, "application/json", b)
		if err != nil {
			logChan <- fmt.Sprintf("http post: %v", err)
			continue
		}
		resp.Body.Close()
	}
	return nil
}

func requestHandler(p interface{}) (ret string) {
	var results modules.Result
	defer func() {
		if e := recover(); e != nil {
			results.Errors = append(results.Errors, fmt.Sprintf("%v", e))
			results.Success = false
			err, _ := json.Marshal(results)
			ret = string(err)
			return
		}
	}()
	e := elements{Ok: true}
	resp, err := buildResults(e, &results)
	if err != nil {
		panic(err)
	}
	return string(resp)
}

type config struct {
	Dispatch struct {
		HTTPURL     string `json:"httpurl"`
		ChannelSize int    `json:"channelsize"`
	} `json:"dispatch"`
}

// PersistModConfig returns a new configuration structure for this module.
func (r *run) PersistModConfig() interface{} {
	return &config{}
}

// RunPersist is the entry point for persistent execution of the module.
func (r *run) RunPersist(in modules.ModuleReader, out modules.ModuleWriter) {
	alertChan = make(chan string, 64)
	logChan = make(chan string, 64)
	regChan := make(chan string, 64)
	handlerErrChan = make(chan error, 64)
	configChan = make(chan modules.ConfigParams, 1)

	go moduleMain()
	l, spec, err := modules.GetPersistListener("dispatch")
	if err != nil {
		handlerErrChan <- err
	} else {
		regChan <- spec
	}
	go modules.HandlePersistRequest(l, requestHandler, handlerErrChan)
	modules.DefaultPersistHandlers(in, out, logChan, handlerErrChan, regChan,
		alertChan, configChan)
}

// Run is the entry point for a standard (e.g., query) based invocation of the module.
func (r *run) Run(in modules.ModuleReader) (resStr string) {
	defer func() {
		if e := recover(); e != nil {
			// return error in json
			r.Results.Errors = append(r.Results.Errors, fmt.Sprintf("%v", e))
			r.Results.Success = false
			err, _ := json.Marshal(r.Results)
			resStr = string(err)
			return
		}
	}()
	runtime.GOMAXPROCS(1)
	sockspec, err := modules.ReadPersistInputParameters(in, &r.Parameters)
	if err != nil {
		panic(err)
	}
	err = r.ValidateParameters()
	if err != nil {
		panic(err)
	}
	resStr = modules.SendPersistRequest(r.Parameters, sockspec)
	return
}

// ValidateParameters validates the parameters set in the runner for the module.
func (r *run) ValidateParameters() (err error) {
	return
}

// PrintResults returns the results of a query of this module in human readable form.
func (r *run) PrintResults(result modules.Result, foundOnly bool) (prints []string, err error) {
	var (
		elem elements
	)

	err = result.GetElements(&elem)
	if err != nil {
		panic(err)
	}
	resStr := fmt.Sprintf("ok:%v", elem.Ok)
	prints = append(prints, resStr)
	if !foundOnly {
		for _, we := range result.Errors {
			prints = append(prints, we)
		}
	}
	return
}

type elements struct {
	Ok bool `json:"ok"`
}

// Parameters defines any query parameters used in this module.
type Parameters struct {
}

func newParameters() *Parameters {
	return &Parameters{}
}
<commit_msg>dispatch: generate warning if messages are being dropped<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]

// Package dispatch implements alert dispatching for the agent as a module.
// Persistent modules which generate alerts will have these alerts forwarded
// to this module if the dispatch module is active. The dispatch module can then
// forward the alerts on based on it's configuration.
package dispatch /* import "mig.ninja/mig/modules/dispatch" */

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime"
	"time"

	"mig.ninja/mig"
	"mig.ninja/mig/modules"
)

type module struct {
}

// NewRun returns a new instance of a modules.Runner for this module.
func (m *module) NewRun() modules.Runner {
	return new(run)
}

func init() {
	modules.Register("dispatch", new(module))
}

type run struct {
	Parameters Parameters
	Results    modules.Result
}

func buildResults(e elements, r *modules.Result) (buf []byte, err error) {
	r.Success = true
	r.Elements = e
	r.FoundAnything = true
	buf, err = json.Marshal(r)
	return
}

var logChan chan string
var alertChan chan string
var handlerErrChan chan error
var configChan chan modules.ConfigParams

// We keep a copy of the agent environment, hostname, and tags here so we can
// use them in the dispatched output
var env mig.AgentEnv
var tags map[string]string
var agtHostname string

// messageBuf is a queue used to store incoming messages, and is drained by
// runDispatch
var messageBuf chan string

// lastDrop stores the last time a message was dropped by the dispatch module
var lastDrop time.Time
var dropCounter int

// Dispatch record describes the formatting of JSON data submitted from the dispatch
// module.
type DispatchRecord struct {
	Hostname string            `json:"hostname"`
	Env      mig.AgentEnv      `json:"environment"`
	Tags     map[string]string `json:"tags"`
	Details  interface{}       `json:"details"`
}

func (d *DispatchRecord) fromString(msg string) error {
	d.Hostname = agtHostname
	d.Env = env
	d.Tags = tags
	err := json.Unmarshal([]byte(msg), &d.Details)
	return err
}

func moduleMain() {
	var cfg config

	incfg := <-configChan
	agtHostname = incfg.Hostname
	tags = incfg.Tags
	buf, err := json.Marshal(incfg.Env)
	if err != nil {
		handlerErrChan <- err
		return
	}
	err = json.Unmarshal(buf, &env)
	if err != nil {
		handlerErrChan <- err
		return
	}
	buf, err = json.Marshal(incfg.Config)
	if err != nil {
		handlerErrChan <- err
		return
	}
	err = json.Unmarshal(buf, &cfg)
	if err != nil {
		handlerErrChan <- err
		return
	}
	logChan <- "module received configuration"
	if cfg.Dispatch.ChannelSize == 0 {
		cfg.Dispatch.ChannelSize = 1024
		logChan <- "channel size not specified, defaulting to 1024"
	}

	messageBuf = make(chan string, cfg.Dispatch.ChannelSize)

	// Register the dispatch function, which will be called when the module
	// recieves an alert message from the agent.
	modules.RegisterDispatchFunction(dispatchIn)

	// Start the dispatch dequeue function.
	err = runDispatch(cfg)
	if err != nil {
		handlerErrChan <- err
		return
	}
}

func dispatchIn(msg string) {
	select {
	case messageBuf <- msg:
	default:
		dropCounter++
		// If we can't queue the message it is just dropped
		now := time.Now()
		if now.After(lastDrop.Add(time.Duration(time.Minute * 5))) {
			logChan <- fmt.Sprintf("warning, dispatch module dropping messages "+
				"(buffer full, %v dropped since last warning)", dropCounter)
			lastDrop = now
			dropCounter = 0
		}
	}
}

func runDispatch(cfg config) error {
	var (
		httpClient http.Client
		dr         DispatchRecord
	)

	for {
		msg := <-messageBuf
		dr.fromString(msg)
		buf, err := json.Marshal(dr)
		if err != nil {
			logChan <- fmt.Sprintf("create dispatch record: %v", err)
			continue
		}
		b := bytes.NewBuffer(buf)
		resp, err := httpClient.Post(cfg.Dispatch.HTTPURL, "application/json", b)
		if err != nil {
			logChan <- fmt.Sprintf("http post: %v", err)
			continue
		}
		resp.Body.Close()
	}
	return nil
}

func requestHandler(p interface{}) (ret string) {
	var results modules.Result
	defer func() {
		if e := recover(); e != nil {
			results.Errors = append(results.Errors, fmt.Sprintf("%v", e))
			results.Success = false
			err, _ := json.Marshal(results)
			ret = string(err)
			return
		}
	}()
	e := elements{Ok: true}
	resp, err := buildResults(e, &results)
	if err != nil {
		panic(err)
	}
	return string(resp)
}

type config struct {
	Dispatch struct {
		HTTPURL     string `json:"httpurl"`
		ChannelSize int    `json:"channelsize"`
	} `json:"dispatch"`
}

// PersistModConfig returns a new configuration structure for this module.
func (r *run) PersistModConfig() interface{} {
	return &config{}
}

// RunPersist is the entry point for persistent execution of the module.
func (r *run) RunPersist(in modules.ModuleReader, out modules.ModuleWriter) {
	alertChan = make(chan string, 64)
	logChan = make(chan string, 64)
	regChan := make(chan string, 64)
	handlerErrChan = make(chan error, 64)
	configChan = make(chan modules.ConfigParams, 1)

	go moduleMain()
	l, spec, err := modules.GetPersistListener("dispatch")
	if err != nil {
		handlerErrChan <- err
	} else {
		regChan <- spec
	}
	go modules.HandlePersistRequest(l, requestHandler, handlerErrChan)
	modules.DefaultPersistHandlers(in, out, logChan, handlerErrChan, regChan,
		alertChan, configChan)
}

// Run is the entry point for a standard (e.g., query) based invocation of the module.
func (r *run) Run(in modules.ModuleReader) (resStr string) {
	defer func() {
		if e := recover(); e != nil {
			// return error in json
			r.Results.Errors = append(r.Results.Errors, fmt.Sprintf("%v", e))
			r.Results.Success = false
			err, _ := json.Marshal(r.Results)
			resStr = string(err)
			return
		}
	}()
	runtime.GOMAXPROCS(1)
	sockspec, err := modules.ReadPersistInputParameters(in, &r.Parameters)
	if err != nil {
		panic(err)
	}
	err = r.ValidateParameters()
	if err != nil {
		panic(err)
	}
	resStr = modules.SendPersistRequest(r.Parameters, sockspec)
	return
}

// ValidateParameters validates the parameters set in the runner for the module.
func (r *run) ValidateParameters() (err error) {
	return
}

// PrintResults returns the results of a query of this module in human readable form.
func (r *run) PrintResults(result modules.Result, foundOnly bool) (prints []string, err error) {
	var (
		elem elements
	)

	err = result.GetElements(&elem)
	if err != nil {
		panic(err)
	}
	resStr := fmt.Sprintf("ok:%v", elem.Ok)
	prints = append(prints, resStr)
	if !foundOnly {
		for _, we := range result.Errors {
			prints = append(prints, we)
		}
	}
	return
}

type elements struct {
	Ok bool `json:"ok"`
}

// Parameters defines any query parameters used in this module.
type Parameters struct {
}

func newParameters() *Parameters {
	return &Parameters{}
}
<|endoftext|>