<commit_before>package instancemetadata

import (
	"errors"
	"fmt"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestNew(t *testing.T) {
	var mdf MetadataFetcher
	mdf = New(true)
	if mdf == nil {
		t.Fail()
	}
}

type FakeMetadataFetcher struct {
	FAvailable bool
	Meta       map[string]string
}

func (m FakeMetadataFetcher) Available() bool {
	return m.FAvailable
}

func (m FakeMetadataFetcher) GetMetadata(key string) (string, error) {
	v, ok := m.Meta[key]
	if ok {
		return v, nil
	}
	return v, errors.New(fmt.Sprintf("Key %s unknown", key))
}

func getFakeMetadataFetcher(a bool) MetadataFetcher {
	fakeM := FakeMetadataFetcher{
		FAvailable: a,
	}
	fakeM.Meta = make(map[string]string)
	fakeM.Meta["placement/availability-zone"] = "us-west-1a"
	fakeM.Meta["instance-id"] = "i-1234"
	fakeM.Meta["mac"] = "06:1d:ea:6f:8c:6e"
	fakeM.Meta["local-ipv4"] = "127.0.0.1"
	fakeM.Meta["network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id"] = "subnet-28b0e940"
	return fakeM
}

func TestGetSubnetIdMacFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "mac")
	_, err := getSubnetId(mdf)
	assert.NotNil(t, err)
}

func TestGetSubnetIdMacFail2(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id")
	_, err := getSubnetId(mdf)
	assert.NotNil(t, err)
}

func TestGetSubnetIdMacOk(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	val, err := getSubnetId(mdf)
	if assert.Nil(t, err) {
		assert.Equal(t, val, "subnet-28b0e940")
	}
}

func TestFetchMetadataFail(t *testing.T) {
	_, err := FetchMetadata(getFakeMetadataFetcher(false))
	if assert.NotNil(t, err) {
		assert.Equal(t, err.Error(), "No metadata service")
	}
}

func TestFetchMetadataAzFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "placement/availability-zone")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadataInstanceFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "instance-id")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadataMacFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "mac")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadata(t *testing.T) {
	m, err := FetchMetadata(getFakeMetadataFetcher(true))
	assert.Nil(t, err)
	assert.Equal(t, m.Instance, "i-1234")
	assert.Equal(t, m.Subnet, "subnet-28b0e940")
	assert.Equal(t, m.AvailabilityZone, "us-west-1a")
	assert.Equal(t, m.Region, "us-west-1")
}
<commit_msg>Clean up<commit_after>package instancemetadata

import (
	"errors"
	"fmt"
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestNew(t *testing.T) {
	var mdf MetadataFetcher
	mdf = New(true)
	assert.NotNil(t, mdf)
}

type FakeMetadataFetcher struct {
	FAvailable bool
	Meta       map[string]string
}

func (m FakeMetadataFetcher) Available() bool {
	return m.FAvailable
}

func (m FakeMetadataFetcher) GetMetadata(key string) (string, error) {
	v, ok := m.Meta[key]
	if ok {
		return v, nil
	}
	return v, errors.New(fmt.Sprintf("Key %s unknown", key))
}

func getFakeMetadataFetcher(a bool) MetadataFetcher {
	fakeM := FakeMetadataFetcher{
		FAvailable: a,
	}
	fakeM.Meta = make(map[string]string)
	fakeM.Meta["placement/availability-zone"] = "us-west-1a"
	fakeM.Meta["instance-id"] = "i-1234"
	fakeM.Meta["mac"] = "06:1d:ea:6f:8c:6e"
	fakeM.Meta["local-ipv4"] = "127.0.0.1"
	fakeM.Meta["network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id"] = "subnet-28b0e940"
	return fakeM
}

func TestGetSubnetIdMacFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "mac")
	_, err := getSubnetId(mdf)
	assert.NotNil(t, err)
}

func TestGetSubnetIdMacFail2(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "network/interfaces/macs/06:1d:ea:6f:8c:6e/subnet-id")
	_, err := getSubnetId(mdf)
	assert.NotNil(t, err)
}

func TestGetSubnetIdMacOk(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	val, err := getSubnetId(mdf)
	if assert.Nil(t, err) {
		assert.Equal(t, val, "subnet-28b0e940")
	}
}

func TestFetchMetadataFail(t *testing.T) {
	_, err := FetchMetadata(getFakeMetadataFetcher(false))
	if assert.NotNil(t, err) {
		assert.Equal(t, err.Error(), "No metadata service")
	}
}

func TestFetchMetadataAzFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "placement/availability-zone")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadataInstanceFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "instance-id")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadataMacFail(t *testing.T) {
	mdf := getFakeMetadataFetcher(true)
	delete(mdf.(FakeMetadataFetcher).Meta, "mac")
	_, err := FetchMetadata(mdf)
	assert.NotNil(t, err)
}

func TestFetchMetadata(t *testing.T) {
	m, err := FetchMetadata(getFakeMetadataFetcher(true))
	assert.Nil(t, err)
	assert.Equal(t, m.Instance, "i-1234")
	assert.Equal(t, m.Subnet, "subnet-28b0e940")
	assert.Equal(t, m.AvailabilityZone, "us-west-1a")
	assert.Equal(t, m.Region, "us-west-1")
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Tim Heckman
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package formatter

import (
	"encoding/json"
	"fmt"
	"math/big"
	"strconv"
)

// PrintTokens prints the generated tokens in the same format as the Cassandra token generator
func PrintTokens(t [][]*big.Int, w int) {
	for i, tokenList := range t {
		// print the header
		fmt.Println(fmt.Sprintf("DC #%d:", i+1))

		// get the width of the largest number to properly space the column
		nnWidth := len(strconv.Itoa(len(tokenList)))

		// print each node in the datacenter
		for ni, nt := range tokenList {
			fmt.Println(fmt.Sprintf("  Node #%*d: % *d", nnWidth, ni+1, w+1, nt))
		}
	}
}

// PrintJSON prints the results of the token generation in a JSON format
func PrintJSON(t [][]*big.Int, prettyPrint bool) {
	var dcList []*string
	var jsonBytes []byte
	var err error

	data := make(map[string]interface{})

	data["keys"] = make([]*string, 0)

	for x, v := range t {
		// set the key name for this datacenter
		dcStr := fmt.Sprintf("dc%d", x+1)

		// append the key name to the array for specifying datacenter order
		dcList = append(dcList, &dcStr)

		// create the entry in the map for this datacenter
		data[dcStr] = v
	}

	data["keys"] = dcList

	if prettyPrint {
		jsonBytes, err = json.MarshalIndent(data, "", "  ")
	} else {
		jsonBytes, err = json.Marshal(data)
	}

	if err != nil {
		fmt.Println("error printing json:", err)
		return
	}

	fmt.Println(string(jsonBytes))
}
<commit_msg>improvements to the formatter<commit_after>// Copyright (c) 2014 Tim Heckman
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package formatter

import (
	"encoding/json"
	"fmt"
	"math/big"
	"strconv"
)

// PrintTokens prints the generated tokens in the same format as the Cassandra token generator
func PrintTokens(t [][]*big.Int, w int) {
	for i, tokenList := range t {
		// print the header
		fmt.Printf("DC #%d:\n", i+1)

		// get the width of the largest number to properly space the column
		nnWidth := len(strconv.Itoa(len(tokenList)))

		// print each node in the datacenter
		for ni, nt := range tokenList {
			fmt.Printf("  Node #%*d: % *d\n", nnWidth, ni+1, w+1, nt)
		}
	}
}

func jsonMarshal(v map[string]interface{}, pp bool) ([]byte, err) {
	if pp {
		return json.MarshalIndent(v, "", "  ")
	} else {
		return json.Marshal(v)
	}
}

// PrintJSON prints the results of the token generation in a JSON format
func PrintJSON(t [][]*big.Int, prettyPrint bool) {
	data := make(map[string]interface{})

	data["keys"] = make([]*string, 0)

	var dcList []*string

	for x, v := range t {
		// set the key name for this datacenter
		dcStr := fmt.Sprintf("dc%d", x+1)

		// append the key name to the array for specifying datacenter order
		dcList = append(dcList, &dcStr)

		// create the entry in the map for this datacenter
		data[dcStr] = v
	}

	data["keys"] = dcList

	jsonBytes, err := jsonMarshal(data, prettyPrint)

	if err != nil {
		m := make(map[string]string)
		m["error"] = err.Error()

		j, err2 := jsonMarshal(m, prettyPrint)

		// if this is not nilwe've hit a very unexpected and very
		// unavoidable error this really should never be hit unless
		// something really bizarre happens...
		if err2 != nil {
			panic(fmt.Sprintf("unavoidable error; '%v' when JSON Marshaling error for data map. Original: '%v'\n", err2.Error(), err.Error()))
		}

		fmt.Println(string(j))
		return
	}

	fmt.Println(string(jsonBytes))
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package save

import (
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
	"sort"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestFSSuccessorFinder(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type fsNodeSlice []*fsNode

func (p fsNodeSlice) Len() int {
	return len(p)
}

func (p fsNodeSlice) Less(i, j int) bool {
	return p[i].RelPath < p[j].RelPath
}

func (p fsNodeSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func sortNodes(graphNodes []graph.Node) (nodes fsNodeSlice) {
	for _, n := range graphNodes {
		nodes = append(nodes, n.(*fsNode))
	}

	sort.Sort(nodes)
	return
}

// A valid os.FileInfo for a directory.
var gDirInfo os.FileInfo

func init() {
	var err error
	gDirInfo, err = os.Stat("/")
	if err != nil {
		log.Fatalf("Stat: %v", err)
	}
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FSSuccessorFinderTest struct {
	ctx context.Context

	// A temporary directory that is cleaned up at the end of the test. This is
	// the base path with which the successor finder is configured.
	dir string

	// The exclusions with which to configure the successor finder.
	exclusions []*regexp.Regexp

	sf graph.SuccessorFinder
}

var _ SetUpInterface = &FSSuccessorFinderTest{}
var _ TearDownInterface = &FSSuccessorFinderTest{}

func init() { RegisterTestSuite(&FSSuccessorFinderTest{}) }

func (t *FSSuccessorFinderTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Create the base directory.
	var err error
	t.dir, err = ioutil.TempDir("", "file_system_visistor_test")
	AssertEq(nil, err)

	// And the successor finder.
	t.resetVisistor()
}

func (t *FSSuccessorFinderTest) TearDown() {
	var err error

	// Clean up the junk we left in the file system.
	err = os.RemoveAll(t.dir)
	AssertEq(nil, err)
}

func (t *FSSuccessorFinderTest) resetVisistor() {
	t.sf = newSuccessorFinder(t.dir, t.exclusions)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FSSuccessorFinderTest) NonExistentPath() {
	node := &fsNode{
		RelPath: "foo/bar",
		Info:    gDirInfo,
	}

	_, err := t.sf.FindDirectSuccessors(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node.RelPath)))
	ExpectThat(err, Error(HasSubstr("no such file")))
}

func (t *FSSuccessorFinderTest) VisitRootNode() {
	var err error

	// Create two children.
	err = ioutil.WriteFile(path.Join(t.dir, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(t.dir, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	node := &fsNode{
		RelPath: "",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))
	ExpectEq("bar", pfis[0].RelPath)
	ExpectEq("foo", pfis[1].RelPath)
}

func (t *FSSuccessorFinderTest) VisitNonRootNode() {
	var err error

	// Make a few levels of sub-directories.
	d := path.Join(t.dir, "sub/dirs")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the directory.
	node := &fsNode{
		RelPath: "sub/dirs",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))
	ExpectEq("sub/dirs/bar", pfis[0].RelPath)
	ExpectEq("sub/dirs/foo", pfis[1].RelPath)
}

func (t *FSSuccessorFinderTest) VisitFileNode() {
	var err error

	// Create a file.
	f := path.Join(t.dir, "foo")
	err = ioutil.WriteFile(f, []byte{}, 0500)
	AssertEq(nil, err)

	fi, err := os.Stat(f)
	AssertEq(nil, err)

	// Visit it.
	node := &fsNode{
		RelPath: "foo",
		Info:    fi,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	ExpectThat(successors, ElementsAre())
}

func (t *FSSuccessorFinderTest) Files() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte("taco"), 0400)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte("burrito"), 0400)
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/bar", pfi.RelPath)
	ExpectEq("bar", pfi.Info.Name())
	ExpectEq(len("burrito"), pfi.Info.Size())

	pfi = pfis[1]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectEq(len("taco"), pfi.Info.Size())
}

func (t *FSSuccessorFinderTest) Directories() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create children.
	err = os.Mkdir(path.Join(d, "foo"), 0400)
	AssertEq(nil, err)

	err = os.Mkdir(path.Join(d, "bar"), 0400)
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/bar", pfi.RelPath)
	ExpectEq("bar", pfi.Info.Name())
	ExpectTrue(pfi.Info.IsDir())

	pfi = pfis[1]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectTrue(pfi.Info.IsDir())
}

func (t *FSSuccessorFinderTest) Symlinks() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create a child.
	err = os.Symlink("blah/blah", path.Join(d, "foo"))
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(1, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectFalse(pfi.Info.IsDir())
}

func (t *FSSuccessorFinderTest) Exclusions() {
	var err error

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create some children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0700)
	AssertEq(nil, err)

	err = os.Mkdir(path.Join(d, "bar"), 0700)
	AssertEq(nil, err)

	err = os.Symlink("blah/blah", path.Join(d, "baz"))
	AssertEq(nil, err)

	// Exclude all of them.
	t.exclusions = []*regexp.Regexp{
		regexp.MustCompile("dir/foo"),
		regexp.MustCompile("(bar|baz)"),
	}

	t.resetVisistor()

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)

	AssertEq(nil, err)
	ExpectThat(successors, ElementsAre())
}
<commit_msg>Added tests for parentage.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package save

