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

package flate

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

type deflateTest struct {
	in    []byte
	level int
	out   []byte
}

type deflateInflateTest struct {
	in []byte
}

type reverseBitsTest struct {
	in       uint16
	bitCount uint8
	out      uint16
}

var deflateTests = []*deflateTest{
	&deflateTest{[]byte{}, 0, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, -1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, DefaultCompression, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 4, []byte{18, 4, 4, 0, 0, 255, 255}},

	&deflateTest{[]byte{0x11}, 0, []byte{0, 1, 0, 254, 255, 17, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 0, []byte{0, 2, 0, 253, 255, 17, 18, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 0,
		[]byte{0, 8, 0, 247, 255, 17, 17, 17, 17, 17, 17, 17, 17, 1, 0, 0, 255, 255},
	},
	&deflateTest{[]byte{}, 1, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 1, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 1, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
	&deflateTest{[]byte{}, 9, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 9, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 9, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 9, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
}

var deflateInflateTests = []*deflateInflateTest{
	&deflateInflateTest{[]byte{}},
	&deflateInflateTest{[]byte{0x11}},
	&deflateInflateTest{[]byte{0x11, 0x12}},
	&deflateInflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}},
	&deflateInflateTest{[]byte{0x11, 0x10, 0x13, 0x41, 0x21, 0x21, 0x41, 0x13, 0x87, 0x78, 0x13}},
	&deflateInflateTest{getLargeDataChunk()},
}

var reverseBitsTests = []*reverseBitsTest{
	&reverseBitsTest{1, 1, 1},
	&reverseBitsTest{1, 2, 2},
	&reverseBitsTest{1, 3, 4},
	&reverseBitsTest{1, 4, 8},
	&reverseBitsTest{1, 5, 16},
	&reverseBitsTest{17, 5, 17},
	&reverseBitsTest{257, 9, 257},
	&reverseBitsTest{29, 5, 23},
}

func getLargeDataChunk() []byte {
	result := make([]byte, 100000)
	for i := range result {
		result[i] = byte(int64(i) * int64(i) & 0xFF)
	}
	return result
}

func TestDeflate(t *testing.T) {
	for _, h := range deflateTests {
		buffer := bytes.NewBuffer(nil)
		w := NewWriter(buffer, h.level)
		w.Write(h.in)
		w.Close()
		if bytes.Compare(buffer.Bytes(), h.out) != 0 {
			t.Errorf("buffer is wrong; level = %v, buffer.Bytes() = %v, expected output = %v",
				h.level, buffer.Bytes(), h.out)
		}
	}
}

type syncBuffer struct {
	buf    bytes.Buffer
	mu     sync.RWMutex
	closed bool
	ready  chan bool
}

func newSyncBuffer() *syncBuffer {
	return &syncBuffer{ready: make(chan bool, 1)}
}

func (b *syncBuffer) Read(p []byte) (n int, err os.Error) {
	for {
		b.mu.RLock()
		n, err = b.buf.Read(p)
		b.mu.RUnlock()
		if n > 0 || b.closed {
			return
		}
		<-b.ready
	}
	panic("unreachable")
}

func (b *syncBuffer) signal() {
	select {
	case b.ready <- true:
	default:
	}
}

func (b *syncBuffer) Write(p []byte) (n int, err os.Error) {
	n, err = b.buf.Write(p)
	b.signal()
	return
}

func (b *syncBuffer) WriteMode() {
	b.mu.Lock()
}

func (b *syncBuffer) ReadMode() {
	b.mu.Unlock()
	b.signal()
}

func (b *syncBuffer) Close() os.Error {
	b.closed = true
	b.signal()
	return nil
}

func testSync(t *testing.T, level int, input []byte, name string) {
	if len(input) == 0 {
		return
	}

	t.Logf("--testSync %d, %d, %s", level, len(input), name)
	buf := newSyncBuffer()
	buf1 := new(bytes.Buffer)
	buf.WriteMode()
	w := NewWriter(io.MultiWriter(buf, buf1), level)
	r := NewReader(buf)

	// Write half the input and read back.
	for i := 0; i < 2; i++ {
		var lo, hi int
		if i == 0 {
			lo, hi = 0, (len(input)+1)/2
		} else {
			lo, hi = (len(input)+1)/2, len(input)
		}
		t.Logf("#%d: write %d-%d", i, lo, hi)
		if _, err := w.Write(input[lo:hi]); err != nil {
			t.Errorf("testSync: write: %v", err)
			return
		}
		if i == 0 {
			if err := w.Flush(); err != nil {
				t.Errorf("testSync: flush: %v", err)
				return
			}
		} else {
			if err := w.Close(); err != nil {
				t.Errorf("testSync: close: %v", err)
			}
		}
		buf.ReadMode()
		out := make([]byte, hi-lo+1)
		m, err := io.ReadAtLeast(r, out, hi-lo)
		t.Logf("#%d: read %d", i, m)
		if m != hi-lo || err != nil {
			t.Errorf("testSync/%d (%d, %d, %s): read %d: %d, %v (%d left)", i, level, len(input), name, hi-lo, m, err, buf.buf.Len())
			return
		}
		if !bytes.Equal(input[lo:hi], out[:hi-lo]) {
			t.Errorf("testSync/%d: read wrong bytes: %x vs %x", i, input[lo:hi], out[:hi-lo])
			return
		}
		if i == 0 && buf.buf.Len() != 0 {
			t.Errorf("testSync/%d (%d, %d, %s): extra data after %d", i, level, len(input), name, hi-lo)
		}
		buf.WriteMode()
	}
	buf.ReadMode()
	out := make([]byte, 10)
	if n, err := r.Read(out); n > 0 || err != os.EOF {
		t.Errorf("testSync (%d, %d, %s): final Read: %d, %v (hex: %x)", level, len(input), name, n, err, out[0:n])
	}
	if buf.buf.Len() != 0 {
		t.Errorf("testSync (%d, %d, %s): extra data at end", level, len(input), name)
	}
	r.Close()

	// stream should work for ordinary reader too
	r = NewReader(buf1)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("testSync: read: %s", err)
		return
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("testSync: decompress(compress(data)) != data: level=%d input=%s", level, name)
	}
}

func testToFromWithLevel(t *testing.T, level int, input []byte, name string) os.Error {
	buffer := bytes.NewBuffer(nil)
	w := NewWriter(buffer, level)
	w.Write(input)
	w.Close()
	r := NewReader(buffer)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("read: %s", err)
		return err
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("decompress(compress(data)) != data: level=%d input=%s", level, name)
	}

	testSync(t, level, input, name)
	return nil
}

func testToFrom(t *testing.T, input []byte, name string) {
	for i := 0; i < 10; i++ {
		testToFromWithLevel(t, i, input, name)
	}
}

func TestDeflateInflate(t *testing.T) {
	for i, h := range deflateInflateTests {
		testToFrom(t, h.in, fmt.Sprintf("#%d", i))
	}
}

func TestReverseBits(t *testing.T) {
	for _, h := range reverseBitsTests {
		if v := reverseBits(h.in, h.bitCount); v != h.out {
			t.Errorf("reverseBits(%v,%v) = %v, want %v",
				h.in, h.bitCount, v, h.out)
		}
	}
}

func TestDeflateInflateString(t *testing.T) {
	gold, err := ioutil.ReadFile("../testdata/e.txt")
	if err != nil {
		t.Error(err)
	}
	testToFromWithLevel(t, 1, gold, "2.718281828...")
}
<commit_msg>compress/flate: fix test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package flate

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

type deflateTest struct {
	in    []byte
	level int
	out   []byte
}

type deflateInflateTest struct {
	in []byte
}

type reverseBitsTest struct {
	in       uint16
	bitCount uint8
	out      uint16
}

var deflateTests = []*deflateTest{
	&deflateTest{[]byte{}, 0, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, -1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, DefaultCompression, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 4, []byte{18, 4, 4, 0, 0, 255, 255}},

	&deflateTest{[]byte{0x11}, 0, []byte{0, 1, 0, 254, 255, 17, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 0, []byte{0, 2, 0, 253, 255, 17, 18, 1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 0,
		[]byte{0, 8, 0, 247, 255, 17, 17, 17, 17, 17, 17, 17, 17, 1, 0, 0, 255, 255},
	},
	&deflateTest{[]byte{}, 1, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 1, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 1, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 1, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
	&deflateTest{[]byte{}, 9, []byte{1, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11}, 9, []byte{18, 4, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x12}, 9, []byte{18, 20, 2, 4, 0, 0, 255, 255}},
	&deflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}, 9, []byte{18, 132, 2, 64, 0, 0, 0, 255, 255}},
}

var deflateInflateTests = []*deflateInflateTest{
	&deflateInflateTest{[]byte{}},
	&deflateInflateTest{[]byte{0x11}},
	&deflateInflateTest{[]byte{0x11, 0x12}},
	&deflateInflateTest{[]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11}},
	&deflateInflateTest{[]byte{0x11, 0x10, 0x13, 0x41, 0x21, 0x21, 0x41, 0x13, 0x87, 0x78, 0x13}},
	&deflateInflateTest{getLargeDataChunk()},
}

var reverseBitsTests = []*reverseBitsTest{
	&reverseBitsTest{1, 1, 1},
	&reverseBitsTest{1, 2, 2},
	&reverseBitsTest{1, 3, 4},
	&reverseBitsTest{1, 4, 8},
	&reverseBitsTest{1, 5, 16},
	&reverseBitsTest{17, 5, 17},
	&reverseBitsTest{257, 9, 257},
	&reverseBitsTest{29, 5, 23},
}

func getLargeDataChunk() []byte {
	result := make([]byte, 100000)
	for i := range result {
		result[i] = byte(int64(i) * int64(i) & 0xFF)
	}
	return result
}

func TestDeflate(t *testing.T) {
	for _, h := range deflateTests {
		buffer := bytes.NewBuffer(nil)
		w := NewWriter(buffer, h.level)
		w.Write(h.in)
		w.Close()
		if bytes.Compare(buffer.Bytes(), h.out) != 0 {
			t.Errorf("buffer is wrong; level = %v, buffer.Bytes() = %v, expected output = %v",
				h.level, buffer.Bytes(), h.out)
		}
	}
}

type syncBuffer struct {
	buf    bytes.Buffer
	mu     sync.RWMutex
	closed bool
	ready  chan bool
}

func newSyncBuffer() *syncBuffer {
	return &syncBuffer{ready: make(chan bool, 1)}
}

func (b *syncBuffer) Read(p []byte) (n int, err os.Error) {
	for {
		b.mu.RLock()
		n, err = b.buf.Read(p)
		b.mu.RUnlock()
		if n > 0 || b.closed {
			return
		}
		<-b.ready
	}
	panic("unreachable")
}

func (b *syncBuffer) signal() {
	select {
	case b.ready <- true:
	default:
	}
}