import (
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
	"sort"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/internal/graph"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
)

func TestFSSuccessorFinder(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type fsNodeSlice []*fsNode

func (p fsNodeSlice) Len() int {
	return len(p)
}

func (p fsNodeSlice) Less(i, j int) bool {
	return p[i].RelPath < p[j].RelPath
}

func (p fsNodeSlice) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}

func sortNodes(graphNodes []graph.Node) (nodes fsNodeSlice) {
	for _, n := range graphNodes {
		nodes = append(nodes, n.(*fsNode))
	}

	sort.Sort(nodes)
	return
}

// A valid os.FileInfo for a directory.
var gDirInfo os.FileInfo

func init() {
	var err error
	gDirInfo, err = os.Stat("/")
	if err != nil {
		log.Fatalf("Stat: %v", err)
	}
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type FSSuccessorFinderTest struct {
	ctx context.Context

	// A temporary directory that is cleaned up at the end of the test. This is
	// the base path with which the successor finder is configured.
	dir string

	// The exclusions with which to configure the successor finder.
	exclusions []*regexp.Regexp

	sf graph.SuccessorFinder
}

var _ SetUpInterface = &FSSuccessorFinderTest{}
var _ TearDownInterface = &FSSuccessorFinderTest{}

func init() { RegisterTestSuite(&FSSuccessorFinderTest{}) }

func (t *FSSuccessorFinderTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Create the base directory.
	var err error
	t.dir, err = ioutil.TempDir("", "file_system_visistor_test")
	AssertEq(nil, err)

	// And the successor finder.
	t.resetVisistor()
}

func (t *FSSuccessorFinderTest) TearDown() {
	var err error

	// Clean up the junk we left in the file system.
	err = os.RemoveAll(t.dir)
	AssertEq(nil, err)
}

func (t *FSSuccessorFinderTest) resetVisistor() {
	t.sf = newSuccessorFinder(t.dir, t.exclusions)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *FSSuccessorFinderTest) NonExistentPath() {
	node := &fsNode{
		RelPath: "foo/bar",
		Info:    gDirInfo,
	}

	_, err := t.sf.FindDirectSuccessors(t.ctx, node)
	ExpectThat(err, Error(HasSubstr(node.RelPath)))
	ExpectThat(err, Error(HasSubstr("no such file")))
}

func (t *FSSuccessorFinderTest) VisitRootNode() {
	var err error

	// Create two children.
	err = ioutil.WriteFile(path.Join(t.dir, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(t.dir, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the root.
	node := &fsNode{
		RelPath: "",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))
	ExpectEq("bar", pfis[0].RelPath)
	ExpectEq("foo", pfis[1].RelPath)
}

func (t *FSSuccessorFinderTest) VisitNonRootNode() {
	var err error

	// Make a few levels of sub-directories.
	d := path.Join(t.dir, "sub/dirs")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0500)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte{}, 0500)
	AssertEq(nil, err)

	// Visit the directory.
	node := &fsNode{
		RelPath: "sub/dirs",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))
	ExpectEq("sub/dirs/bar", pfis[0].RelPath)
	ExpectEq("sub/dirs/foo", pfis[1].RelPath)
}

func (t *FSSuccessorFinderTest) VisitFileNode() {
	var err error

	// Create a file.
	f := path.Join(t.dir, "foo")
	err = ioutil.WriteFile(f, []byte{}, 0500)
	AssertEq(nil, err)

	fi, err := os.Stat(f)
	AssertEq(nil, err)

	// Visit it.
	node := &fsNode{
		RelPath: "foo",
		Info:    fi,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	ExpectThat(successors, ElementsAre())
}

func (t *FSSuccessorFinderTest) Files() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create two children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte("taco"), 0400)
	AssertEq(nil, err)

	err = ioutil.WriteFile(path.Join(d, "bar"), []byte("burrito"), 0400)
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/bar", pfi.RelPath)
	ExpectEq("bar", pfi.Info.Name())
	ExpectEq(len("burrito"), pfi.Info.Size())
	ExpectEq(node, pfi.Parent)

	pfi = pfis[1]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectEq(len("taco"), pfi.Info.Size())
	ExpectEq(node, pfi.Parent)
}

func (t *FSSuccessorFinderTest) Directories() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create children.
	err = os.Mkdir(path.Join(d, "foo"), 0400)
	AssertEq(nil, err)

	err = os.Mkdir(path.Join(d, "bar"), 0400)
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(2, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/bar", pfi.RelPath)
	ExpectEq("bar", pfi.Info.Name())
	ExpectTrue(pfi.Info.IsDir())
	ExpectEq(node, pfi.Parent)

	pfi = pfis[1]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectTrue(pfi.Info.IsDir())
	ExpectEq(node, pfi.Parent)
}

func (t *FSSuccessorFinderTest) Symlinks() {
	var err error
	var pfi *fsNode

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create a child.
	err = os.Symlink("blah/blah", path.Join(d, "foo"))
	AssertEq(nil, err)

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)
	AssertEq(nil, err)

	// Check the output.
	pfis := sortNodes(successors)
	AssertEq(1, len(pfis))

	pfi = pfis[0]
	ExpectEq("dir/foo", pfi.RelPath)
	ExpectEq("foo", pfi.Info.Name())
	ExpectFalse(pfi.Info.IsDir())
	ExpectEq(node, pfi.Parent)
}

func (t *FSSuccessorFinderTest) Exclusions() {
	var err error

	// Make a sub-directory.
	d := path.Join(t.dir, "dir")

	err = os.MkdirAll(d, 0700)
	AssertEq(nil, err)

	// Create some children.
	err = ioutil.WriteFile(path.Join(d, "foo"), []byte{}, 0700)
	AssertEq(nil, err)

	err = os.Mkdir(path.Join(d, "bar"), 0700)
	AssertEq(nil, err)

	err = os.Symlink("blah/blah", path.Join(d, "baz"))
	AssertEq(nil, err)

	// Exclude all of them.
	t.exclusions = []*regexp.Regexp{
		regexp.MustCompile("dir/foo"),
		regexp.MustCompile("(bar|baz)"),
	}

	t.resetVisistor()

	// Visit.
	node := &fsNode{
		RelPath: "dir",
		Info:    gDirInfo,
	}

	successors, err := t.sf.FindDirectSuccessors(t.ctx, node)

	AssertEq(nil, err)
	ExpectThat(successors, ElementsAre())
}
<|endoftext|>


<commit_before>package redis

import (
	"fmt"
	"time"
)

// PubSub implements Pub/Sub commands as described in
// http://redis.io/topics/pubsub.
type PubSub struct {
	*baseClient
}

func (c *Client) PubSub() *PubSub {
	return &PubSub{
		baseClient: &baseClient{
			opt:      c.opt,
			connPool: newSingleConnPool(c.connPool, false),
		},
	}
}

func (c *Client) Publish(channel, message string) *IntCmd {
	req := NewIntCmd("PUBLISH", channel, message)
	c.Process(req)
	return req
}

type Message struct {
	Channel string
	Payload string
}

func (m *Message) String() string {
	return fmt.Sprintf("Message<%s: %s>", m.Channel, m.Payload)
}

type PMessage struct {
	Channel string
	Pattern string
	Payload string
}

func (m *PMessage) String() string {
	return fmt.Sprintf("PMessage<%s: %s>", m.Channel, m.Payload)
}

type Subscription struct {
	Kind    string
	Channel string
	Count   int
}

func (m *Subscription) String() string {
	return fmt.Sprintf("%s: %s", m.Kind, m.Channel)
}

func (c *PubSub) Receive() (interface{}, error) {
	return c.ReceiveTimeout(0)
}

func (c *PubSub) ReceiveTimeout(timeout time.Duration) (interface{}, error) {
	cn, err := c.conn()
	if err != nil {
		return nil, err
	}
	cn.ReadTimeout = timeout

	cmd := NewSliceCmd()
	if err := cmd.parseReply(cn.rd); err != nil {
		return nil, err
	}

	reply := cmd.Val()

	msgName := reply[0].(string)
	switch msgName {
	case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
		return &Subscription{
			Kind:    msgName,
			Channel: reply[1].(string),
			Count:   int(reply[2].(int64)),
		}, nil
	case "message":
		return &Message{
			Channel: reply[1].(string),
			Payload: reply[2].(string),
		}, nil
	case "pmessage":
		return &PMessage{
			Pattern: reply[1].(string),
			Channel: reply[2].(string),
			Payload: reply[3].(string),
		}, nil
	}

	return nil, fmt.Errorf("redis: unsupported message name: %q", msgName)
}

func (c *PubSub) subscribe(cmd string, channels ...string) error {
	cn, err := c.conn()
	if err != nil {
		return err
	}

	args := append([]string{cmd}, channels...)
	req := NewSliceCmd(args...)
	return cn.writeCmds(req)
}

func (c *PubSub) Subscribe(channels ...string) error {
	return c.subscribe("SUBSCRIBE", channels...)
}

func (c *PubSub) PSubscribe(patterns ...string) error {
	return c.subscribe("PSUBSCRIBE", patterns...)
}

func (c *PubSub) unsubscribe(cmd string, channels ...string) error {
	cn, err := c.conn()
	if err != nil {
		return err
	}

	args := append([]string{cmd}, channels...)
	req := NewSliceCmd(args...)
	return cn.writeCmds(req)
}

func (c *PubSub) Unsubscribe(channels ...string) error {
	return c.unsubscribe("UNSUBSCRIBE", channels...)
}

func (c *PubSub) PUnsubscribe(patterns ...string) error {
	return c.unsubscribe("PUNSUBSCRIBE", patterns...)
}
<commit_msg>More comments.<commit_after>package redis

import (
	"fmt"
	"time"
)

// PubSub implements Pub/Sub commands as described in
// http://redis.io/topics/pubsub.
type PubSub struct {
	*baseClient
}

func (c *Client) PubSub() *PubSub {
	return &PubSub{
		baseClient: &baseClient{
			opt:      c.opt,
			connPool: newSingleConnPool(c.connPool, false),
		},
	}
}

func (c *Client) Publish(channel, message string) *IntCmd {
	req := NewIntCmd("PUBLISH", channel, message)
	c.Process(req)
	return req
}

// Message received as result of a PUBLISH command issued by another client.
type Message struct {
	Channel string
	Payload string
}

func (m *Message) String() string {
	return fmt.Sprintf("Message<%s: %s>", m.Channel, m.Payload)
}

// Message matching a pattern-matching subscription received as result
// of a PUBLISH command issued by another client.
type PMessage struct {
	Channel string
	Pattern string
	Payload string
}

func (m *PMessage) String() string {
	return fmt.Sprintf("PMessage<%s: %s>", m.Channel, m.Payload)
}

// Message received after a successful subscription to channel.
type Subscription struct {
	// Can be "subscribe", "unsubscribe", "psubscribe" or "punsubscribe".
	Kind string
	// Channel name we have subscribed to.
	Channel string
	// Number of channels we are currently subscribed to.
	Count int
}

func (m *Subscription) String() string {
	return fmt.Sprintf("%s: %s", m.Kind, m.Channel)
}

func (c *PubSub) Receive() (interface{}, error) {
	return c.ReceiveTimeout(0)
}

func (c *PubSub) ReceiveTimeout(timeout time.Duration) (interface{}, error) {
	cn, err := c.conn()
	if err != nil {
		return nil, err
	}
	cn.ReadTimeout = timeout

	cmd := NewSliceCmd()
	if err := cmd.parseReply(cn.rd); err != nil {
		return nil, err
	}

	reply := cmd.Val()

	msgName := reply[0].(string)
	switch msgName {
	case "subscribe", "unsubscribe", "psubscribe", "punsubscribe":
		return &Subscription{
			Kind:    msgName,
			Channel: reply[1].(string),
			Count:   int(reply[2].(int64)),
		}, nil
	case "message":
		return &Message{
			Channel: reply[1].(string),
			Payload: reply[2].(string),
		}, nil
	case "pmessage":
		return &PMessage{
			Pattern: reply[1].(string),
			Channel: reply[2].(string),
			Payload: reply[3].(string),
		}, nil
	}

	return nil, fmt.Errorf("redis: unsupported message name: %q", msgName)
}

func (c *PubSub) subscribe(cmd string, channels ...string) error {
	cn, err := c.conn()
	if err != nil {
		return err
	}

	args := append([]string{cmd}, channels...)
	req := NewSliceCmd(args...)
	return cn.writeCmds(req)
}

func (c *PubSub) Subscribe(channels ...string) error {
	return c.subscribe("SUBSCRIBE", channels...)
}

func (c *PubSub) PSubscribe(patterns ...string) error {
	return c.subscribe("PSUBSCRIBE", patterns...)
}

func (c *PubSub) unsubscribe(cmd string, channels ...string) error {
	cn, err := c.conn()
	if err != nil {
		return err
	}

	args := append([]string{cmd}, channels...)
	req := NewSliceCmd(args...)
	return cn.writeCmds(req)
}

func (c *PubSub) Unsubscribe(channels ...string) error {
	return c.unsubscribe("UNSUBSCRIBE", channels...)
}

func (c *PubSub) PUnsubscribe(patterns ...string) error {
	return c.unsubscribe("PUNSUBSCRIBE", patterns...)
}
<|endoftext|>


<commit_before>package py

/*
#include "Python.h"
*/
import "C"
import (
	"fmt"
	"pfi/sensorbee/sensorbee/data"
	"runtime"
	"unsafe"
)

// ObjectFunc is a bind of `PyObject` used as `PyFunc`
type ObjectFunc struct {
	Object
}

// invokeDirect calls name's function. User needs to call DecRef.
// This returns an Object even if result values are more than one.
// For example, use to get the object of the class instance that method returned.
func invokeDirect(pyObj *C.PyObject, name string, args []data.Value,
	kwdArgs data.Map) (Object, error) {
	type Result struct {
		val Object
		err error
	}
	ch := make(chan *Result, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				ch <- &Result{Object{}, fmt.Errorf(
					"cannot call '%v' due to panic: %v", name, r)}
			}
		}()

		runtime.LockOSThread()
		state := GILState_Ensure()
		defer GILState_Release(state)

		ret, err := callMethod(pyObj, name, kwdArgs, args...)
		ch <- &Result{ret, err}
	}()
	res := <-ch

	return res.val, res.err
}

// invoke name's function. TODO should be placed at internal package.
func invoke(pyObj *C.PyObject, name string, args []data.Value, kwdArgs data.Map) (
	data.Value, error) {
	type Result struct {
		val data.Value
		err error
	}
	ch := make(chan *Result, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				ch <- &Result{data.Null{}, fmt.Errorf(
					"cannot call '%v' due to panic: %v", name, r)}
			}
		}()

		runtime.LockOSThread()
		state := GILState_Ensure()
		defer GILState_Release(state)

		ret, err := callMethod(pyObj, name, kwdArgs, args...)
		if err != nil {
			ch <- &Result{data.Null{}, err}
			return
		}
		defer ret.decRef()

		po, err := fromPyTypeObject(ret.p)
		ch <- &Result{po, err}
	}()
	res := <-ch

	return res.val, res.err
}

// callMethod calls `name` method on `pyObj`. This function is not locked GIL
func callMethod(pyObj *C.PyObject, name string, kwdArgs data.Map,
	args ...data.Value) (Object, error) {
	pyFunc, err := getPyFunc(pyObj, name)
	if err != nil {
		return Object{}, fmt.Errorf("fail to get '%v' function: %v", name,
			err.Error())
	}
	defer pyFunc.decRef()

	// no named arguments
	pyArg, err := convertArgsGo2Py(args)
	if err != nil {
		return Object{}, fmt.Errorf(
			"fail to convert argument in calling '%v' function: %v", name,
			err.Error())
	}
	defer pyArg.decRef()

	// named arguments
	var ret Object
	if len(kwdArgs) == 0 {
		ret, err = pyFunc.callObject(pyArg)
	} else {
		pyKwdArg, err := newPyObj(kwdArgs)
		if err != nil {
			return Object{}, fmt.Errorf(
				"fail to convert named arguments in calling '%v' function: %v",
				name, err.Error())
		}
		ret, err = pyFunc.callObjectWithKwd(pyKwdArg, pyArg)
	}

	if err != nil {
		return Object{}, fmt.Errorf("fail to call '%v' function: %v", name,
			err.Error())
	}
	return ret, nil
}

// TODO should be placed at internal package
func getPyFunc(pyObj *C.PyObject, name string) (ObjectFunc, error) {
	cFunc := C.CString(name)
	defer C.free(unsafe.Pointer(cFunc))

	pyFunc := C.PyObject_GetAttrString(pyObj, cFunc)
	if pyFunc == nil {
		return ObjectFunc{}, getPyErr()
	}

	if ok := C.PyCallable_Check(pyFunc); ok == 0 {
		return ObjectFunc{}, fmt.Errorf("'%v' is not callable object", name)
	}

	return ObjectFunc{Object{p: pyFunc}}, nil
}

// callObject executes python function, using `PyObject_CallObject`. Returns a
// `PyObject` even if result values are more than one. When a value will be set
// directory, and values will be set as a `PyTuple` object.
func (f *ObjectFunc) callObject(arg Object) (po Object, err error) {
	po = Object{
		p: C.PyObject_CallObject(f.p, arg.p),
	}
	if po.p == nil {
		err = getPyErr()
	}
	return
}

// callObjectWithKwd executes python function, using `PyObject_Call`. Error
// specification is same as `callObject`.
func (f *ObjectFunc) callObjectWithKwd(kwdArg Object, arg Object) (po Object,
	err error) {
	po = Object{
		p: C.PyObject_Call(f.p, arg.p, kwdArg.p),
	}
	if po.p == nil {
		err = getPyErr()
	}
	return
}

func convertArgsGo2Py(args []data.Value) (Object, error) {
	pyArg := C.PyTuple_New(C.Py_ssize_t(len(args)))
	if pyArg == nil {
		return Object{}, getPyErr()
	}
	shouldDecRef := true
	defer func() {
		if shouldDecRef {
			C.Py_DecRef(pyArg)
		}
	}()
	for i, v := range args {
		o, err := newPyObj(v)
		if err != nil {
			return Object{}, err
		}
		// PyTuple object takes over the value's reference, and not need to
		// decrease reference counter.
		C.PyTuple_SetItem(pyArg, C.Py_ssize_t(i), o.p)
	}
	shouldDecRef = false
	return Object{pyArg}, nil
}
<commit_msg>modify callmethod signature, remove named return value<commit_after>package py

/*
#include "Python.h"
*/
import "C"
import (
	"fmt"
	"pfi/sensorbee/sensorbee/data"
	"runtime"
	"unsafe"
)

// ObjectFunc is a bind of `PyObject` used as `PyFunc`
type ObjectFunc struct {
	Object
}

// invokeDirect calls name's function. User needs to call DecRef.
// This returns an Object even if result values are more than one.
// For example, use to get the object of the class instance that method returned.
func invokeDirect(pyObj *C.PyObject, name string, args []data.Value,
	kwdArgs data.Map) (Object, error) {
	type Result struct {
		val Object
		err error
	}
	ch := make(chan *Result, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				ch <- &Result{Object{}, fmt.Errorf(
					"cannot call '%v' due to panic: %v", name, r)}
			}
		}()

		runtime.LockOSThread()
		state := GILState_Ensure()
		defer GILState_Release(state)

		ret, err := callMethod(pyObj, name, args, kwdArgs)
		ch <- &Result{ret, err}
	}()
	res := <-ch

	return res.val, res.err
}