func (b *syncBuffer) Write(p []byte) (n int, err os.Error) {
	n, err = b.buf.Write(p)
	b.signal()
	return
}

func (b *syncBuffer) WriteMode() {
	b.mu.Lock()
}

func (b *syncBuffer) ReadMode() {
	b.mu.Unlock()
	b.signal()
}

func (b *syncBuffer) Close() os.Error {
	b.closed = true
	b.signal()
	return nil
}

func testSync(t *testing.T, level int, input []byte, name string) {
	if len(input) == 0 {
		return
	}

	t.Logf("--testSync %d, %d, %s", level, len(input), name)
	buf := newSyncBuffer()
	buf1 := new(bytes.Buffer)
	buf.WriteMode()
	w := NewWriter(io.MultiWriter(buf, buf1), level)
	r := NewReader(buf)

	// Write half the input and read back.
	for i := 0; i < 2; i++ {
		var lo, hi int
		if i == 0 {
			lo, hi = 0, (len(input)+1)/2
		} else {
			lo, hi = (len(input)+1)/2, len(input)
		}
		t.Logf("#%d: write %d-%d", i, lo, hi)
		if _, err := w.Write(input[lo:hi]); err != nil {
			t.Errorf("testSync: write: %v", err)
			return
		}
		if i == 0 {
			if err := w.Flush(); err != nil {
				t.Errorf("testSync: flush: %v", err)
				return
			}
		} else {
			if err := w.Close(); err != nil {
				t.Errorf("testSync: close: %v", err)
			}
		}
		buf.ReadMode()
		out := make([]byte, hi-lo+1)
		m, err := io.ReadAtLeast(r, out, hi-lo)
		t.Logf("#%d: read %d", i, m)
		if m != hi-lo || err != nil {
			t.Errorf("testSync/%d (%d, %d, %s): read %d: %d, %v (%d left)", i, level, len(input), name, hi-lo, m, err, buf.buf.Len())
			return
		}
		if !bytes.Equal(input[lo:hi], out[:hi-lo]) {
			t.Errorf("testSync/%d: read wrong bytes: %x vs %x", i, input[lo:hi], out[:hi-lo])
			return
		}
		// This test originally checked that after reading
		// the first half of the input, there was nothing left
		// in the read buffer (buf.buf.Len() != 0) but that is
		// not necessarily the case: the write Flush may emit
		// some extra framing bits that are not necessary
		// to process to obtain the first half of the uncompressed
		// data.  The test ran correctly most of the time, because
		// the background goroutine had usually read even
		// those extra bits by now, but it's not a useful thing to
		// check.
		buf.WriteMode()
	}
	buf.ReadMode()
	out := make([]byte, 10)
	if n, err := r.Read(out); n > 0 || err != os.EOF {
		t.Errorf("testSync (%d, %d, %s): final Read: %d, %v (hex: %x)", level, len(input), name, n, err, out[0:n])
	}
	if buf.buf.Len() != 0 {
		t.Errorf("testSync (%d, %d, %s): extra data at end", level, len(input), name)
	}
	r.Close()

	// stream should work for ordinary reader too
	r = NewReader(buf1)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("testSync: read: %s", err)
		return
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("testSync: decompress(compress(data)) != data: level=%d input=%s", level, name)
	}
}

func testToFromWithLevel(t *testing.T, level int, input []byte, name string) os.Error {
	buffer := bytes.NewBuffer(nil)
	w := NewWriter(buffer, level)
	w.Write(input)
	w.Close()
	r := NewReader(buffer)
	out, err := ioutil.ReadAll(r)
	if err != nil {
		t.Errorf("read: %s", err)
		return err
	}
	r.Close()
	if !bytes.Equal(input, out) {
		t.Errorf("decompress(compress(data)) != data: level=%d input=%s", level, name)
	}

	testSync(t, level, input, name)
	return nil
}

func testToFrom(t *testing.T, input []byte, name string) {
	for i := 0; i < 10; i++ {
		testToFromWithLevel(t, i, input, name)
	}
}

func TestDeflateInflate(t *testing.T) {
	for i, h := range deflateInflateTests {
		testToFrom(t, h.in, fmt.Sprintf("#%d", i))
	}
}

func TestReverseBits(t *testing.T) {
	for _, h := range reverseBitsTests {
		if v := reverseBits(h.in, h.bitCount); v != h.out {
			t.Errorf("reverseBits(%v,%v) = %v, want %v",
				h.in, h.bitCount, v, h.out)
		}
	}
}

func TestDeflateInflateString(t *testing.T) {
	gold, err := ioutil.ReadFile("../testdata/e.txt")
	if err != nil {
		t.Error(err)
	}
	testToFromWithLevel(t, 1, gold, "2.718281828...")
}
<|endoftext|>


<commit_before>package seccomp

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
	"golang.org/x/sys/unix"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path"
	"syscall"
)

// #include "sys/ptrace.h"
import "C"

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

type SystemCallArgs []int