// invoke name's function. TODO should be placed at internal package.
func invoke(pyObj *C.PyObject, name string, args []data.Value, kwdArgs data.Map) (
	data.Value, error) {
	type Result struct {
		val data.Value
		err error
	}
	ch := make(chan *Result, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				ch <- &Result{data.Null{}, fmt.Errorf(
					"cannot call '%v' due to panic: %v", name, r)}
			}
		}()

		runtime.LockOSThread()
		state := GILState_Ensure()
		defer GILState_Release(state)

		ret, err := callMethod(pyObj, name, args, kwdArgs)
		if err != nil {
			ch <- &Result{data.Null{}, err}
			return
		}
		defer ret.decRef()

		po, err := fromPyTypeObject(ret.p)
		ch <- &Result{po, err}
	}()
	res := <-ch

	return res.val, res.err
}

// callMethod calls `name` method on `pyObj`. This function is not locked GIL
func callMethod(pyObj *C.PyObject, name string, args []data.Value,
	kwdArgs data.Map) (Object, error) {
	pyFunc, err := getPyFunc(pyObj, name)
	if err != nil {
		return Object{}, fmt.Errorf("fail to get '%v' function: %v", name,
			err.Error())
	}
	defer pyFunc.decRef()

	// no named arguments
	pyArg, err := convertArgsGo2Py(args)
	if err != nil {
		return Object{}, fmt.Errorf(
			"fail to convert argument in calling '%v' function: %v", name,
			err.Error())
	}
	defer pyArg.decRef()

	// named arguments
	var ret Object
	if len(kwdArgs) == 0 {
		ret, err = pyFunc.callObject(pyArg)
	} else {
		pyKwdArg, err := newPyObj(kwdArgs)
		if err != nil {
			return Object{}, fmt.Errorf(
				"fail to convert named arguments in calling '%v' function: %v",
				name, err.Error())
		}
		ret, err = pyFunc.callObjectWithKwd(pyKwdArg, pyArg)
	}

	if err != nil {
		return Object{}, fmt.Errorf("fail to call '%v' function: %v", name,
			err.Error())
	}
	return ret, nil
}

// TODO should be placed at internal package
func getPyFunc(pyObj *C.PyObject, name string) (ObjectFunc, error) {
	cFunc := C.CString(name)
	defer C.free(unsafe.Pointer(cFunc))

	pyFunc := C.PyObject_GetAttrString(pyObj, cFunc)
	if pyFunc == nil {
		return ObjectFunc{}, getPyErr()
	}

	if ok := C.PyCallable_Check(pyFunc); ok == 0 {
		return ObjectFunc{}, fmt.Errorf("'%v' is not callable object", name)
	}

	return ObjectFunc{Object{p: pyFunc}}, nil
}

// callObject executes python function, using `PyObject_CallObject`. Returns a
// `PyObject` even if result values are more than one. When a value will be set
// directory, and values will be set as a `PyTuple` object.
func (f *ObjectFunc) callObject(arg Object) (Object, error) {
	po := C.PyObject_CallObject(f.p, arg.p)
	if po == nil {
		return Object{}, getPyErr()
	}
	return Object{p: po}, nil
}

// callObjectWithKwd executes python function, using `PyObject_Call`. Error
// specification is same as `callObject`.
func (f *ObjectFunc) callObjectWithKwd(kwdArg Object, arg Object) (Object, error) {
	po := C.PyObject_Call(f.p, arg.p, kwdArg.p)
	if po == nil {
		return Object{}, getPyErr()
	}
	return Object{p: po}, nil
}

func convertArgsGo2Py(args []data.Value) (Object, error) {
	pyArg := C.PyTuple_New(C.Py_ssize_t(len(args)))
	if pyArg == nil {
		return Object{}, getPyErr()
	}
	shouldDecRef := true
	defer func() {
		if shouldDecRef {
			C.Py_DecRef(pyArg)
		}
	}()
	for i, v := range args {
		o, err := newPyObj(v)
		if err != nil {
			return Object{}, err
		}
		// PyTuple object takes over the value's reference, and not need to
		// decrease reference counter.
		C.PyTuple_SetItem(pyArg, C.Py_ssize_t(i), o.p)
	}
	shouldDecRef = false
	return Object{pyArg}, nil
}
<|endoftext|>


<commit_before>package ransac

import (
	"math"
	"math/rand"
)

type calcError func(map[string]float64, [2]float64) float64
type calcModel func([][2]float64) map[string]float64

// Problem is the formulation of the ransac problem.
type Problem struct {
	data  [][2]float64
	fit   calcError
	model calcModel
}

func (p Problem) sample(sampleSize int) [][2]float64 {
	sample := make([][2]float64, sampleSize)
	currentSample := 0

	for currentSample < sampleSize {
		randomIndex := random(0, len(p.data))
		exists := existsInData(sample, p.data[randomIndex])
		if !exists {
			sample[currentSample] = p.data[randomIndex]
			currentSample++
		}
	}
	return sample
}

func (p Problem) calcModelError(model map[string]float64) float64 {
	var ssd float64
	for _, point := range p.data {
		error := p.fit(model, point)
		ssd += error * error
	}
	return ssd
}

// SetModel sets a function which describes the mode.
func (p *Problem) SetModel(fn calcModel) {
	p.model = fn
}

// SetModelError sets the function, which calculates the error.
func (p *Problem) SetModelError(fn calcError) {
	p.fit = fn
}

// SetData sets the function, which sets the dataset.
func (p *Problem) SetData(data [][2]float64) {
	p.data = data
}

func (p Problem) classifyInliers(model map[string]float64, sample [][2]float64, maxError float64) ([][2]float64, [][2]float64) {
	var inliers [][2]float64
	var outliers [][2]float64

	for _, point := range p.data {
		if !existsInData(sample, point) {
			error := p.fit(model, point)
			if error <= maxError {
				inliers = append(inliers, point)
			} else {
				outliers = append(outliers, point)
			}
		}
	}
	return inliers, outliers
}

// Estimate does the actual work of fitting.
func (p Problem) Estimate(maxIterations, sampleSize int, inliersRatioLimit float64, maxError float64, improveWithConsensusSet bool) (map[string]float64, [][2]float64, [][2]float64, float64) {
	var bestInliers [][2]float64
	var bestOutliers [][2]float64
	var bestModel map[string]float64
	var bestError float64 = math.Inf(1)

	for iteration := 0; iteration <= maxIterations; iteration++ {
		sample := p.sample(sampleSize)
		model := p.model(sample)
		inliers, outliers := p.classifyInliers(model, sample, maxError)
		inliersRatio := float64(len(inliers)) / float64(len(p.data))
		if inliersRatio >= inliersRatioLimit {
			candidateModel := model
			if improveWithConsensusSet {
				candidateModel = p.model(inliers)
			}
			candidateError := p.calcModelError(model)
			if candidateError < bestError {
				bestInliers = inliers
				bestOutliers = outliers
				bestModel = candidateModel
				bestError = candidateError
			}
		}
	}

	return bestModel, bestInliers, bestOutliers, bestError
}

func random(min, max int) int {
	return rand.Intn(max-min) + min
}

func existsInData(data [][2]float64, sample [2]float64) bool {
	for _, point := range data {
		if point[0] == sample[0] && point[1] == sample[1] {
			return true
		}
	}
	return false
}
<commit_msg>performance improvements<commit_after>package ransac

import (
	"math"
	"math/rand"
)

type calcError func(map[string]float64, [2]float64) float64
type calcModel func([][2]float64) map[string]float64

// Problem is the formulation of the ransac problem.
type Problem struct {
	data       [][2]float64
	dataLength int
	fit        calcError
	model      calcModel
}

func (p *Problem) sample(sampleSize int) [][2]float64 {
	sample := make([][2]float64, sampleSize)
	currentSample := 0

	for currentSample < sampleSize {
		randomIndex := random(0, p.dataLength)
		exists := existsInData(sample, p.data[randomIndex])
		if !exists {
			sample[currentSample] = p.data[randomIndex]
			currentSample++
		}
	}
	return sample
}

func (p *Problem) calcModelError(model map[string]float64) float64 {
	var ssd float64
	for _, point := range p.data {
		error := p.fit(model, point)
		ssd += error * error
	}
	return ssd
}

// SetModel sets a function which describes the mode.
func (p *Problem) SetModel(fn calcModel) {
	p.model = fn
}

// SetModelError sets the function, which calculates the error.
func (p *Problem) SetModelError(fn calcError) {
	p.fit = fn
}

// SetData sets the function, which sets the dataset.
func (p *Problem) SetData(data [][2]float64) {
	p.data = data
	p.dataLength = len(data)
}

func (p *Problem) classifyInliers(model map[string]float64, sample [][2]float64, maxError float64) ([][2]float64, [][2]float64) {
	var inliers [][2]float64
	var outliers [][2]float64

	for _, point := range p.data {
		if !existsInData(sample, point) {
			error := p.fit(model, point)
			if error <= maxError {
				inliers = append(inliers, point)
			} else {
				outliers = append(outliers, point)
			}
		}
	}
	return inliers, outliers
}

// Estimate does the actual work of fitting.
func (p *Problem) Estimate(maxIterations, sampleSize int, inliersRatioLimit float64, maxError float64, improveWithConsensusSet bool) (map[string]float64, [][2]float64, [][2]float64, float64) {
	var bestInliers [][2]float64
	var bestOutliers [][2]float64
	var bestModel map[string]float64
	var bestError float64 = math.Inf(1)
	dataLength := float64(p.dataLength)

	for iteration := 0; iteration <= maxIterations; iteration++ {
		sample := p.sample(sampleSize)
		model := p.model(sample)
		inliers, outliers := p.classifyInliers(model, sample, maxError)
		inliersRatio := float64(len(inliers)) / dataLength
		if inliersRatio >= inliersRatioLimit {
			candidateModel := model
			if improveWithConsensusSet {
				candidateModel = p.model(inliers)
			}
			candidateError := p.calcModelError(candidateModel)
			if candidateError < bestError {
				bestInliers = inliers
				bestOutliers = outliers
				bestModel = candidateModel
				bestError = candidateError
			}
		}
	}

	return bestModel, bestInliers, bestOutliers, bestError
}

func random(min, max int) int {
	return rand.Intn(max-min) + min
}

func existsInData(data [][2]float64, sample [2]float64) bool {
	for _, point := range data {
		if point[0] == sample[0] && point[1] == sample[1] {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"github.com/miekg/dns"
	"github.com/montanaflynn/stats"
	"log"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

type Msg struct {
	Times    []float64
	ErrCount int
}

const VERSION = "1.0.0"

func main() {
	var passes int
	var concurrency int
	var resolver string
	var name string
	log.Printf("rpounder resolver benchmark %s starting", VERSION)
	flag.IntVar(&passes, "p", 100, "Number of passes")
	flag.IntVar(&concurrency, "c", 10, "Concurrent processes")
	flag.StringVar(&resolver, "r", "localhost:53", "resolver(s) ip/port. Separate multiple resolvers with spaces or commas")
	flag.StringVar(&name, "n", "example.com", "Hostname(s) to look up. Separate multiple hostnames with spaces or commas")
	flag.Parse()
	passes_per_goproc := passes / concurrency
	if passes_per_goproc == 0 {
		log.Fatalf("concurrency needs to be greater than passes")
	}
	mod := passes % concurrency
	c := make(chan Msg)
	q := make(chan struct{})

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigs
		log.Printf("=== INTERRUPT: sig = %v", sig)
		close(q)
	}()

	for i := 0; i < concurrency; i++ {
		go func() {
			num := passes_per_goproc
			if mod > 0 {
				num += 1
				mod--
			}
			pound(num, resolver, name, c, q)
		}()
	}

	i := 0
	times := make([]float64, 0, passes)
	errs := 0
	for msg := range c {
		errs += msg.ErrCount
		times = append(times, msg.Times...)
		i++
		if i >= concurrency {
			break
		}
	}
	// Print out percentiles
	log.Printf("Percentage of the requests served within a certain time:")
	pctls := []int{50, 66, 75, 80, 90, 95, 99}
	for _, rank := range pctls {
		v, _ := stats.Percentile(times, float64(rank))
		log.Printf("%d%%  => %s", rank, time.Duration(v))
	}
	min, _ := stats.Min(times)
	max, _ := stats.Max(times)
	log.Printf("Fastest: %s  -- Slowest: %s", time.Duration(min), time.Duration(max))
	log.Printf("Total passes: %d. Total errors: %d", len(times), errs)
}

func pound(passes int, res, hostnames string, c chan Msg, q chan struct{}) {
	sprex := regexp.MustCompile("[\\s\\,]+")
	lookup_hosts := sprex.Split(hostnames, -1)
	msg := Msg{make([]float64, 0, passes), 0}
	defer func() {
		c <- msg
	}()
	resolvers := sprex.Split(res, -1)
	resolver := NewResolver(resolvers)
	var err error
	hi := 0
	for i := 0; i < passes; i++ {
		select {
		case <-q:
			return
		default:
		}
		hostname := lookup_hosts[hi]
		hi++
		if hi >= len(lookup_hosts) {
			hi = 0
		}
		start := time.Now()
		_, err = resolver.LookupHost(hostname)
		end := float64(time.Since(start))
		if err != nil {
			msg.ErrCount++
		}
		msg.Times = append(msg.Times, end)
	}
}

// DnsResolver represents a dns resolver
// Credit where credit is due -- this code was taken and modified from
// https://github.com/bogdanovich/dns_resolver:
//The MIT License (MIT)

//Copyright (c) 2015 Anton Bogdanovich

//Permission is hereby granted, free of charge, to any person obtaining a copy
//of this software and associated documentation files (the "Software"), to deal
//in the Software without restriction, including without limitation the rights
//to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//copies of the Software, and to permit persons to whom the Software is
//furnished to do so, subject to the following conditions:

//The above copyright notice and this permission notice shall be included in all
//copies or substantial portions of the Software.

//THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//SOFTWARE.

type DnsResolver struct {
	Servers []string
	r       *rand.Rand
}

// New initializes DnsResolver.
func NewResolver(servers []string) *DnsResolver {
	for i := range servers {
		if len(strings.Split(servers[i], ":")) < 2 {
			servers[i] = net.JoinHostPort(servers[i], "53")
		}
	}
	return &DnsResolver{servers, rand.New(rand.NewSource(time.Now().UnixNano()))}
}

// LookupHost returns IP addresses of provied host.
func (r *DnsResolver) LookupHost(host string) ([]net.IP, error) {
	return r.lookupHost(host)
}

func (r *DnsResolver) lookupHost(host string) ([]net.IP, error) {
	m1 := new(dns.Msg)
	m1.Id = dns.Id()
	m1.RecursionDesired = true
	m1.Question = make([]dns.Question, 1)
	m1.Question[0] = dns.Question{dns.Fqdn(host), dns.TypeA, dns.ClassINET}
	in, err := dns.Exchange(m1, r.Servers[r.r.Intn(len(r.Servers))])

	result := []net.IP{}

	if err != nil {
		return result, err
	}

	if in != nil && in.Rcode != dns.RcodeSuccess {
		return result, errors.New(dns.RcodeToString[in.Rcode])
	}

	for _, record := range in.Answer {
		if t, ok := record.(*dns.A); ok {
			result = append(result, t.A)
		}
	}
	return result, err
}
<commit_msg>redo resolver logic<commit_after>package main

import (
	"errors"
	"flag"
	"github.com/miekg/dns"
	"github.com/montanaflynn/stats"
	"log"
	"net"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"syscall"
	"time"
)

type Msg struct {
	Times    []float64
	ErrCount int
}

const VERSION = "1.0.0"

func main() {
	var passes int
	var concurrency int
	var resolver string
	var name string
	log.Printf("rpounder resolver benchmark %s starting", VERSION)
	flag.IntVar(&passes, "p", 100, "Number of passes")
	flag.IntVar(&concurrency, "c", 10, "Concurrent processes")
	flag.StringVar(&resolver, "r", "localhost:53", "resolver(s) ip/port. Separate multiple resolvers with spaces or commas")
	flag.StringVar(&name, "n", "example.com", "Hostname(s) to look up. Separate multiple hostnames with spaces or commas")
	flag.Parse()
	passes_per_goproc := passes / concurrency
	if passes_per_goproc == 0 {
		log.Fatalf("concurrency needs to be greater than passes")
	}
	mod := passes % concurrency
	c := make(chan Msg)
	q := make(chan struct{})

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigs
		log.Printf("=== INTERRUPT: sig = %v", sig)
		close(q)
	}()

	for i := 0; i < concurrency; i++ {
		go func() {
			num := passes_per_goproc
			if mod > 0 {
				num += 1
				mod--
			}
			pound(num, resolver, name, c, q)
		}()
	}

	i := 0
	times := make([]float64, 0, passes)
	errs := 0
	for msg := range c {
		errs += msg.ErrCount
		times = append(times, msg.Times...)
		i++
		if i >= concurrency {
			break
		}
	}
	// Print out percentiles
	log.Printf("Percentage of the requests served within a certain time:")
	pctls := []int{50, 66, 75, 80, 90, 95, 99}
	for _, rank := range pctls {
		v, _ := stats.Percentile(times, float64(rank))
		log.Printf("%d%%  => %s", rank, time.Duration(v))
	}
	min, _ := stats.Min(times)
	max, _ := stats.Max(times)
	log.Printf("Fastest: %s  -- Slowest: %s", time.Duration(min), time.Duration(max))
	log.Printf("Total passes: %d. Total errors: %d", len(times), errs)
}

func pound(passes int, res, hostnames string, c chan Msg, q chan struct{}) {
	sprex := regexp.MustCompile("[\\s\\,]+")
	lookup_hosts := sprex.Split(hostnames, -1)
	msg := Msg{make([]float64, 0, passes), 0}
	defer func() {
		c <- msg
	}()
	resolvers := sprex.Split(res, -1)
	resolver := NewResolver(resolvers)
	var err error
	hi := 0
	for i := 0; i < passes; i++ {
		select {
		case <-q:
			return
		default:
		}
		hostname := lookup_hosts[hi]
		hi++
		if hi >= len(lookup_hosts) {
			hi = 0
		}
		start := time.Now()
		_, err = resolver.LookupHost(hostname)
		end := float64(time.Since(start))
		if err != nil {
			msg.ErrCount++
		}
		msg.Times = append(msg.Times, end)
	}
}

type Resolver struct {
	servers []string
	idx     int
}

// New initializes DnsResolver.
func NewResolver(servers []string) *Resolver {
	for i := range servers {
		if len(strings.Split(servers[i], ":")) < 2 {
			servers[i] = net.JoinHostPort(servers[i], "53")
		}
	}
	return &Resolver{servers, 0}
}

// LookupHost returns A record for host. Round robin if multiple resolvers
func (r *Resolver) LookupHost(host string) ([]net.IP, error) {
	defer func() {
		r.idx++
	}()
	if r.idx >= len(r.servers) {
		r.idx = 0
	}
	m := &dns.Msg{}
	m.Id = dns.Id()
	m.RecursionDesired = true
	m.Question = make([]dns.Question, 1)
	m.Question[0] = dns.Question{dns.Fqdn(host), dns.TypeA, dns.ClassINET}
	in, err := dns.Exchange(m, r.servers[r.idx])

	result := []net.IP{}

	if err != nil {
		return result, err
	}

	if in != nil && in.Rcode != dns.RcodeSuccess {
		return result, errors.New(dns.RcodeToString[in.Rcode])
	}

	for _, record := range in.Answer {
		if t, ok := record.(*dns.A); ok {
			result = append(result, t.A)
		}
	}
	return result, err
}
<|endoftext|>


<commit_before>// This source file is part of the Packet Guardian project.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package server

import (
	"fmt"
	"net/http"
	"net/http/pprof"
	"runtime"
	"strings"

	"github.com/dchest/captcha"
	assetfs "github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/context"
	"github.com/julienschmidt/httprouter"
	"github.com/lfkeitel/verbose/v4"

	"github.com/packet-guardian/dhcp-lib"
	"github.com/packet-guardian/packet-guardian/src/auth"
	"github.com/packet-guardian/packet-guardian/src/bindata"
	"github.com/packet-guardian/packet-guardian/src/common"
	"github.com/packet-guardian/packet-guardian/src/controllers"
	"github.com/packet-guardian/packet-guardian/src/controllers/api"
	"github.com/packet-guardian/packet-guardian/src/models"
	"github.com/packet-guardian/packet-guardian/src/models/stores"
	mid "github.com/packet-guardian/packet-guardian/src/server/middleware"
)

func LoadRoutes(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	r.Handler("GET", "/", midStack(e, stores, &rootHandler{leases: stores.Leases}))
	r.ServeFiles(
		"/public/*filepath",
		&assetfs.AssetFS{
			Asset:     bindata.GetAsset,
			AssetDir:  bindata.GetAssetDir,
			AssetInfo: bindata.GetAssetInfo,
			Prefix:    "public"})

	authController := controllers.NewAuthController(e, stores.Users)
	r.Handler("GET", "/login", midStack(e, stores, http.HandlerFunc(authController.LoginHandler)))
	r.Handler("POST", "/login", midStack(e, stores, http.HandlerFunc(authController.LoginHandler)))
	r.Handler("GET", "/logout", midStack(e, stores, http.HandlerFunc(authController.LogoutHandler)))

	manageController := controllers.NewManagerController(e, stores.Devices, stores.Leases)
	r.Handler("GET", "/register", midStack(e, stores, http.HandlerFunc(manageController.RegistrationHandler)))
	r.Handler("GET", "/manage", midStack(e, stores, mid.CheckAuth(http.HandlerFunc(manageController.ManageHandler))))

	guestController := controllers.NewGuestController(e, stores.Users, stores.Devices, stores.Leases)
	r.Handler("GET", "/register/guest", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.RegistrationHandler), e, stores.Leases)))
	r.Handler("POST", "/register/guest", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.RegistrationHandler), e, stores.Leases)))
	r.Handler("GET", "/register/guest/verify", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.VerificationHandler), e, stores.Leases)))
	r.Handler("POST", "/register/guest/verify", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.VerificationHandler), e, stores.Leases)))

	r.Handler("GET", "/admin/*a", midStack(e, stores, adminRouter(e, stores)))
	r.Handler("GET", "/api/*a", midStack(e, stores, apiRouter(e, stores)))
	r.Handler("POST", "/api/*a", midStack(e, stores, apiRouter(e, stores)))
	r.Handler("DELETE", "/api/*a", midStack(e, stores, apiRouter(e, stores)))

	r.Handler("GET", "/captcha/*a", captcha.Server(captcha.StdWidth, captcha.StdHeight))

	if e.IsDev() {
		r.Handler("GET", "/dev/*a", midStack(e, stores, devRouter(e)))
		r.Handler("GET", "/debug/*a", midStack(e, stores, debugRouter(e)))
		e.Log.Debug("Profiling enabled")
	}

	h := mid.Logging(r, e) // Logging
	h = mid.Panic(h, e)    // Panic catcher
	return h
}