func Tracer() {

	var train = false
	var cmd string
	var cmdArgs []string
	var p *oz.Profile

	var noprofile = flag.Bool("train", false, "Training mode")
	var debug = flag.Bool("debug", false, "Debug")
	var appendpolicy = flag.Bool("append", false, "Append to existing policy if exists")
	var trainingoutput = flag.String("output", "", "Training policy output file")

	flag.Parse()

	var args = flag.Args()

	if *noprofile == true {
		train = true

		// TODO: remove hardcoded path and read prefix from /etc/oz.conf

		cmd = "/usr/bin/oz-seccomp"
		cmdArgs = append([]string{"-mode=train"}, args...)
	} else {
		p = new(oz.Profile)
		if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
			log.Error("unable to decode profile data: %v", err)
			os.Exit(1)
		}
		if p.Seccomp.Mode == oz.PROFILE_SECCOMP_TRAIN {
			train = true
		}
		*debug = p.Seccomp.Debug
		cmd = args[0]
		cmdArgs = args[1:]
	}

	var cpid = 0
	done := false

	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	if *noprofile == false {

		pi, err := c.StdinPipe()
		if err != nil {
			fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
			os.Exit(1)
		}
		jdata, err := json.Marshal(p)
		if err != nil {
			fmt.Errorf("Unable to marshal seccomp state: %+v", err)
			os.Exit(1)
		}
		io.Copy(pi, bytes.NewBuffer(jdata))
		pi.Close()
	}
	children := make(map[int]bool)
	renderFunctions := getRenderingFunctions()

	trainingset := make(map[int]bool)
	trainingargs := make(map[int]map[int][]uint)

	if err := c.Start(); err == nil {
		cpid = c.Process.Pid
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4) err:%v pid:%i", err, pid)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			pflags := unix.PTRACE_O_TRACESECCOMP
			pflags |= unix.PTRACE_O_TRACEFORK
			pflags |= unix.PTRACE_O_TRACEVFORK
			pflags |= unix.PTRACE_O_TRACECLONE
			pflags |= C.PTRACE_O_EXITKILL

			syscall.PtraceSetOptions(pid, pflags)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4) err:%v pid:%i children:%v\n", err, pid, children)
				done = true
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if s.Signaled() == true {
				log.Error("Pid signaled, pid: %v signal: %v", pid, s)
				delete(children, pid)
				continue
			}
			switch uint32(s) >> 8 {

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8):
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)

				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}

				systemcall, err := syscallByNum(getSyscallNumber(regs))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}

				/* Render the system call invocation */

				r := getSyscallRegisterArgs(regs)
				call := ""

				if train == true {
					trainingset[getSyscallNumber(regs)] = true
					if systemcall.captureArgs != nil {
						for c, i := range systemcall.captureArgs {
							if i == 1 {
								if trainingargs[getSyscallNumber(regs)] == nil {
									trainingargs[getSyscallNumber(regs)] = make(map[int][]uint)
								}
								if contains(trainingargs[getSyscallNumber(regs)][c], uint(r[c])) == false {
									trainingargs[getSyscallNumber(regs)][c] = append(trainingargs[getSyscallNumber(regs)][c], uint(r[c]))
								}
							}
						}
					}
				}

				if f, ok := renderFunctions[getSyscallNumber(regs)]; ok {
					call, err = f(pid, r)
					if err != nil {
						log.Info("%v", err)
						continue
					}
					if *debug == true {
						call += "\n  " + renderSyscallBasic(pid, systemcall, regs)
					}
				} else {
					call = renderSyscallBasic(pid, systemcall, regs)
				}

				log.Info("seccomp hit on sandbox pid %v (%v) syscall %v (%v):\n  %s", pid, getProcessCmdLine(pid), systemcall.name, systemcall.num, call)
				continue

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXIT << 8):
				if *debug == true {
					log.Error("Ptrace exit event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				delete(children, pid)
				continue

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_CLONE << 8):
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				if *debug == true {
					log.Error("Ptrace clone event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_FORK << 8):
				if *debug == true {
					log.Error("PTrace fork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK << 8):
				if *debug == true {
					log.Error("Ptrace vfork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK_DONE << 8):
				if *debug == true {
					log.Error("Ptrace vfork done event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true

				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXEC << 8):
				if *debug == true {
					log.Error("Ptrace exec event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_STOP << 8):
				if *debug == true {
					log.Error("Ptrace stop event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP):
				if *debug == true {
					log.Error("SIGTRAP detected in pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGCHLD):
				if *debug == true {
					log.Error("SIGCHLD detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGSTOP):
				if *debug == true {
					log.Error("SIGSTOP detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGSEGV):
				if *debug == true {
					log.Error("SIGSEGV detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				err = syscall.Kill(pid, 9)
				if err != nil {
					log.Error("kill: %v", err)
					os.Exit(1)
				}
				delete(children, pid)
				continue
			default:
				y := s.StopSignal()
				if *debug == true {
					log.Error("Child stopped for unknown reasons pid %v status %v signal %i (%s)", pid, s, y, getProcessCmdLine(pid))
				}
				continue
			}
		}

		if train == true {
			var u *user.User
			var e error
			u, e = user.Current()
			var resolvedpath = ""

			if e != nil {
				log.Error("user.Current(): %v", e)
			}

			if *trainingoutput != "" {
				resolvedpath = *trainingoutput
			} else {
				if *noprofile == false {
					resolvedpath, e = fs.ResolvePathNoGlob(p.Seccomp.TrainOutput, -1, u)
					if e != nil {
						log.Error("resolveVars(): %v", e)
					}
				} else {
					s := fmt.Sprintf("${HOME}/%s-%d.seccomp", fname(os.Args[2]), cpid)
					resolvedpath, e = fs.ResolvePathNoGlob(s, -1, u)
				}
			}
			policyout := "execve:1\n"
			for call := range trainingset {
				done := false
				for c := range trainingargs {
					if c == call {
						for a, v := range trainingargs[c] {
							sc, _ := syscallByNum(call)
							policyout += fmt.Sprintf("%s:%s\n", sc.name, genArgs(uint(a), (v)))
							done = true
						}
					}
				}
				if done == false {
					sc, _ := syscallByNum(call)
					policyout += fmt.Sprintf("%s:1\n", sc.name)
				}
			}
			if *appendpolicy == true {
				log.Error("Not yet implemented.")
			}
			f, err := os.OpenFile(resolvedpath, os.O_CREATE|os.O_RDWR, 0600)
			if err == nil {
				_, err := f.WriteString(policyout)
				if err != nil {
					log.Error("Error writing policy file: %v", err)
				}
				err = f.Close()
				if err != nil {
					log.Error("Error closing policy file: %v", err)
				}
			} else {
				log.Error("Error opening policy file \"%s\": %v", resolvedpath, err)
			}
		}
	}
}

func genArgs(a uint, vals []uint) string {
	s := ""
	for idx, x := range vals {
		s += fmt.Sprintf(" arg%d == %d ", a, x)
		if idx < len(vals)-1 {
			s += "||"
		}
	}
	return s
}

func contains(slice []uint, val uint) bool {
	for _, x := range slice {
		if val == x {
			return true
		}
	}
	return false
}
func fname(p string) string {
	_, fname := path.Split(p)
	return fname
}
<commit_msg>Use proper prefix path for finding oz binaries<commit_after>package seccomp

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path"
	"syscall"

	"golang.org/x/sys/unix"

	"github.com/subgraph/oz"
	"github.com/subgraph/oz/fs"
)

// #include "sys/ptrace.h"
import "C"

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

type SystemCallArgs []int

func Tracer() {

	var train = false
	var cmd string
	var cmdArgs []string
	var p *oz.Profile

	var noprofile = flag.Bool("train", false, "Training mode")
	var debug = flag.Bool("debug", false, "Debug")
	var appendpolicy = flag.Bool("append", false, "Append to existing policy if exists")
	var trainingoutput = flag.String("output", "", "Training policy output file")

	flag.Parse()

	var args = flag.Args()

	OzConfig, err := oz.LoadConfig(oz.DefaultConfigPath)
	if err != nil {
		log.Error("unable to load oz config")
		os.Exit(1)
	}

	if *noprofile == true {
		train = true

		// TODO: remove hardcoded path and read prefix from /etc/oz.conf

		cmd = path.Join(OzConfig.PrefixPath, "bin", "oz-seccomp")
		cmdArgs = append([]string{"-mode=train"}, args...)
	} else {
		p = new(oz.Profile)
		if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
			log.Error("unable to decode profile data: %v", err)
			os.Exit(1)
		}
		if p.Seccomp.Mode == oz.PROFILE_SECCOMP_TRAIN {
			train = true
		}
		*debug = p.Seccomp.Debug
		cmd = args[0]
		cmdArgs = args[1:]
	}

	var cpid = 0
	done := false

	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	if *noprofile == false {

		pi, err := c.StdinPipe()
		if err != nil {
			fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
			os.Exit(1)
		}
		jdata, err := json.Marshal(p)
		if err != nil {
			fmt.Errorf("Unable to marshal seccomp state: %+v", err)
			os.Exit(1)
		}
		io.Copy(pi, bytes.NewBuffer(jdata))
		pi.Close()
	}
	children := make(map[int]bool)
	renderFunctions := getRenderingFunctions()

	trainingset := make(map[int]bool)
	trainingargs := make(map[int]map[int][]uint)

	if err := c.Start(); err == nil {
		cpid = c.Process.Pid
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4) err:%v pid:%i", err, pid)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			pflags := unix.PTRACE_O_TRACESECCOMP
			pflags |= unix.PTRACE_O_TRACEFORK
			pflags |= unix.PTRACE_O_TRACEVFORK
			pflags |= unix.PTRACE_O_TRACECLONE
			pflags |= C.PTRACE_O_EXITKILL

			syscall.PtraceSetOptions(pid, pflags)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4) err:%v pid:%i children:%v\n", err, pid, children)
				done = true
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if s.Signaled() == true {
				log.Error("Pid signaled, pid: %v signal: %v", pid, s)
				delete(children, pid)
				continue
			}
			switch uint32(s) >> 8 {

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8):
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)

				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}

				systemcall, err := syscallByNum(getSyscallNumber(regs))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}

				/* Render the system call invocation */

				r := getSyscallRegisterArgs(regs)
				call := ""

				if train == true {
					trainingset[getSyscallNumber(regs)] = true
					if systemcall.captureArgs != nil {
						for c, i := range systemcall.captureArgs {
							if i == 1 {
								if trainingargs[getSyscallNumber(regs)] == nil {
									trainingargs[getSyscallNumber(regs)] = make(map[int][]uint)
								}
								if contains(trainingargs[getSyscallNumber(regs)][c], uint(r[c])) == false {
									trainingargs[getSyscallNumber(regs)][c] = append(trainingargs[getSyscallNumber(regs)][c], uint(r[c]))
								}
							}
						}
					}
				}

				if f, ok := renderFunctions[getSyscallNumber(regs)]; ok {
					call, err = f(pid, r)
					if err != nil {
						log.Info("%v", err)
						continue
					}
					if *debug == true {
						call += "\n  " + renderSyscallBasic(pid, systemcall, regs)
					}
				} else {
					call = renderSyscallBasic(pid, systemcall, regs)
				}

				log.Info("seccomp hit on sandbox pid %v (%v) syscall %v (%v):\n  %s", pid, getProcessCmdLine(pid), systemcall.name, systemcall.num, call)
				continue

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXIT << 8):
				if *debug == true {
					log.Error("Ptrace exit event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				delete(children, pid)
				continue

			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_CLONE << 8):
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				if *debug == true {
					log.Error("Ptrace clone event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_FORK << 8):
				if *debug == true {
					log.Error("PTrace fork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK << 8):
				if *debug == true {
					log.Error("Ptrace vfork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK_DONE << 8):
				if *debug == true {
					log.Error("Ptrace vfork done event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				newpid, err := syscall.PtraceGetEventMsg(pid)
				if err != nil {
					log.Error("PTrace event message retrieval failed: %v", err)
				}
				children[int(newpid)] = true

				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXEC << 8):
				if *debug == true {
					log.Error("Ptrace exec event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_STOP << 8):
				if *debug == true {
					log.Error("Ptrace stop event detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGTRAP):
				if *debug == true {
					log.Error("SIGTRAP detected in pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGCHLD):
				if *debug == true {
					log.Error("SIGCHLD detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGSTOP):
				if *debug == true {
					log.Error("SIGSTOP detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				continue
			case uint32(unix.SIGSEGV):
				if *debug == true {
					log.Error("SIGSEGV detected pid %v (%s)", pid, getProcessCmdLine(pid))
				}
				err = syscall.Kill(pid, 9)
				if err != nil {
					log.Error("kill: %v", err)
					os.Exit(1)
				}
				delete(children, pid)
				continue
			default:
				y := s.StopSignal()
				if *debug == true {
					log.Error("Child stopped for unknown reasons pid %v status %v signal %i (%s)", pid, s, y, getProcessCmdLine(pid))
				}
				continue
			}
		}

		if train == true {
			var u *user.User
			var e error
			u, e = user.Current()
			var resolvedpath = ""

			if e != nil {
				log.Error("user.Current(): %v", e)
			}

			if *trainingoutput != "" {
				resolvedpath = *trainingoutput
			} else {
				if *noprofile == false {
					resolvedpath, e = fs.ResolvePathNoGlob(p.Seccomp.TrainOutput, -1, u)
					if e != nil {
						log.Error("resolveVars(): %v", e)
					}
				} else {
					s := fmt.Sprintf("${HOME}/%s-%d.seccomp", fname(os.Args[2]), cpid)
					resolvedpath, e = fs.ResolvePathNoGlob(s, -1, u)
				}
			}
			policyout := "execve:1\n"
			for call := range trainingset {
				done := false
				for c := range trainingargs {
					if c == call {
						for a, v := range trainingargs[c] {
							sc, _ := syscallByNum(call)
							policyout += fmt.Sprintf("%s:%s\n", sc.name, genArgs(uint(a), (v)))
							done = true
						}
					}
				}
				if done == false {
					sc, _ := syscallByNum(call)
					policyout += fmt.Sprintf("%s:1\n", sc.name)
				}
			}
			if *appendpolicy == true {
				log.Error("Not yet implemented.")
			}
			f, err := os.OpenFile(resolvedpath, os.O_CREATE|os.O_RDWR, 0600)
			if err == nil {
				_, err := f.WriteString(policyout)
				if err != nil {
					log.Error("Error writing policy file: %v", err)
				}
				err = f.Close()
				if err != nil {
					log.Error("Error closing policy file: %v", err)
				}
			} else {
				log.Error("Error opening policy file \"%s\": %v", resolvedpath, err)
			}
		}
	}
}

func genArgs(a uint, vals []uint) string {
	s := ""
	for idx, x := range vals {
		s += fmt.Sprintf(" arg%d == %d ", a, x)
		if idx < len(vals)-1 {
			s += "||"
		}
	}
	return s
}

func contains(slice []uint, val uint) bool {
	for _, x := range slice {
		if val == x {
			return true
		}
	}
	return false
}
func fname(p string) string {
	_, fname := path.Split(p)
	return fname
}
<|endoftext|>


<commit_before>package resource

import (
	"encoding/json"
	"fmt"

	"github.com/hashicorp/terraform-plugin-sdk/internal/addrs"
	"github.com/zclconf/go-cty/cty"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/internal/configs/hcl2shim"

	"github.com/hashicorp/terraform-plugin-sdk/internal/states"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
)

// shimState takes a new *states.State and reverts it to a legacy state for the provider ACC tests
func shimNewState(newState *states.State, providers map[string]terraform.ResourceProvider) (*terraform.State, error) {
	state := terraform.NewState()

	// in the odd case of a nil state, let the helper packages handle it
	if newState == nil {
		return nil, nil
	}

	for _, newMod := range newState.Modules {
		mod := state.AddModule(newMod.Addr)

		for name, out := range newMod.OutputValues {
			outputType := ""
			val := hcl2shim.ConfigValueFromHCL2(out.Value)
			ty := out.Value.Type()
			switch {
			case ty == cty.String:
				outputType = "string"
			case ty.IsTupleType() || ty.IsListType():
				outputType = "list"
			case ty.IsMapType():
				outputType = "map"
			}

			mod.Outputs[name] = &terraform.OutputState{
				Type:      outputType,
				Value:     val,
				Sensitive: out.Sensitive,
			}
		}

		for _, res := range newMod.Resources {
			resType := res.Addr.Type
			providerType := res.ProviderConfig.ProviderConfig.Type

			resource := getResource(providers, providerType, res.Addr)

			for key, i := range res.Instances {
				resState := &terraform.ResourceState{
					Type:     resType,
					Provider: res.ProviderConfig.String(),
				}

				// We should always have a Current instance here, but be safe about checking.
				if i.Current != nil {
					flatmap, err := shimmedAttributes(i.Current, resource)
					if err != nil {
						return nil, fmt.Errorf("error decoding state for %q: %s", resType, err)
					}

					var meta map[string]interface{}
					if i.Current.Private != nil {
						err := json.Unmarshal(i.Current.Private, &meta)
						if err != nil {
							return nil, err
						}
					}

					resState.Primary = &terraform.InstanceState{
						ID:         flatmap["id"],
						Attributes: flatmap,
						Tainted:    i.Current.Status == states.ObjectTainted,
						Meta:       meta,
					}

					if i.Current.SchemaVersion != 0 {
						resState.Primary.Meta = map[string]interface{}{
							"schema_version": i.Current.SchemaVersion,
						}
					}

					for _, dep := range i.Current.Dependencies {
						resState.Dependencies = append(resState.Dependencies, dep.String())
					}

					// convert the indexes to the old style flapmap indexes
					idx := ""
					switch key.(type) {
					case addrs.IntKey:
						// don't add numeric index values to resources with a count of 0
						if len(res.Instances) > 1 {
							idx = fmt.Sprintf(".%d", key)
						}
					case addrs.StringKey:
						idx = "." + key.String()
					}

					mod.Resources[res.Addr.String()+idx] = resState
				}

				// add any deposed instances
				for _, dep := range i.Deposed {
					flatmap, err := shimmedAttributes(dep, resource)
					if err != nil {
						return nil, fmt.Errorf("error decoding deposed state for %q: %s", resType, err)
					}

					var meta map[string]interface{}
					if dep.Private != nil {
						err := json.Unmarshal(dep.Private, &meta)
						if err != nil {
							return nil, err
						}
					}

					deposed := &terraform.InstanceState{
						ID:         flatmap["id"],
						Attributes: flatmap,
						Tainted:    dep.Status == states.ObjectTainted,
						Meta:       meta,
					}
					if dep.SchemaVersion != 0 {
						deposed.Meta = map[string]interface{}{
							"schema_version": dep.SchemaVersion,
						}
					}

					resState.Deposed = append(resState.Deposed, deposed)
				}
			}
		}
	}

	return state, nil
}

func getResource(providers map[string]terraform.ResourceProvider, providerName string, addr addrs.Resource) *schema.Resource {
	p := providers[providerName]
	if p == nil {
		panic(fmt.Sprintf("provider %q not found in test step", providerName))
	}

	// this is only for tests, so should only see schema.Providers
	provider := p.(*schema.Provider)

	switch addr.Mode {
	case addrs.ManagedResourceMode:
		resource := provider.ResourcesMap[addr.Type]
		if resource != nil {
			return resource
		}
	case addrs.DataResourceMode:
		resource := provider.DataSourcesMap[addr.Type]
		if resource != nil {
			return resource
		}
	}

	panic(fmt.Sprintf("resource %s not found in test step", addr.Type))
}

func shimmedAttributes(instance *states.ResourceInstanceObjectSrc, res *schema.Resource) (map[string]string, error) {
	flatmap := instance.AttrsFlat
	if flatmap != nil {
		return flatmap, nil
	}

	// if we have json attrs, they need to be decoded
	rio, err := instance.Decode(res.CoreConfigSchema().ImpliedType())
	if err != nil {
		return nil, err
	}

	instanceState, err := res.ShimInstanceStateFromValue(rio.Value)
	if err != nil {
		return nil, err
	}

	return instanceState.Attributes, nil
}
<commit_msg>don't lose track of private data in the ACC tests<commit_after>package resource

import (
	"encoding/json"
	"fmt"

	"github.com/hashicorp/terraform-plugin-sdk/internal/addrs"
	"github.com/zclconf/go-cty/cty"

	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/internal/configs/hcl2shim"

	"github.com/hashicorp/terraform-plugin-sdk/internal/states"
	"github.com/hashicorp/terraform-plugin-sdk/terraform"
)

// shimState takes a new *states.State and reverts it to a legacy state for the provider ACC tests
func shimNewState(newState *states.State, providers map[string]terraform.ResourceProvider) (*terraform.State, error) {
	state := terraform.NewState()

	// in the odd case of a nil state, let the helper packages handle it
	if newState == nil {
		return nil, nil
	}

	for _, newMod := range newState.Modules {
		mod := state.AddModule(newMod.Addr)

		for name, out := range newMod.OutputValues {
			outputType := ""
			val := hcl2shim.ConfigValueFromHCL2(out.Value)
			ty := out.Value.Type()
			switch {
			case ty == cty.String:
				outputType = "string"
			case ty.IsTupleType() || ty.IsListType():
				outputType = "list"
			case ty.IsMapType():
				outputType = "map"
			}

			mod.Outputs[name] = &terraform.OutputState{
				Type:      outputType,
				Value:     val,
				Sensitive: out.Sensitive,
			}
		}

		for _, res := range newMod.Resources {
			resType := res.Addr.Type
			providerType := res.ProviderConfig.ProviderConfig.Type

			resource := getResource(providers, providerType, res.Addr)

			for key, i := range res.Instances {
				resState := &terraform.ResourceState{
					Type:     resType,
					Provider: res.ProviderConfig.String(),
				}

				// We should always have a Current instance here, but be safe about checking.
				if i.Current != nil {
					flatmap, err := shimmedAttributes(i.Current, resource)
					if err != nil {
						return nil, fmt.Errorf("error decoding state for %q: %s", resType, err)
					}

					var meta map[string]interface{}
					if i.Current.Private != nil {
						err := json.Unmarshal(i.Current.Private, &meta)
						if err != nil {
							return nil, err
						}
					}

					resState.Primary = &terraform.InstanceState{
						ID:         flatmap["id"],
						Attributes: flatmap,
						Tainted:    i.Current.Status == states.ObjectTainted,
						Meta:       meta,
					}

					if i.Current.SchemaVersion != 0 {
						if resState.Primary.Meta == nil {
							resState.Primary.Meta = map[string]interface{}{}
						}
						resState.Primary.Meta["schema_version"] = i.Current.SchemaVersion
					}

					for _, dep := range i.Current.Dependencies {
						resState.Dependencies = append(resState.Dependencies, dep.String())
					}

					// convert the indexes to the old style flapmap indexes
					idx := ""
					switch key.(type) {
					case addrs.IntKey:
						// don't add numeric index values to resources with a count of 0
						if len(res.Instances) > 1 {
							idx = fmt.Sprintf(".%d", key)
						}
					case addrs.StringKey:
						idx = "." + key.String()
					}

					mod.Resources[res.Addr.String()+idx] = resState
				}

				// add any deposed instances
				for _, dep := range i.Deposed {
					flatmap, err := shimmedAttributes(dep, resource)
					if err != nil {
						return nil, fmt.Errorf("error decoding deposed state for %q: %s", resType, err)
					}

					var meta map[string]interface{}
					if dep.Private != nil {
						err := json.Unmarshal(dep.Private, &meta)
						if err != nil {
							return nil, err
						}
					}

					deposed := &terraform.InstanceState{
						ID:         flatmap["id"],
						Attributes: flatmap,
						Tainted:    dep.Status == states.ObjectTainted,
						Meta:       meta,
					}
					if dep.SchemaVersion != 0 {
						deposed.Meta = map[string]interface{}{
							"schema_version": dep.SchemaVersion,
						}
					}

					resState.Deposed = append(resState.Deposed, deposed)
				}
			}
		}
	}

	return state, nil
}

func getResource(providers map[string]terraform.ResourceProvider, providerName string, addr addrs.Resource) *schema.Resource {
	p := providers[providerName]
	if p == nil {
		panic(fmt.Sprintf("provider %q not found in test step", providerName))
	}

	// this is only for tests, so should only see schema.Providers
	provider := p.(*schema.Provider)

	switch addr.Mode {
	case addrs.ManagedResourceMode:
		resource := provider.ResourcesMap[addr.Type]
		if resource != nil {
			return resource
		}
	case addrs.DataResourceMode:
		resource := provider.DataSourcesMap[addr.Type]
		if resource != nil {
			return resource
		}
	}

	panic(fmt.Sprintf("resource %s not found in test step", addr.Type))
}

func shimmedAttributes(instance *states.ResourceInstanceObjectSrc, res *schema.Resource) (map[string]string, error) {
	flatmap := instance.AttrsFlat
	if flatmap != nil {
		return flatmap, nil
	}

	// if we have json attrs, they need to be decoded
	rio, err := instance.Decode(res.CoreConfigSchema().ImpliedType())
	if err != nil {
		return nil, err
	}

	instanceState, err := res.ShimInstanceStateFromValue(rio.Value)
	if err != nil {
		return nil, err
	}

	return instanceState.Attributes, nil
}
<|endoftext|>


<commit_before>package helpers

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/bwmarrin/discordgo"
)

const (
	LEFT_ARROW_EMOJI  = "⬅"
	RIGHT_ARROW_EMOJI = "➡"
	X_EMOJI           = "🇽"
)

// map of messageID to pagedEmbedMessage
var pagedEmbededMessages map[string]*pagedEmbedMessage

type pagedEmbedMessage struct {
	files           []*discordgo.File
	fullEmbed       *discordgo.MessageEmbed
	totalNumOfPages int
	currentPage     int
	fieldsPerPage   int
	color           int
	messageID       string
	channelID       string
	userId          string //user who triggered the message
	msgType         string // "image" - will cause the embed to page through the files instead of fields
}

func init() {
	pagedEmbededMessages = make(map[string]*pagedEmbedMessage)
}

// will remove all reactions from all paged embed messages.
//  mainly used on bot uninit to clean embeds
func RemoveReactionsFromPagedEmbeds() {
	// TODO: sync group? Without blocking bot shutdown
	for _, pagedEmbed := range pagedEmbededMessages {
		go func() {
			cache.GetSession().MessageReactionsRemoveAll(pagedEmbed.channelID, pagedEmbed.messageID)
		}()
	}
}

// GetPagedMessage will return the paged message if there is one, otherwill will return nil
func GetPagedMessage(messageID string) *pagedEmbedMessage {
	pagedMessaged, _ := pagedEmbededMessages[messageID]
	return pagedMessaged
}

// CreatePagedMessage creates the paged messages
func SendPagedMessage(msg *discordgo.Message, embed *discordgo.MessageEmbed, fieldsPerPage int) error {

	// if there aren't multiple fields to be paged through, or if the amount of fields is less than the requested fields per page
	//  just send a normal embed
	if len(embed.Fields) < 2 || len(embed.Fields) <= fieldsPerPage {
		SendEmbed(msg.ChannelID, embed)
		return nil
	}

	// fields per page can not be less than 1
	if fieldsPerPage < 1 {
		return errors.New("fieldsPerPage can not be less than 1")
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   fieldsPerPage,
		totalNumOfPages: int(math.Ceil(float64(len(embed.Fields)) / float64(fieldsPerPage))),
		userId:          msg.Author.ID,
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// SendPagedImageMessage creates the paged image messages
func SendPagedImageMessage(msg *discordgo.Message, msgSend *discordgo.MessageSend, fieldsPerPage int) error {
	if msgSend.Embed == nil {
		return errors.New("parameter msgSend must contain an embed")
	}

	// make sure the image url is set to the name of the first file incease it wasn't set
	msgSend.Embed.Image.URL = fmt.Sprintf("attachment://%s", msgSend.Files[0].Name)

	// check if there are multiple files, not just send it normally
	if len(msgSend.Files) < 2 {
		SendComplex(msg.ChannelID, msgSend)
		return nil
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       msgSend.Embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   fieldsPerPage,
		totalNumOfPages: len(msgSend.Files),
		files:           msgSend.Files,
		userId:          msg.Author.ID,
		msgType:         "image",
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// UpdateMessagePage will update the page based on the given direction and current page
//  reactions must be the left or right arrow
func (p *pagedEmbedMessage) UpdateMessagePage(reaction *discordgo.MessageReactionAdd) {
	// check for valid reaction
	if LEFT_ARROW_EMOJI != reaction.Emoji.Name && RIGHT_ARROW_EMOJI != reaction.Emoji.Name && X_EMOJI != reaction.Emoji.Name {
		return
	}

	// check if user who made the embed message is closing it
	if X_EMOJI == reaction.Emoji.Name {
		if reaction.UserID == p.userId {
			delete(pagedEmbededMessages, reaction.MessageID)
			cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)
		}
		return
	}

	// update current page based on direction
	if LEFT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage--
		if p.currentPage == 0 {
			p.currentPage = p.totalNumOfPages
		}
	}
	if RIGHT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage++
		if p.currentPage > p.totalNumOfPages {
			p.currentPage = 1
		}
	}

	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// updated stats
	if p.msgType == "image" {
		// image embeds can't be edited, need to delete and remate it
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)

		// if fields were sent with image embed, handle those
		if len(p.fullEmbed.Fields) > 0 {

			// get start and end fields based on current page and fields per page
			startField := (p.currentPage - 1) * p.fieldsPerPage
			endField := startField + p.fieldsPerPage
			if endField > len(p.fullEmbed.Fields) {
				endField = len(p.fullEmbed.Fields)
			}

			tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		}

		// we need to split and reset the reader since a reader can only be used once
		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		// change the url of the embed to point to the new image
		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)

		// update footer and send message
		tempEmbed.Footer = p.getEmbedFooter()
		sentMessage, _ := SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

		// update map with new message id since
		originalmsgID := p.messageID
		p.messageID = sentMessage[0].ID
		p.addReactionsToMessage()
		pagedEmbededMessages[sentMessage[0].ID] = p
		delete(pagedEmbededMessages, originalmsgID)
	} else {

		// get start and end fields based on current page and fields per page
		startField := (p.currentPage - 1) * p.fieldsPerPage
		endField := startField + p.fieldsPerPage
		if endField > len(p.fullEmbed.Fields) {
			endField = len(p.fullEmbed.Fields)
		}

		tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		tempEmbed.Footer = p.getEmbedFooter()
		EditEmbed(p.channelID, p.messageID, tempEmbed)

		cache.GetSession().MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
	}

}

// setupAndSendFirstMessage
func (p *pagedEmbedMessage) setupAndSendFirstMessage() {
	var sentMessage []*discordgo.Message
	var err error

	// copy the embeded message so changes can be made to it
	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// set footer which will hold information about the page it is on
	tempEmbed.Footer = p.getEmbedFooter()

	if p.msgType == "image" {

		// if fields were sent with image embed, handle those
		if len(p.fullEmbed.Fields) > 0 {

			// get start and end fields based on current page and fields per page
			startField := (p.currentPage - 1) * p.fieldsPerPage
			endField := startField + p.fieldsPerPage
			if endField > len(p.fullEmbed.Fields) {
				endField = len(p.fullEmbed.Fields)
			}

			tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		}

		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)
		sentMessage, _ = SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

	} else {
		// reduce fields to the fields per page
		tempEmbed.Fields = tempEmbed.Fields[:p.fieldsPerPage]

		sentMessage, err = SendEmbed(p.channelID, tempEmbed)
	}

	if err != nil {
		// should probably handle this at some point lul
		return
	}
	p.messageID = sentMessage[0].ID

	p.addReactionsToMessage()
}

// getEmbedFooter is a simlple helper function to return the footer for the embed message
func (p *pagedEmbedMessage) getEmbedFooter() *discordgo.MessageEmbedFooter {
	return &discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("Page: %d / %d", p.currentPage, p.totalNumOfPages),
	}
}

func (p *pagedEmbedMessage) addReactionsToMessage() {
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, LEFT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, RIGHT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, X_EMOJI)
}
<commit_msg>[pagedEmbed] adds goto page functionality for paged embeds<commit_after>package helpers

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"
	"regexp"
	"strconv"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/bwmarrin/discordgo"
)

const (
	LEFT_ARROW_EMOJI  = "⬅"
	RIGHT_ARROW_EMOJI = "➡"
	X_EMOJI           = "🇽"
	NAV_NUMBERS       = "🔢"
)

// map of messageID to pagedEmbedMessage
var pagedEmbededMessages map[string]*pagedEmbedMessage
var validReactions map[string]bool

type pagedEmbedMessage struct {
	files               []*discordgo.File
	fullEmbed           *discordgo.MessageEmbed
	totalNumOfPages     int
	currentPage         int
	fieldsPerPage       int
	color               int
	messageID           string
	channelID           string
	userId              string //user who triggered the message
	msgType             string // "image" - will cause the embed to page through the files instead of fields
	waitingForPageInput bool
}

func init() {
	pagedEmbededMessages = make(map[string]*pagedEmbedMessage)

	validReactions = map[string]bool{
		LEFT_ARROW_EMOJI:  true,
		RIGHT_ARROW_EMOJI: true,
		X_EMOJI:           true,
		NAV_NUMBERS:       true,
	}

}

// will remove all reactions from all paged embed messages.
//  mainly used on bot uninit to clean embeds
func RemoveReactionsFromPagedEmbeds() {
	// TODO: sync group? Without blocking bot shutdown
	for _, pagedEmbed := range pagedEmbededMessages {
		go func(pagedEmbed *pagedEmbedMessage) {
			cache.GetSession().MessageReactionsRemoveAll(pagedEmbed.channelID, pagedEmbed.messageID)
		}(pagedEmbed)
	}
}

// GetPagedMessage will return the paged message if there is one, otherwill will return nil
func GetPagedMessage(messageID string) *pagedEmbedMessage {
	pagedMessaged, _ := pagedEmbededMessages[messageID]
	return pagedMessaged
}

// CreatePagedMessage creates the paged messages
func SendPagedMessage(msg *discordgo.Message, embed *discordgo.MessageEmbed, fieldsPerPage int) error {

	// if there aren't multiple fields to be paged through, or if the amount of fields is less than the requested fields per page
	//  just send a normal embed
	if len(embed.Fields) < 2 || len(embed.Fields) <= fieldsPerPage {
		SendEmbed(msg.ChannelID, embed)
		return nil
	}

	// fields per page can not be less than 1
	if fieldsPerPage < 1 {
		return errors.New("fieldsPerPage can not be less than 1")
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		waitingForPageInput: false,
		fullEmbed:           embed,
		channelID:           msg.ChannelID,
		currentPage:         1,
		fieldsPerPage:       fieldsPerPage,
		totalNumOfPages:     int(math.Ceil(float64(len(embed.Fields)) / float64(fieldsPerPage))),
		userId:              msg.Author.ID,
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// SendPagedImageMessage creates the paged image messages
func SendPagedImageMessage(msg *discordgo.Message, msgSend *discordgo.MessageSend, fieldsPerPage int) error {
	if msgSend.Embed == nil {
		return errors.New("parameter msgSend must contain an embed")
	}

	// make sure the image url is set to the name of the first file incease it wasn't set
	msgSend.Embed.Image.URL = fmt.Sprintf("attachment://%s", msgSend.Files[0].Name)

	// check if there are multiple files, not just send it normally
	if len(msgSend.Files) < 2 {
		SendComplex(msg.ChannelID, msgSend)
		return nil
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		waitingForPageInput: false,
		fullEmbed:           msgSend.Embed,
		channelID:           msg.ChannelID,
		currentPage:         1,
		fieldsPerPage:       fieldsPerPage,
		totalNumOfPages:     len(msgSend.Files),
		files:               msgSend.Files,
		userId:              msg.Author.ID,
		msgType:             "image",
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// UpdateMessagePage will update the page based on the given direction and current page
//  reactions must be the left or right arrow
func (p *pagedEmbedMessage) UpdateMessagePage(reaction *discordgo.MessageReactionAdd) {

	// check for valid reaction
	if !validReactions[reaction.Emoji.Name] || reaction.UserID != p.userId {
		return
	}

	// check if user who made the embed message is closing it
	if X_EMOJI == reaction.Emoji.Name {
		delete(pagedEmbededMessages, reaction.MessageID)
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)
		return
	}

	// check if user wants to navigate to a specific page
	if NAV_NUMBERS == reaction.Emoji.Name {
		if p.waitingForPageInput {
			return
		}

		// if the page entered is valid and not already the current page, then update
		if page, err := p.getUserInputPage(); err == nil && p.currentPage != page {
			p.currentPage = page
		} else {
			return
		}
	}

	// update current page based on direction
	if LEFT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage--
		if p.currentPage == 0 {
			p.currentPage = p.totalNumOfPages
		}
	}
	if RIGHT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage++
		if p.currentPage > p.totalNumOfPages {
			p.currentPage = 1
		}
	}

	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// updated stats
	if p.msgType == "image" {
		// image embeds can't be edited, need to delete and remate it
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)

		// if fields were sent with image embed, handle those
		if len(p.fullEmbed.Fields) > 0 {

			// get start and end fields based on current page and fields per page
			startField := (p.currentPage - 1) * p.fieldsPerPage
			endField := startField + p.fieldsPerPage
			if endField > len(p.fullEmbed.Fields) {
				endField = len(p.fullEmbed.Fields)
			}

			tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		}

		// we need to split and reset the reader since a reader can only be used once
		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		// change the url of the embed to point to the new image
		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)

		// update footer and send message
		tempEmbed.Footer = p.getEmbedFooter()
		sentMessage, _ := SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

		// update map with new message id since
		originalmsgID := p.messageID
		p.messageID = sentMessage[0].ID
		p.addReactionsToMessage()
		pagedEmbededMessages[sentMessage[0].ID] = p
		delete(pagedEmbededMessages, originalmsgID)
	} else {

		// get start and end fields based on current page and fields per page
		startField := (p.currentPage - 1) * p.fieldsPerPage
		endField := startField + p.fieldsPerPage
		if endField > len(p.fullEmbed.Fields) {
			endField = len(p.fullEmbed.Fields)
		}

		tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		tempEmbed.Footer = p.getEmbedFooter()
		EditEmbed(p.channelID, p.messageID, tempEmbed)

		cache.GetSession().MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
	}
}

// setupAndSendFirstMessage
func (p *pagedEmbedMessage) setupAndSendFirstMessage() {
	var sentMessage []*discordgo.Message
	var err error

	// copy the embeded message so changes can be made to it
	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// set footer which will hold information about the page it is on
	tempEmbed.Footer = p.getEmbedFooter()

	if p.msgType == "image" {

		// if fields were sent with image embed, handle those
		if len(p.fullEmbed.Fields) > 0 {

			// get start and end fields based on current page and fields per page
			startField := (p.currentPage - 1) * p.fieldsPerPage
			endField := startField + p.fieldsPerPage
			if endField > len(p.fullEmbed.Fields) {
				endField = len(p.fullEmbed.Fields)
			}

			tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		}

		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)
		sentMessage, _ = SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

	} else {
		// reduce fields to the fields per page
		tempEmbed.Fields = tempEmbed.Fields[:p.fieldsPerPage]

		sentMessage, err = SendEmbed(p.channelID, tempEmbed)
	}

	if err != nil {
		// should probably handle this at some point lul
		return
	}
	p.messageID = sentMessage[0].ID

	p.addReactionsToMessage()
}

// getEmbedFooter is a simlple helper function to return the footer for the embed message
func (p *pagedEmbedMessage) getEmbedFooter() *discordgo.MessageEmbedFooter {
	return &discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("Page: %d / %d", p.currentPage, p.totalNumOfPages),
	}
}

func (p *pagedEmbedMessage) addReactionsToMessage() {
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, LEFT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, RIGHT_ARROW_EMOJI)

	if p.totalNumOfPages > 5 {
		cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, NAV_NUMBERS)
	}

	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, X_EMOJI)
}

// getUserInputPage waits for the user to enter a page
func (p *pagedEmbedMessage) getUserInputPage() (int, error) {
	queryMsg, err := SendMessage(p.channelID, "<@"+p.userId+"> Which page would you like to open? <:blobidea:317047867036663809>")
	if err != nil {
		return 0, err
	}

	defer cache.GetSession().ChannelMessageDelete(queryMsg[0].ChannelID, queryMsg[0].ID)

	timeoutChan := make(chan int)
	go func() {
		time.Sleep(time.Second * 45)
		timeoutChan <- 0
	}()

	p.waitingForPageInput = true
	for {
		select {
		case userMsg := <-waitForUserMessage():

			// check for user who opened embed
			if userMsg.Author.ID != p.userId {
				continue
			}

			// get page number from user text
			re := regexp.MustCompile("[0-9]+")
			if userEnteredNum, err := strconv.Atoi(re.FindString(userMsg.Content)); err == nil {

				// delete user message and remove reaction
				go cache.GetSession().ChannelMessageDelete(userMsg.ChannelID, userMsg.ID)
				go cache.GetSession().MessageReactionRemove(p.channelID, p.messageID, NAV_NUMBERS, p.userId)

				p.waitingForPageInput = false
				if userEnteredNum > 0 && userEnteredNum <= p.totalNumOfPages {

					return userEnteredNum, nil
				} else {
					return 0, errors.New("Page out of embed message range")
				}
			}
		case <-timeoutChan:
			go cache.GetSession().MessageReactionRemove(p.channelID, p.messageID, NAV_NUMBERS, p.userId)
			p.waitingForPageInput = false
			return 0, errors.New("Timed out")
		}
	}
}

func waitForUserMessage() chan *discordgo.MessageCreate {
	out := make(chan *discordgo.MessageCreate)
	cache.GetSession().AddHandlerOnce(func(_ *discordgo.Session, e *discordgo.MessageCreate) {
		out <- e
	})
	return out
}
<|endoftext|>


<commit_before>package agent

import (
	"github.com/hashicorp/serf/testutil"
	"github.com/mitchellh/cli"
	"log"
	"testing"
	"time"
)

func TestCommand_implements(t *testing.T) {
	var _ cli.Command = new(Command)
}

func TestCommandRun(t *testing.T) {
	shutdownCh := make(chan struct{})
	defer close(shutdownCh)

	ui := new(cli.MockUi)
	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         ui,
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-rpc-addr", getRPCAddr(),
	}

	resultCh := make(chan int)
	go func() {
		resultCh <- c.Run(args)
	}()

	testutil.Yield()

	// Verify it runs "forever"
	select {
	case <-resultCh:
		t.Fatalf("ended too soon, err: %s", ui.ErrorWriter.String())
	case <-time.After(50 * time.Millisecond):
	}

	// Send a shutdown request
	shutdownCh <- struct{}{}

	select {
	case code := <-resultCh:
		if code != 0 {
			t.Fatalf("bad code: %d", code)
		}
	case <-time.After(50 * time.Millisecond):
		t.Fatalf("timeout")
	}
}

func TestCommandRun_rpc(t *testing.T) {
	doneCh := make(chan struct{})
	shutdownCh := make(chan struct{})
	defer func() {
		close(shutdownCh)
		<-doneCh
	}()

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	rpcAddr := getRPCAddr()
	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-rpc-addr", rpcAddr,
	}

	go func() {
		code := c.Run(args)
		if code != 0 {
			log.Printf("bad: %d", code)
		}

		close(doneCh)
	}()

	testutil.Yield()

	client, err := NewRPCClient(rpcAddr)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer client.Close()

	members, err := client.Members()
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(members) != 1 {
		t.Fatalf("bad: %#v", members)
	}
}

func TestCommandRun_join(t *testing.T) {
	a1 := testAgent(nil)
	if err := a1.Start(); err != nil {
		t.Fatalf("err: %s", err)
	}
	defer a1.Shutdown()

	doneCh := make(chan struct{})
	shutdownCh := make(chan struct{})
	defer func() {
		close(shutdownCh)
		<-doneCh
	}()

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-join", a1.conf.MemberlistConfig.BindAddr,
		"-replay",
	}

	go func() {
		code := c.Run(args)
		if code != 0 {
			log.Printf("bad: %d", code)
		}

		close(doneCh)
	}()

	testutil.Yield()

	if len(a1.Serf().Members()) != 2 {
		t.Fatalf("bad: %#v", a1.Serf().Members())
	}
}

func TestCommandRun_joinFail(t *testing.T) {
	shutdownCh := make(chan struct{})
	defer close(shutdownCh)

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-join", testutil.GetBindAddr().String(),
	}

	code := c.Run(args)
	if code == 0 {
		t.Fatal("should fail")
	}
}
<commit_msg>Adding test for advertise addr<commit_after>package agent

import (
	"bytes"
	"github.com/hashicorp/serf/testutil"
	"github.com/mitchellh/cli"
	"log"
	"testing"
	"time"
)

func TestCommand_implements(t *testing.T) {
	var _ cli.Command = new(Command)
}

func TestCommandRun(t *testing.T) {
	shutdownCh := make(chan struct{})
	defer close(shutdownCh)

	ui := new(cli.MockUi)
	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         ui,
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-rpc-addr", getRPCAddr(),
	}

	resultCh := make(chan int)
	go func() {
		resultCh <- c.Run(args)
	}()

	testutil.Yield()

	// Verify it runs "forever"
	select {
	case <-resultCh:
		t.Fatalf("ended too soon, err: %s", ui.ErrorWriter.String())
	case <-time.After(50 * time.Millisecond):
	}

	// Send a shutdown request
	shutdownCh <- struct{}{}

	select {
	case code := <-resultCh:
		if code != 0 {
			t.Fatalf("bad code: %d", code)
		}
	case <-time.After(50 * time.Millisecond):
		t.Fatalf("timeout")
	}
}

func TestCommandRun_rpc(t *testing.T) {
	doneCh := make(chan struct{})
	shutdownCh := make(chan struct{})
	defer func() {
		close(shutdownCh)
		<-doneCh
	}()

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	rpcAddr := getRPCAddr()
	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-rpc-addr", rpcAddr,
	}

	go func() {
		code := c.Run(args)
		if code != 0 {
			log.Printf("bad: %d", code)
		}

		close(doneCh)
	}()

	testutil.Yield()

	client, err := NewRPCClient(rpcAddr)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer client.Close()

	members, err := client.Members()
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(members) != 1 {
		t.Fatalf("bad: %#v", members)
	}
}

func TestCommandRun_join(t *testing.T) {
	a1 := testAgent(nil)
	if err := a1.Start(); err != nil {
		t.Fatalf("err: %s", err)
	}
	defer a1.Shutdown()

	doneCh := make(chan struct{})
	shutdownCh := make(chan struct{})
	defer func() {
		close(shutdownCh)
		<-doneCh
	}()

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-join", a1.conf.MemberlistConfig.BindAddr,
		"-replay",
	}

	go func() {
		code := c.Run(args)
		if code != 0 {
			log.Printf("bad: %d", code)
		}

		close(doneCh)
	}()

	testutil.Yield()

	if len(a1.Serf().Members()) != 2 {
		t.Fatalf("bad: %#v", a1.Serf().Members())
	}
}

func TestCommandRun_joinFail(t *testing.T) {
	shutdownCh := make(chan struct{})
	defer close(shutdownCh)

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-join", testutil.GetBindAddr().String(),
	}

	code := c.Run(args)
	if code == 0 {
		t.Fatal("should fail")
	}
}

func TestCommandRun_advertiseAddr(t *testing.T) {
	doneCh := make(chan struct{})
	shutdownCh := make(chan struct{})
	defer func() {
		close(shutdownCh)
		<-doneCh
	}()

	c := &Command{
		ShutdownCh: shutdownCh,
		Ui:         new(cli.MockUi),
	}

	rpcAddr := getRPCAddr()
	args := []string{
		"-bind", testutil.GetBindAddr().String(),
		"-rpc-addr", rpcAddr,
		"-advertise", "127.0.0.10:12345",
	}

	go func() {
		code := c.Run(args)
		if code != 0 {
			log.Printf("bad: %d", code)
		}

		close(doneCh)
	}()

	testutil.Yield()

	client, err := NewRPCClient(rpcAddr)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer client.Close()

	members, err := client.Members()
	if err != nil {
		t.Fatalf("err: %s", err)
	}

	if len(members) != 1 {
		t.Fatalf("bad: %#v", members)
	}

	// Check the addr and port is as advertised!
	m := members[0]
	if bytes.Compare(m.Addr, []byte{127, 0, 0, 10}) != 0 {
		t.Fatalf("bad: %#v", m)
	}
	if m.Port != 12345 {
		t.Fatalf("bad: %#v", m)
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"io"
	"sync"
)

// GatedWriter is an io.Writer implementation that buffers all of its
// data into an internal buffer until it is told to let data through.
type GatedWriter struct {
	Writer io.Writer

	buf   [][]byte
	flush bool
	lock  sync.RWMutex
}

// Flush tells the GatedWriter to flush any buffered data and to stop
// buffering.
func (w *GatedWriter) Flush() {
	w.lock.Lock()
	w.flush = true
	w.lock.Unlock()

	for _, p := range w.buf {
		w.Write(p)
	}
	w.buf = nil
}

func (w *GatedWriter) Write(p []byte) (n int, err error) {
        // Once we flush we no longer synchronize writers since there's
        // no use of the internal buffer. This is the happy path.
        w.lock.RLock()
        if w.flush {
                w.lock.RUnlock()
                return w.Writer.Write(p)
        }
        w.lock.RUnlock()

        // Now take the write lock.
        w.lock.Lock()
        defer w.lock.Unlock()

        // Things could have changed between the locking operations, so we
        // have to check one more time.
        if w.flush {
                return w.Writer.Write(p)
        }

        // Buffer up the written data.
        p2 := make([]byte, len(p))
        copy(p2, p)
        w.buf = append(w.buf, p2)
        return len(p), nil
}
<commit_msg>Runs go-fmt (incidental to this PR).<commit_after>package agent

import (
	"io"
	"sync"
)

// GatedWriter is an io.Writer implementation that buffers all of its
// data into an internal buffer until it is told to let data through.
type GatedWriter struct {
	Writer io.Writer

	buf   [][]byte
	flush bool
	lock  sync.RWMutex
}

// Flush tells the GatedWriter to flush any buffered data and to stop
// buffering.
func (w *GatedWriter) Flush() {
	w.lock.Lock()
	w.flush = true
	w.lock.Unlock()

	for _, p := range w.buf {
		w.Write(p)
	}
	w.buf = nil
}

func (w *GatedWriter) Write(p []byte) (n int, err error) {
	// Once we flush we no longer synchronize writers since there's
	// no use of the internal buffer. This is the happy path.
	w.lock.RLock()
	if w.flush {
		w.lock.RUnlock()
		return w.Writer.Write(p)
	}
	w.lock.RUnlock()

	// Now take the write lock.
	w.lock.Lock()
	defer w.lock.Unlock()

	// Things could have changed between the locking operations, so we
	// have to check one more time.
	if w.flush {
		return w.Writer.Write(p)
	}

	// Buffer up the written data.
	p2 := make([]byte, len(p))
	copy(p2, p)
	w.buf = append(w.buf, p2)
	return len(p), nil
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// go test *.go -bench=".*" -benchmem

package gqueue_test

import (
	"github.com/gogf/gf/g/container/gqueue"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

func TestQueue_Len(t *testing.T) {

	max := 100
	for n := 10; n < max; n++ {
		q1 := gqueue.New(max)
		for i := 0; i < max; i++ {
			q1.Push(i)
		}
		gtest.Assert(q1.Len(), max)
		gtest.Assert(q1.Size(), max)
	}
}

func TestQueue_Basic(t *testing.T) {
	q := gqueue.New()
	for i := 0; i < 100; i++ {
		q.Push(i)
	}
	gtest.Assert(q.Pop(), 0)
	gtest.Assert(q.Pop(), 1)
}

func TestQueue_Pop(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	q1.Push(3)
	q1.Push(4)
	i1 := q1.Pop()
	gtest.Assert(i1, 1)

}

func TestQueue_Close(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	gtest.Assert(q1.Len(), 2)
	q1.Close()
	q1.Pop()
	gtest.Assert(q1.Len(), 2)
}
<commit_msg>Update gqueue_unit_test.go<commit_after>// Copyright 2017 gf Author(https://github.com/gogf/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.

// go test *.go -bench=".*" -benchmem

package gqueue_test

import (
	"github.com/gogf/gf/g/container/gqueue"
	"github.com/gogf/gf/g/test/gtest"
	"testing"
)

func TestQueue_Len(t *testing.T) {

	max := 100
	for n := 10; n < max; n++ {
		q1 := gqueue.New(max)
		for i := 0; i < max; i++ {
			q1.Push(i)
		}
		gtest.Assert(q1.Len(), max)
		gtest.Assert(q1.Size(), max)
	}
}

func TestQueue_Basic(t *testing.T) {
	q := gqueue.New()
	for i := 0; i < 100; i++ {
		q.Push(i)
	}
	gtest.Assert(q.Pop(), 0)
	gtest.Assert(q.Pop(), 1)
}

func TestQueue_Pop(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	q1.Push(3)
	q1.Push(4)
	i1 := q1.Pop()
	gtest.Assert(i1, 1)

}

func TestQueue_Close(t *testing.T) {
	q1 := gqueue.New()
	q1.Push(1)
	q1.Push(2)
	gtest.Assert(q1.Len(), 2)
	q1.Close()
}
<|endoftext|>


<commit_before>package geotrigger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
)

const (
	geotrigger_base_url = "https://geotrigger.arcgis.com"
	ago_base_url        = "https://www.arcgis.com"
	ago_token_route     = "/sharing/oauth2/token"
	ago_register_route  = "/sharing/oauth2/registerDevice"
	version             = "1.0.0"
)

var defEnv = &environment{
	geotrigger_base_url,
	ago_base_url,
}

// The Session interface obfuscates whether we are a device or an application,
// both of which implement the interface slightly differently.
type session interface {
	request(string, map[string]interface{}, interface{}) error
	info() map[string]string
	// A session is also a TokenManager
	tokenManager
	// used internally when token expires
	refresh(string) error
	// used internally for changing URLs at runtime for testing
	setEnv(*environment)
}

type environment struct {
	geotriggerURL string
	agoURL        string
}

type errorResponse struct {
	Error errorJSON `json:"error"`
}

type errorJSON struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// func type for passing in to `post`. called when we get a 498 invalid token
type refreshHandler func() (string, error)

// funcs below manage http specifically for geotrigger service and AGO credentials
func doRefresh(session session, token string) (string, error) {
	error := session.refresh(token)
	var refreshResult *tokenRequest
	var accessToken string
	if error == nil {
		accessToken = session.getAccessToken()
		refreshResult = newTokenRequest(refreshComplete, false)
	} else {
		refreshResult = newTokenRequest(refreshFailed, false)
	}

	go session.tokenRequest(refreshResult)

	return accessToken, error
}

func geotriggerPost(env *environment, session session, route string, params map[string]interface{},
	responseJSON interface{}) error {
	body, err := json.Marshal(params)
	if err != nil {
		return fmt.Errorf("Error while marshaling params into JSON for route: %s. %s", route, err)
	}

	// This func gets a blocking call if we get a 498 from the geotrigger server
	refreshFunc := func() (string, error) {
		tr := newTokenRequest(refreshNeeded, true)
		go session.tokenRequest(tr)

		tokenResp := <-tr.tokenResponses

		if tokenResp.isAccessToken {
			// refresh request denied, another routine has already refreshed!
			// go ahead and use this access token
			return tokenResp.token, nil
		}

		// refresh request approved, get a fresh token
		return doRefresh(session, tokenResp.token)
	}

	tr := newTokenRequest(accessNeeded, true)
	go session.tokenRequest(tr)

	tokenResp := <-tr.tokenResponses

	var token string
	if tokenResp.isAccessToken {
		// we have access, go ahead and use it
		token = tokenResp.token
	} else {
		// access request denied, the token has expired. go get a fresh one
		token, err = doRefresh(session, tokenResp.token)
	}

	if err != nil {
		return fmt.Errorf("Error while trying to refresh token before hitting route: %s. %s",
			route, err)
	}

	req, err := http.NewRequest("POST", routeConcat(env.geotriggerURL, route), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("Error creating GeotriggerPost for route %s. %s", route, err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-GT-Client-Name", "geotrigger-go")
	req.Header.Set("X-GT-Client-Version", version)

	return post(req, body, responseJSON, refreshFunc)
}

func agoPost(env *environment, route string, body []byte, responseJSON interface{}) error {
	req, err := http.NewRequest("POST", routeConcat(env.agoURL, route), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("Error creating AgoPost for route %s. %s", route, err)
	}

	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	return post(req, body, responseJSON, func() (string, error) {
		return "", errors.New("Expired token response from AGO. This is basically a 500.")
	})
}

func post(req *http.Request, body []byte, responseJSON interface{}, refreshFunc refreshHandler) error {
	path := req.URL.Path

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("Error while posting to: %s. Error: %s", path, err)
	}

	if resp.StatusCode != 200 {
		return fmt.Errorf("Received status code %d from %s.", resp.StatusCode, path)
	}

	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("Could not read response body from %s. %s", path, err)
	}

	if errResponse := errorCheck(contents); errResponse != nil {
		if errResponse.Error.Code == 498 {
			if token, err := refreshFunc(); err == nil {
				// time to refresh!
				// the body of the req cannot be reused, because it has already been read
				// and the standard lib can't rewind the pointer on the same content.
				// so, we have passed the underlying []byte down here so we can
				// make a new reader from it. This is a bit unsafe (we are skipping
				// the NewRequest constructor), but since the data is the same, all should be well.
				var bodyReader io.Reader
				bodyReader = bytes.NewReader(body)
				rc, ok := bodyReader.(io.ReadCloser)
				if !ok && bodyReader != nil {
					rc = ioutil.NopCloser(bodyReader)
				}
				req.Body = rc
				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
				return post(req, body, responseJSON, refreshFunc)
			} else {
				return err
			}
		} else {
			return fmt.Errorf("Error from %s, code: %d. Message: %s",
				path, errResponse.Error.Code, errResponse.Error.Message)
		}
	}

	return parseJSONResponse(contents, responseJSON)
}

func errorCheck(resp []byte) *errorResponse {
	var errorContainer errorResponse
	if err := json.Unmarshal(resp, &errorContainer); err != nil {
		// Don't return an error here, as it is possible for the response
		// to not be parsed into an errorResponse, causing an error to be thrown, but still
		// be valid, ie: the root element of the response is an array.
		// We are just looking to see if we can spot a known server error anyway.
		return nil
	}

	if errorContainer.Error.Code > 0 && len(errorContainer.Error.Message) > 0 {
		return &errorContainer
	}

	return nil
}

func parseJSONResponse(resp []byte, responseJSON interface{}) error {
	t := reflect.TypeOf(responseJSON)
	if t == nil || t.Kind() != reflect.Ptr {
		return fmt.Errorf("Provided responseJSON interface should be a pointer (to struct or map).")
	}

	if err := json.Unmarshal(resp, responseJSON); err != nil {
		return fmt.Errorf("Error parsing response: %s  Error: %s", string(resp), err)
	}

	return nil
}

func routeConcat(baseURL, route string) string {
	var buffer bytes.Buffer
	buffer.WriteString(baseURL)

	if route[:1] != "/" {
		buffer.WriteString("/")
	}

	buffer.WriteString(route)

	return buffer.String()
}

func testEnv(gtURL, agoURL string) *environment {
	return &environment{gtURL, agoURL}
}
<commit_msg>server not always returning code with error... Fixes #6.<commit_after>package geotrigger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"reflect"
)

const (
	geotrigger_base_url = "https://geotrigger.arcgis.com"
	ago_base_url        = "https://www.arcgis.com"
	ago_token_route     = "/sharing/oauth2/token"
	ago_register_route  = "/sharing/oauth2/registerDevice"
	version             = "1.0.0"
)

var defEnv = &environment{
	geotrigger_base_url,
	ago_base_url,
}

// The Session interface obfuscates whether we are a device or an application,
// both of which implement the interface slightly differently.
type session interface {
	request(string, map[string]interface{}, interface{}) error
	info() map[string]string
	// A session is also a TokenManager
	tokenManager
	// used internally when token expires
	refresh(string) error
	// used internally for changing URLs at runtime for testing
	setEnv(*environment)
}

type environment struct {
	geotriggerURL string
	agoURL        string
}

type errorResponse struct {
	Error errorJSON `json:"error"`
}

type errorJSON struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// func type for passing in to `post`. called when we get a 498 invalid token
type refreshHandler func() (string, error)

// funcs below manage http specifically for geotrigger service and AGO credentials
func doRefresh(session session, token string) (string, error) {
	error := session.refresh(token)
	var refreshResult *tokenRequest
	var accessToken string
	if error == nil {
		accessToken = session.getAccessToken()
		refreshResult = newTokenRequest(refreshComplete, false)
	} else {
		refreshResult = newTokenRequest(refreshFailed, false)
	}

	go session.tokenRequest(refreshResult)

	return accessToken, error
}

func geotriggerPost(env *environment, session session, route string, params map[string]interface{},
	responseJSON interface{}) error {
	body, err := json.Marshal(params)
	if err != nil {
		return fmt.Errorf("Error while marshaling params into JSON for route: %s. %s", route, err)
	}

	// This func gets a blocking call if we get a 498 from the geotrigger server
	refreshFunc := func() (string, error) {
		tr := newTokenRequest(refreshNeeded, true)
		go session.tokenRequest(tr)

		tokenResp := <-tr.tokenResponses

		if tokenResp.isAccessToken {
			// refresh request denied, another routine has already refreshed!
			// go ahead and use this access token
			return tokenResp.token, nil
		}

		// refresh request approved, get a fresh token
		return doRefresh(session, tokenResp.token)
	}

	tr := newTokenRequest(accessNeeded, true)
	go session.tokenRequest(tr)

	tokenResp := <-tr.tokenResponses

	var token string
	if tokenResp.isAccessToken {
		// we have access, go ahead and use it
		token = tokenResp.token
	} else {
		// access request denied, the token has expired. go get a fresh one
		token, err = doRefresh(session, tokenResp.token)
	}

	if err != nil {
		return fmt.Errorf("Error while trying to refresh token before hitting route: %s. %s",
			route, err)
	}

	req, err := http.NewRequest("POST", routeConcat(env.geotriggerURL, route), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("Error creating GeotriggerPost for route %s. %s", route, err)
	}

	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-GT-Client-Name", "geotrigger-go")
	req.Header.Set("X-GT-Client-Version", version)

	return post(req, body, responseJSON, refreshFunc)
}

func agoPost(env *environment, route string, body []byte, responseJSON interface{}) error {
	req, err := http.NewRequest("POST", routeConcat(env.agoURL, route), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("Error creating AgoPost for route %s. %s", route, err)
	}

	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	return post(req, body, responseJSON, func() (string, error) {
		return "", errors.New("Expired token response from AGO. This is basically a 500.")
	})
}

func post(req *http.Request, body []byte, responseJSON interface{}, refreshFunc refreshHandler) error {
	path := req.URL.Path

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("Error while posting to: %s. Error: %s", path, err)
	}

	if resp.StatusCode != 200 {
		return fmt.Errorf("Received status code %d from %s.", resp.StatusCode, path)
	}

	defer resp.Body.Close()
	contents, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("Could not read response body from %s. %s", path, err)
	}

	if errResponse := errorCheck(contents); errResponse != nil {
		if errResponse.Error.Code == 498 {
			if token, err := refreshFunc(); err == nil {
				// time to refresh!
				// the body of the req cannot be reused, because it has already been read
				// and the standard lib can't rewind the pointer on the same content.
				// so, we have passed the underlying []byte down here so we can
				// make a new reader from it. This is a bit unsafe (we are skipping
				// the NewRequest constructor), but since the data is the same, all should be well.
				var bodyReader io.Reader
				bodyReader = bytes.NewReader(body)
				rc, ok := bodyReader.(io.ReadCloser)
				if !ok && bodyReader != nil {
					rc = ioutil.NopCloser(bodyReader)
				}
				req.Body = rc
				req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
				return post(req, body, responseJSON, refreshFunc)
			} else {
				return err
			}
		} else {
			return fmt.Errorf("Error from %s, code: %d. Message: %s",
				path, errResponse.Error.Code, errResponse.Error.Message)
		}
	}

	return parseJSONResponse(contents, responseJSON)
}

func errorCheck(resp []byte) *errorResponse {
	var errorContainer errorResponse
	if err := json.Unmarshal(resp, &errorContainer); err != nil {
		// Don't return an error here, as it is possible for the response
		// to not be parsed into an errorResponse, causing an error to be thrown, but still
		// be valid, ie: the root element of the response is an array.
		// We are just looking to see if we can spot a known server error anyway.
		return nil
	}

	// We don't inspect the value of `Error.Code` here because the server may not return a code value
	if len(errorContainer.Error.Message) > 0 {
		return &errorContainer
	}

	return nil
}

func parseJSONResponse(resp []byte, responseJSON interface{}) error {
	t := reflect.TypeOf(responseJSON)
	if t == nil || t.Kind() != reflect.Ptr {
		return fmt.Errorf("Provided responseJSON interface should be a pointer (to struct or map).")
	}

	if err := json.Unmarshal(resp, responseJSON); err != nil {
		return fmt.Errorf("Error parsing response: %s  Error: %s", string(resp), err)
	}

	return nil
}

func routeConcat(baseURL, route string) string {
	var buffer bytes.Buffer
	buffer.WriteString(baseURL)

	if route[:1] != "/" {
		buffer.WriteString("/")
	}

	buffer.WriteString(route)

	return buffer.String()
}

func testEnv(gtURL, agoURL string) *environment {
	return &environment{gtURL, agoURL}
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"time"

	"github.com/timeglass/snow/monitor"
)

type Timer struct {
	mbu     time.Duration
	timeout time.Duration
	time    time.Duration
	ticking bool
	read    chan chan time.Duration
	inc     chan chan time.Duration
	reset   chan struct{}

	Wakeup <-chan monitor.DirEvent
	*sync.Mutex
}

func NewTimer(mbu time.Duration, to time.Duration) *Timer {
	t := &Timer{
		mbu:     mbu,
		timeout: to,

		read:  make(chan chan time.Duration),
		inc:   make(chan chan time.Duration),
		reset: make(chan struct{}),

		Wakeup: make(chan monitor.DirEvent),
		Mutex:  &sync.Mutex{},
	}

	//handle read& writes
	go func() {
		for {
			select {
			case r := <-t.read:
				r <- t.time
			case i := <-t.inc:
				t.time += <-i
			case <-t.reset:
				t.time = 0
			}
		}
	}()

	//handle timeout and wakeup
	go func() {
		for {
			select {
			case <-time.After(t.timeout):
				t.Stop()
			case <-t.Wakeup:
				t.Start()
			}
		}
	}()

	return t
}

func (t *Timer) Time() time.Duration {
	r := make(chan time.Duration)
	t.read <- r
	return <-r
}

func (t *Timer) Reset() {
	t.reset <- struct{}{}
}

func (t *Timer) Stop() {
	t.Lock()
	defer t.Unlock()
	t.ticking = false
}

func (t *Timer) Start() {
	t.Lock()
	defer t.Unlock()
	if t.ticking {
		return
	}

	t.ticking = true
	go func() {
		for {

			//increment with mbu
			i := make(chan time.Duration)
			t.inc <- i
			i <- t.mbu

			//wait for next mbu to arrive
			<-time.After(t.mbu)

			//previous tick was the last mbu
			t.Lock()
			defer t.Unlock()
			if !t.ticking {
				return
			}

		}
	}()
}
<commit_msg>fixed regression for 0.5.2 [spent 1m10s]<commit_after>package main

import (
	"sync"
	"time"

	"github.com/timeglass/snow/monitor"
)

type Timer struct {
	mbu     time.Duration
	timeout time.Duration
	time    time.Duration
	ticking bool
	read    chan chan time.Duration
	inc     chan chan time.Duration
	reset   chan struct{}

	Wakeup <-chan monitor.DirEvent
	*sync.Mutex
}

func NewTimer(mbu time.Duration, to time.Duration) *Timer {
	t := &Timer{
		mbu:     mbu,
		timeout: to,

		read:  make(chan chan time.Duration),
		inc:   make(chan chan time.Duration),
		reset: make(chan struct{}),

		Wakeup: make(chan monitor.DirEvent),
		Mutex:  &sync.Mutex{},
	}

	//handle read& writes
	go func() {
		for {
			select {
			case r := <-t.read:
				r <- t.time
			case i := <-t.inc:
				t.time += <-i
			case <-t.reset:
				t.time = 0
			}
		}
	}()

	//handle timeout and wakeup
	go func() {
		for {
			select {
			case <-time.After(t.timeout):
				t.Stop()
			case <-t.Wakeup:
				t.Start()
			}
		}
	}()

	return t
}

func (t *Timer) Time() time.Duration {
	r := make(chan time.Duration)
	t.read <- r
	return <-r
}

func (t *Timer) Reset() {
	t.reset <- struct{}{}
}

func (t *Timer) Stop() {
	t.Lock()
	defer t.Unlock()
	t.ticking = false
}

func (t *Timer) Start() {
	t.Lock()
	defer t.Unlock()
	if t.ticking {
		return
	}

	t.ticking = true
	go func() {
		for {

			//increment with mbu
			i := make(chan time.Duration)
			t.inc <- i
			i <- t.mbu

			//wait for next mbu to arrive
			<-time.After(t.mbu)

			//previous tick was the last mbu
			if !t.ticking {
				return
			}
		}
	}()
}
<|endoftext|>