func midStack(e *common.Environment, stores stores.StoreCollection, h http.Handler) http.Handler {
	h = mid.BlacklistCheck(h, e, stores.Devices, stores.Leases) // Enforce a blacklist check
	h = mid.Cache(h, e)                                         // Set cache headers if needed
	h = mid.SetSessionInfo(h, e, stores.Users)                  // Adds Environment and user information to requet context
	h = context.ClearHandler(h)                                 // Clear Gorilla sessions
	return h
}

func devRouter(e *common.Environment) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	devController := controllers.NewDevController(e)
	r.HandlerFunc("GET", "/dev/reloadtemp", devController.ReloadTemplates)

	h := mid.CheckAdmin(r)
	h = mid.CheckAuth(h)
	return h
}

func debugRouter(e *common.Environment) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	r.HandlerFunc("GET", "/debug/pprof", pprof.Index)
	r.HandlerFunc("GET", "/debug/pprof/cmdline", pprof.Cmdline)
	r.HandlerFunc("GET", "/debug/pprof/profile", pprof.Profile)
	r.HandlerFunc("GET", "/debug/pprof/symbol", pprof.Symbol)
	r.HandlerFunc("GET", "/debug/pprof/trace", pprof.Trace)
	// Manually add support for paths linked to by index page at /debug/pprof/
	r.Handler("GET", "/debug/pprof/goroutine", pprof.Handler("goroutine"))
	r.Handler("GET", "/debug/pprof/heap", pprof.Handler("heap"))
	r.Handler("GET", "/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
	r.Handler("GET", "/debug/pprof/block", pprof.Handler("block"))

	r.HandlerFunc("GET", "/debug/heap-stats", heapStats)

	h := mid.CheckAdmin(r)
	h = mid.CheckAuth(h)
	return h
}

func heapStats(w http.ResponseWriter, r *http.Request) {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	fmt.Fprintf(w,
		"HeapSys: %d, HeapAlloc: %d, HeapIdle: %d, HeapReleased: %d\n",
		m.HeapSys,
		m.HeapAlloc,
		m.HeapIdle,
		m.HeapReleased,
	)
}

func adminRouter(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	adminController := controllers.NewAdminController(e, stores)
	r.GET("/admin/", adminController.DashboardHandler)
	r.GET("/admin/search", adminController.SearchHandler)
	r.GET("/admin/manage/user/:username", adminController.ManageHandler)
	r.GET("/admin/manage/device/:mac", adminController.ShowDeviceHandler)
	r.GET("/admin/users", adminController.AdminUserListHandler)
	r.GET("/admin/users/:username", adminController.AdminUserHandler)
	r.GET("/admin/reports", adminController.ReportHandler)
	r.GET("/admin/reports/:report", adminController.ReportHandler)

	h := mid.CheckAdmin(r)
	h = mid.CheckAuth(h)
	return h
}

func apiRouter(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()

	deviceAPIController := api.NewDeviceController(e, stores.Users, stores.Devices, stores.Leases)
	r.POST("/api/device", deviceAPIController.RegistrationHandler)
	r.DELETE("/api/device/user/:username", deviceAPIController.DeleteHandler)
	r.POST("/api/device/reassign", deviceAPIController.ReassignHandler)
	r.POST("/api/device/mac/:mac/description", deviceAPIController.EditDescriptionHandler)
	r.POST("/api/device/mac/:mac/expiration", deviceAPIController.EditExpirationHandler)
	r.POST("/api/device/mac/:mac/flag", deviceAPIController.EditFlaggedHandler)
	r.GET("/api/device/:mac", deviceAPIController.GetDeviceHandler)

	blacklistController := api.NewBlacklistController(e, stores.Users, stores.Devices)
	r.POST("/api/blacklist/user/:username", blacklistController.BlacklistUserHandler)
	r.DELETE("/api/blacklist/user/:username", blacklistController.BlacklistUserHandler)

	r.POST("/api/blacklist/device", blacklistController.BlacklistDeviceHandler)
	r.DELETE("/api/blacklist/device", blacklistController.BlacklistDeviceHandler)

	userAPIController := api.NewUserController(e, stores.Users, stores.Devices)
	r.POST("/api/user", userAPIController.UserHandler)
	r.DELETE("/api/user", userAPIController.UserHandler)
	r.GET("/api/user/:username", userAPIController.UserHandler)

	statusAPIController := api.NewStatusController(e)
	r.GET("/api/status", statusAPIController.GetStatus)

	return mid.CheckAuthAPI(r, stores.Users)
}

type rootHandler struct {
	leases stores.LeaseStore
}

func (h *rootHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if auth.IsLoggedIn(r) {
		sessionUser := models.GetUserFromContext(r)
		if sessionUser.Can(models.ViewAdminPage) {
			http.Redirect(w, r, "/admin", http.StatusTemporaryRedirect)
			return
		}

		http.Redirect(w, r, "/manage", http.StatusTemporaryRedirect)
		return
	}

	e := common.GetEnvironmentFromContext(r)
	ip := common.GetIPFromContext(r)
	reg, err := dhcp.IsRegisteredByIP(h.leases, ip)
	if err != nil {
		e.Log.WithFields(verbose.Fields{
			"error":   err,
			"package": "routes",
			"ip":      ip.String(),
		}).Error("Error getting registration status")
		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
		return
	}

	if reg {
		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
		return
	}

	http.Redirect(w, r, "/register", http.StatusTemporaryRedirect)
}

func notFoundHandler(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, "/api") {
		common.NewEmptyAPIResponse().WriteResponse(w, http.StatusNotFound)
		return
	}
	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}
<commit_msg>Remove /dev route<commit_after>// This source file is part of the Packet Guardian project.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package server

import (
	"fmt"
	"net/http"
	"net/http/pprof"
	"runtime"
	"strings"

	"github.com/dchest/captcha"
	assetfs "github.com/elazarl/go-bindata-assetfs"
	"github.com/gorilla/context"
	"github.com/julienschmidt/httprouter"
	"github.com/lfkeitel/verbose/v4"

	"github.com/packet-guardian/dhcp-lib"
	"github.com/packet-guardian/packet-guardian/src/auth"
	"github.com/packet-guardian/packet-guardian/src/bindata"
	"github.com/packet-guardian/packet-guardian/src/common"
	"github.com/packet-guardian/packet-guardian/src/controllers"
	"github.com/packet-guardian/packet-guardian/src/controllers/api"
	"github.com/packet-guardian/packet-guardian/src/models"
	"github.com/packet-guardian/packet-guardian/src/models/stores"
	mid "github.com/packet-guardian/packet-guardian/src/server/middleware"
)

func LoadRoutes(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	r.Handler("GET", "/", midStack(e, stores, &rootHandler{leases: stores.Leases}))
	r.ServeFiles(
		"/public/*filepath",
		&assetfs.AssetFS{
			Asset:     bindata.GetAsset,
			AssetDir:  bindata.GetAssetDir,
			AssetInfo: bindata.GetAssetInfo,
			Prefix:    "public"})

	authController := controllers.NewAuthController(e, stores.Users)
	r.Handler("GET", "/login", midStack(e, stores, http.HandlerFunc(authController.LoginHandler)))
	r.Handler("POST", "/login", midStack(e, stores, http.HandlerFunc(authController.LoginHandler)))
	r.Handler("GET", "/logout", midStack(e, stores, http.HandlerFunc(authController.LogoutHandler)))

	manageController := controllers.NewManagerController(e, stores.Devices, stores.Leases)
	r.Handler("GET", "/register", midStack(e, stores, http.HandlerFunc(manageController.RegistrationHandler)))
	r.Handler("GET", "/manage", midStack(e, stores, mid.CheckAuth(http.HandlerFunc(manageController.ManageHandler))))

	guestController := controllers.NewGuestController(e, stores.Users, stores.Devices, stores.Leases)
	r.Handler("GET", "/register/guest", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.RegistrationHandler), e, stores.Leases)))
	r.Handler("POST", "/register/guest", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.RegistrationHandler), e, stores.Leases)))
	r.Handler("GET", "/register/guest/verify", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.VerificationHandler), e, stores.Leases)))
	r.Handler("POST", "/register/guest/verify", midStack(e, stores, mid.CheckGuestReg(
		http.HandlerFunc(guestController.VerificationHandler), e, stores.Leases)))

	r.Handler("GET", "/admin/*a", midStack(e, stores, adminRouter(e, stores)))
	r.Handler("GET", "/api/*a", midStack(e, stores, apiRouter(e, stores)))
	r.Handler("POST", "/api/*a", midStack(e, stores, apiRouter(e, stores)))
	r.Handler("DELETE", "/api/*a", midStack(e, stores, apiRouter(e, stores)))

	r.Handler("GET", "/captcha/*a", captcha.Server(captcha.StdWidth, captcha.StdHeight))

	if e.IsDev() {
		r.Handler("GET", "/debug/*a", midStack(e, stores, debugRouter(e)))
		e.Log.Debug("Profiling enabled")
	}

	h := mid.Logging(r, e) // Logging
	h = mid.Panic(h, e)    // Panic catcher
	return h
}

func midStack(e *common.Environment, stores stores.StoreCollection, h http.Handler) http.Handler {
	h = mid.BlacklistCheck(h, e, stores.Devices, stores.Leases) // Enforce a blacklist check
	h = mid.Cache(h, e)                                         // Set cache headers if needed
	h = mid.SetSessionInfo(h, e, stores.Users)                  // Adds Environment and user information to requet context
	h = context.ClearHandler(h)                                 // Clear Gorilla sessions
	return h
}

func debugRouter(e *common.Environment) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	r.HandlerFunc("GET", "/debug/pprof", pprof.Index)
	r.HandlerFunc("GET", "/debug/pprof/cmdline", pprof.Cmdline)
	r.HandlerFunc("GET", "/debug/pprof/profile", pprof.Profile)
	r.HandlerFunc("GET", "/debug/pprof/symbol", pprof.Symbol)
	r.HandlerFunc("GET", "/debug/pprof/trace", pprof.Trace)
	// Manually add support for paths linked to by index page at /debug/pprof/
	r.Handler("GET", "/debug/pprof/goroutine", pprof.Handler("goroutine"))
	r.Handler("GET", "/debug/pprof/heap", pprof.Handler("heap"))
	r.Handler("GET", "/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
	r.Handler("GET", "/debug/pprof/block", pprof.Handler("block"))

	r.HandlerFunc("GET", "/debug/heap-stats", heapStats)

	h := mid.CheckAdmin(r)
	h = mid.CheckAuth(h)
	return h
}

func heapStats(w http.ResponseWriter, r *http.Request) {
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	fmt.Fprintf(w,
		"HeapSys: %d, HeapAlloc: %d, HeapIdle: %d, HeapReleased: %d\n",
		m.HeapSys,
		m.HeapAlloc,
		m.HeapIdle,
		m.HeapReleased,
	)
}

func adminRouter(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()
	r.NotFound = http.HandlerFunc(notFoundHandler)

	adminController := controllers.NewAdminController(e, stores)
	r.GET("/admin/", adminController.DashboardHandler)
	r.GET("/admin/search", adminController.SearchHandler)
	r.GET("/admin/manage/user/:username", adminController.ManageHandler)
	r.GET("/admin/manage/device/:mac", adminController.ShowDeviceHandler)
	r.GET("/admin/users", adminController.AdminUserListHandler)
	r.GET("/admin/users/:username", adminController.AdminUserHandler)
	r.GET("/admin/reports", adminController.ReportHandler)
	r.GET("/admin/reports/:report", adminController.ReportHandler)

	h := mid.CheckAdmin(r)
	h = mid.CheckAuth(h)
	return h
}

func apiRouter(e *common.Environment, stores stores.StoreCollection) http.Handler {
	r := httprouter.New()

	deviceAPIController := api.NewDeviceController(e, stores.Users, stores.Devices, stores.Leases)
	r.POST("/api/device", deviceAPIController.RegistrationHandler)
	r.DELETE("/api/device/user/:username", deviceAPIController.DeleteHandler)
	r.POST("/api/device/reassign", deviceAPIController.ReassignHandler)
	r.POST("/api/device/mac/:mac/description", deviceAPIController.EditDescriptionHandler)
	r.POST("/api/device/mac/:mac/expiration", deviceAPIController.EditExpirationHandler)
	r.POST("/api/device/mac/:mac/flag", deviceAPIController.EditFlaggedHandler)
	r.GET("/api/device/:mac", deviceAPIController.GetDeviceHandler)

	blacklistController := api.NewBlacklistController(e, stores.Users, stores.Devices)
	r.POST("/api/blacklist/user/:username", blacklistController.BlacklistUserHandler)
	r.DELETE("/api/blacklist/user/:username", blacklistController.BlacklistUserHandler)

	r.POST("/api/blacklist/device", blacklistController.BlacklistDeviceHandler)
	r.DELETE("/api/blacklist/device", blacklistController.BlacklistDeviceHandler)

	userAPIController := api.NewUserController(e, stores.Users, stores.Devices)
	r.POST("/api/user", userAPIController.UserHandler)
	r.DELETE("/api/user", userAPIController.UserHandler)
	r.GET("/api/user/:username", userAPIController.UserHandler)

	statusAPIController := api.NewStatusController(e)
	r.GET("/api/status", statusAPIController.GetStatus)

	return mid.CheckAuthAPI(r, stores.Users)
}

type rootHandler struct {
	leases stores.LeaseStore
}

func (h *rootHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if auth.IsLoggedIn(r) {
		sessionUser := models.GetUserFromContext(r)
		if sessionUser.Can(models.ViewAdminPage) {
			http.Redirect(w, r, "/admin", http.StatusTemporaryRedirect)
			return
		}

		http.Redirect(w, r, "/manage", http.StatusTemporaryRedirect)
		return
	}

	e := common.GetEnvironmentFromContext(r)
	ip := common.GetIPFromContext(r)
	reg, err := dhcp.IsRegisteredByIP(h.leases, ip)
	if err != nil {
		e.Log.WithFields(verbose.Fields{
			"error":   err,
			"package": "routes",
			"ip":      ip.String(),
		}).Error("Error getting registration status")
		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
		return
	}

	if reg {
		http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
		return
	}

	http.Redirect(w, r, "/register", http.StatusTemporaryRedirect)
}

func notFoundHandler(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, "/api") {
		common.NewEmptyAPIResponse().WriteResponse(w, http.StatusNotFound)
		return
	}
	http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
}
<|endoftext|>


<commit_before>// 4 october 2014
package main

import (
	"io"
	"encoding/binary"
	"errors"
	"strings"
	"strconv"
)

// TODO adorn io.ErrUnexpectedEOF?

type File struct {
	TailData	*TailData

	r		io.ReadSeeker
	version	string
	skip		int
	err		error
}

func Open(r io.ReadSeeker) *File {
	f := new(File)
	f.r = r
	f.readTailData()
	return f
}

// TODO rename to Error()? (see what stdlib does)
func (f *File) Err() error {
	return f.err
}

func (f *File) seek(offset uint32) {
	if f.err != nil {
		return
	}
	_, f.err = f.r.Seek(int64(offset), 0)
}

func (f *File) seekfull(offset int64, whence int) {
	if f.err != nil {
		return
	}
	_, f.err = f.r.Seek(offset, whence)
}

func (f *File) readu16() uint16 {
	var p [2]byte

	if f.err != nil {
		return 0
	}
	n, err := f.r.Read(p[:])
	if err != nil {
		f.err = err
		return 0
	} else if n != 2 {
		f.err = io.ErrUnexpectedEOF
		return 0
	}
	f.skip -= 2
	return binary.LittleEndian.Uint16(p[:])
}

func (f *File) readu32() uint32 {
	var p [4]byte

	if f.err != nil {
		return 0
	}
	n, err := f.r.Read(p[:])
	if err != nil {
		f.err = err
		return 0
	} else if n != 4 {
		f.err = io.ErrUnexpectedEOF
		return 0
	}
	f.skip -= 4
	return binary.LittleEndian.Uint32(p[:])
}

func (f *File) readSkip() {
	if f.err != nil {
		return
	}
	f.skip = 0
	if f.versionGreaterEqual("1.1.0.0") {
		f.skip = int(f.read7BitEncodedInt())
	}
}

func (f *File) doSkip() {
	if f.err != nil {
		return
	}
	if f.skip > 0 {
		_, f.err = f.r.Seek(int64(f.skip), 1)
	}
	f.skip = 0
}

func (f *File) readCLS() (cls CLS) {
	cls.Count = f.readu32()
	cls.Length = f.readu32()
	cls.StartPos = f.readu32()
	return cls
}

func (f *File) readODL() (odl ODL) {
	odl.Offset = f.readCLS()
	odl.Data = f.readCLS()
	odl.Location = f.readCLS()
	return odl
}

func (f *File) readString() string {
	if f.err != nil {
		return ""
	}
	length := f.read7BitEncodedInt()
	if f.err != nil {
		return ""
	}
	p := make([]byte, length)
	n, err := f.r.Read(p)
	if err != nil {
		f.err = err
		return ""
	} else if n != int(length) {
		f.err = io.ErrUnexpectedEOF
		return ""
	}
	f.skip -= n
	// don't worry, the string is supposed to be UTF-8
	return string(p)
}

func (f *File) read7BitEncodedInt() uint32 {
	var n [1]byte

	if f.err != nil {
		return 0
	}
	out := uint32(0)
	shift := uint32(0)
	for {
		q, err := f.r.Read(n[:])
		if err != nil {
			f.err = err
			return 0
		} else if q != 1 {
			f.err = io.ErrUnexpectedEOF
			return 0
		}
		f.skip--
		out |= uint32(n[0] & 0x7F) << shift
		shift += 7
		if n[0] & 0x80 == 0 {
			break
		}
	}
	return out
}

func (f *File) versionGreaterEqual(against string) bool {
	var err error

	if f.err != nil {
		return false
	}
	vparts := strings.Split(f.version, ".")
	aparts := strings.Split(against, ".")
	if len(vparts) != len(aparts) {
		f.err = errors.New("vparts and aparts have different lengths in File.versionGreaterEqual()")
		return false
	}
	vints := make([]int, len(vparts))
	for i, v := range vparts {
		vints[i], err = strconv.Atoi(v)
		if err != nil {
			// TODO adorn error message
			f.err = err
			return false
		}
	}
	aints := make([]int, len(aparts))
	for i, a := range vparts {
		aints[i], err = strconv.Atoi(a)
		if err != nil {
			// TODO adorn error message
			f.err = err
			return false
		}
	}
	for i := 0; i < len(vints); i++ {
		if vints[i] > aints[i] {
			// versions are greater
			return true
		} else if vints[i] == aints[i] {
			// parts are the same so far...
			continue
		} else {
			// versions are less
			return false
		}
	}
	// versions match
	return true
}

// TODO adorn error messages?
func (f *File) readOffsetArray(offset CLS, data CLS) []uint32 {
	if f.err != nil {
		return nil
	}
	if data.Count == 0 {
		return nil
	}
	n := data.Count / f.TailData.OffsetInterval
	if data.Count % f.TailData.OffsetInterval != 0 {
		n++
	}
	list := make([]uint32, n)
	f.seek(offset.StartPos)
	postv1 := f.versionGreaterEqual("1.1.0.0")
	for i := uint32(0); i < n; i++ {
		if postv1 {
			f.readSkip()
			list[i] = f.read7BitEncodedInt()
			f.doSkip()
		} else {
			list[i] = f.readu32()
		}
	}
	if f.err != nil {
		return nil
	}
	return list
}

func (f *File) realOffset(offsets []uint32, cls CLS) uint32 {
	if f.err != nil {
		return 0
	}
	if !f.versionGreaterEqual("1.1.0.0") {
		return offsets[0]
	}
	return offsets[0] + cls.StartPos
}
<commit_msg>Fixed a small bug in Open(). New setup works!<commit_after>// 4 october 2014
package main

import (
	"io"
	"encoding/binary"
	"errors"
	"strings"
	"strconv"
)

// TODO adorn io.ErrUnexpectedEOF?

type File struct {
	TailData	*TailData

	r		io.ReadSeeker
	version	string
	skip		int
	err		error
}

func Open(r io.ReadSeeker) *File {
	f := new(File)
	f.r = r
	// TODO avoid need for assignment
	f.TailData = f.readTailData()
	return f
}

// TODO rename to Error()? (see what stdlib does)
func (f *File) Err() error {
	return f.err
}

func (f *File) seek(offset uint32) {
	if f.err != nil {
		return
	}
	_, f.err = f.r.Seek(int64(offset), 0)
}

func (f *File) seekfull(offset int64, whence int) {
	if f.err != nil {
		return
	}
	_, f.err = f.r.Seek(offset, whence)
}

func (f *File) readu16() uint16 {
	var p [2]byte

	if f.err != nil {
		return 0
	}
	n, err := f.r.Read(p[:])
	if err != nil {
		f.err = err
		return 0
	} else if n != 2 {
		f.err = io.ErrUnexpectedEOF
		return 0
	}
	f.skip -= 2
	return binary.LittleEndian.Uint16(p[:])
}

func (f *File) readu32() uint32 {
	var p [4]byte

	if f.err != nil {
		return 0
	}
	n, err := f.r.Read(p[:])
	if err != nil {
		f.err = err
		return 0
	} else if n != 4 {
		f.err = io.ErrUnexpectedEOF
		return 0
	}
	f.skip -= 4
	return binary.LittleEndian.Uint32(p[:])
}

func (f *File) readSkip() {
	if f.err != nil {
		return
	}
	f.skip = 0
	if f.versionGreaterEqual("1.1.0.0") {
		f.skip = int(f.read7BitEncodedInt())
	}
}

func (f *File) doSkip() {
	if f.err != nil {
		return
	}
	if f.skip > 0 {
		_, f.err = f.r.Seek(int64(f.skip), 1)
	}
	f.skip = 0
}

func (f *File) readCLS() (cls CLS) {
	cls.Count = f.readu32()
	cls.Length = f.readu32()
	cls.StartPos = f.readu32()
	return cls
}

func (f *File) readODL() (odl ODL) {
	odl.Offset = f.readCLS()
	odl.Data = f.readCLS()
	odl.Location = f.readCLS()
	return odl
}

func (f *File) readString() string {
	if f.err != nil {
		return ""
	}
	length := f.read7BitEncodedInt()
	if f.err != nil {
		return ""
	}
	p := make([]byte, length)
	n, err := f.r.Read(p)
	if err != nil {
		f.err = err
		return ""
	} else if n != int(length) {
		f.err = io.ErrUnexpectedEOF
		return ""
	}
	f.skip -= n
	// don't worry, the string is supposed to be UTF-8
	return string(p)
}

func (f *File) read7BitEncodedInt() uint32 {
	var n [1]byte

	if f.err != nil {
		return 0
	}
	out := uint32(0)
	shift := uint32(0)
	for {
		q, err := f.r.Read(n[:])
		if err != nil {
			f.err = err
			return 0
		} else if q != 1 {
			f.err = io.ErrUnexpectedEOF
			return 0
		}
		f.skip--
		out |= uint32(n[0] & 0x7F) << shift
		shift += 7
		if n[0] & 0x80 == 0 {
			break
		}
	}
	return out
}

func (f *File) versionGreaterEqual(against string) bool {
	var err error

	if f.err != nil {
		return false
	}
	vparts := strings.Split(f.version, ".")
	aparts := strings.Split(against, ".")
	if len(vparts) != len(aparts) {
		f.err = errors.New("vparts and aparts have different lengths in File.versionGreaterEqual()")
		return false
	}
	vints := make([]int, len(vparts))
	for i, v := range vparts {
		vints[i], err = strconv.Atoi(v)
		if err != nil {
			// TODO adorn error message
			f.err = err
			return false
		}
	}
	aints := make([]int, len(aparts))
	for i, a := range vparts {
		aints[i], err = strconv.Atoi(a)
		if err != nil {
			// TODO adorn error message
			f.err = err
			return false
		}
	}
	for i := 0; i < len(vints); i++ {
		if vints[i] > aints[i] {
			// versions are greater
			return true
		} else if vints[i] == aints[i] {
			// parts are the same so far...
			continue
		} else {
			// versions are less
			return false
		}
	}
	// versions match
	return true
}

// TODO adorn error messages?
func (f *File) readOffsetArray(offset CLS, data CLS) []uint32 {
	if f.err != nil {
		return nil
	}
	if data.Count == 0 {
		return nil
	}
	n := data.Count / f.TailData.OffsetInterval
	if data.Count % f.TailData.OffsetInterval != 0 {
		n++
	}
	list := make([]uint32, n)
	f.seek(offset.StartPos)
	postv1 := f.versionGreaterEqual("1.1.0.0")
	for i := uint32(0); i < n; i++ {
		if postv1 {
			f.readSkip()
			list[i] = f.read7BitEncodedInt()
			f.doSkip()
		} else {
			list[i] = f.readu32()
		}
	}
	if f.err != nil {
		return nil
	}
	return list
}

func (f *File) realOffset(offsets []uint32, cls CLS) uint32 {
	if f.err != nil {
		return 0
	}
	if !f.versionGreaterEqual("1.1.0.0") {
		return offsets[0]
	}
	return offsets[0] + cls.StartPos
}
<|endoftext|>


<commit_before>// The plugin package provides the functionality to both expose a Packer
// plugin binary and to connect to an existing Packer plugin binary.
//
// Packer supports plugins in the form of self-contained external static
// Go binaries. These binaries behave in a certain way (enforced by this
// package) and are connected to in a certain way (also enforced by this
// package).
package plugin

import (
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	packrpc "github.com/mitchellh/packer/packer/rpc"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
)

const MagicCookieKey = "PACKER_PLUGIN_MAGIC_COOKIE"
const MagicCookieValue = "d602bf8f470bc67ca7faa0386276bbdd4330efaf76d1a219cb4d6991ca9872b2"

// This serves a single RPC connection on the given RPC server on
// a random port.
func serve(server *rpc.Server) (err error) {
	if os.Getenv(MagicCookieKey) != MagicCookieValue {
		return errors.New("Please do not execute plugins directly. Packer will execute these for you.")
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	minPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MIN_PORT"), 10, 32)
	if err != nil {
		return
	}

	maxPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MAX_PORT"), 10, 32)
	if err != nil {
		return
	}

	log.Printf("Plugin minimum port: %d\n", minPort)
	log.Printf("Plugin maximum port: %d\n", maxPort)

	// Set the RPC port range
	packrpc.PortRange(int(minPort), int(maxPort))

	var address string
	var listener net.Listener
	for port := minPort; port <= maxPort; port++ {
		address = fmt.Sprintf(":%d", port)
		listener, err = net.Listen("tcp", address)
		if err != nil {
			if !strings.Contains(err.Error(), "address already in use") {
				// Not an address already in use error, return.
				return
			} else {
				// Address is in use, just try another
				err = nil
				continue
			}
		}

		break
	}

	defer listener.Close()

	// Output the address to stdout
	log.Printf("Plugin address: %s\n", address)
	fmt.Println(address)
	os.Stdout.Sync()

	// Accept a connection
	log.Println("Waiting for connection...")
	conn, err := listener.Accept()
	if err != nil {
		log.Printf("Error accepting connection: %s\n", err.Error())
		return
	}

	// Serve a single connection
	log.Println("Serving a plugin connection...")
	server.ServeConn(conn)
	return
}

// Registers a signal handler to "swallow" interrupts so that the
// plugin isn't killed. The main host Packer process is responsible
// for killing the plugins when interrupted.
func swallowInterrupts() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)

	go func() {
		<-ch
		log.Println("Received interrupt signal. Ignoring.")
	}()
}

// Serves a builder from a plugin.
func ServeBuilder(builder packer.Builder) {
	log.Println("Preparing to serve a builder plugin...")

	server := rpc.NewServer()
	packrpc.RegisterBuilder(server, builder)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Panic(err)
	}
}

// Serves a command from a plugin.
func ServeCommand(command packer.Command) {
	log.Println("Preparing to serve a command plugin...")

	server := rpc.NewServer()
	packrpc.RegisterCommand(server, command)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Panic(err)
	}
}

// Serves a hook from a plugin.
func ServeHook(hook packer.Hook) {
	log.Println("Preparing to serve a hook plugin...")

	server := rpc.NewServer()
	packrpc.RegisterHook(server, hook)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Panic(err)
	}
}

// Serves a post-processor from a plugin.
func ServePostProcessor(p packer.PostProcessor) {
	log.Println("Preparing to serve a post-processor plugin...")

	server := rpc.NewServer()
	packrpc.RegisterPostProcessor(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Panic(err)
	}
}

// Serves a provisioner from a plugin.
func ServeProvisioner(p packer.Provisioner) {
	log.Println("Preparing to serve a provisioner plugin...")

	server := rpc.NewServer()
	packrpc.RegisterProvisioner(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Panic(err)
	}
}
<commit_msg>packer/plugin: No need for a panic when plugins error<commit_after>// The plugin package provides the functionality to both expose a Packer
// plugin binary and to connect to an existing Packer plugin binary.
//
// Packer supports plugins in the form of self-contained external static
// Go binaries. These binaries behave in a certain way (enforced by this
// package) and are connected to in a certain way (also enforced by this
// package).
package plugin

import (
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	packrpc "github.com/mitchellh/packer/packer/rpc"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
)

const MagicCookieKey = "PACKER_PLUGIN_MAGIC_COOKIE"
const MagicCookieValue = "d602bf8f470bc67ca7faa0386276bbdd4330efaf76d1a219cb4d6991ca9872b2"

// This serves a single RPC connection on the given RPC server on
// a random port.
func serve(server *rpc.Server) (err error) {
	if os.Getenv(MagicCookieKey) != MagicCookieValue {
		return errors.New("Please do not execute plugins directly. Packer will execute these for you.")
	}

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	minPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MIN_PORT"), 10, 32)
	if err != nil {
		return
	}

	maxPort, err := strconv.ParseInt(os.Getenv("PACKER_PLUGIN_MAX_PORT"), 10, 32)
	if err != nil {
		return
	}

	log.Printf("Plugin minimum port: %d\n", minPort)
	log.Printf("Plugin maximum port: %d\n", maxPort)

	// Set the RPC port range
	packrpc.PortRange(int(minPort), int(maxPort))

	var address string
	var listener net.Listener
	for port := minPort; port <= maxPort; port++ {
		address = fmt.Sprintf(":%d", port)
		listener, err = net.Listen("tcp", address)
		if err != nil {
			if !strings.Contains(err.Error(), "address already in use") {
				// Not an address already in use error, return.
				return
			} else {
				// Address is in use, just try another
				err = nil
				continue
			}
		}

		break
	}

	defer listener.Close()

	// Output the address to stdout
	log.Printf("Plugin address: %s\n", address)
	fmt.Println(address)
	os.Stdout.Sync()

	// Accept a connection
	log.Println("Waiting for connection...")
	conn, err := listener.Accept()
	if err != nil {
		log.Printf("Error accepting connection: %s\n", err.Error())
		return
	}

	// Serve a single connection
	log.Println("Serving a plugin connection...")
	server.ServeConn(conn)
	return
}

// Registers a signal handler to "swallow" interrupts so that the
// plugin isn't killed. The main host Packer process is responsible
// for killing the plugins when interrupted.
func swallowInterrupts() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)

	go func() {
		<-ch
		log.Println("Received interrupt signal. Ignoring.")
	}()
}

// Serves a builder from a plugin.
func ServeBuilder(builder packer.Builder) {
	log.Println("Preparing to serve a builder plugin...")

	server := rpc.NewServer()
	packrpc.RegisterBuilder(server, builder)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a command from a plugin.
func ServeCommand(command packer.Command) {
	log.Println("Preparing to serve a command plugin...")

	server := rpc.NewServer()
	packrpc.RegisterCommand(server, command)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a hook from a plugin.
func ServeHook(hook packer.Hook) {
	log.Println("Preparing to serve a hook plugin...")

	server := rpc.NewServer()
	packrpc.RegisterHook(server, hook)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a post-processor from a plugin.
func ServePostProcessor(p packer.PostProcessor) {
	log.Println("Preparing to serve a post-processor plugin...")

	server := rpc.NewServer()
	packrpc.RegisterPostProcessor(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}

// Serves a provisioner from a plugin.
func ServeProvisioner(p packer.Provisioner) {
	log.Println("Preparing to serve a provisioner plugin...")

	server := rpc.NewServer()
	packrpc.RegisterProvisioner(server, p)

	swallowInterrupts()
	if err := serve(server); err != nil {
		log.Printf("ERROR: %s", err)
		os.Exit(1)
	}
}
<|endoftext|>