<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		spaceName := string(space.ProviderId)
		spaceName = strings.Replace(spaceName, " ", "-", -1)
		spaceName = strings.ToLower(spaceName)
		if !names.IsValidSpace(spaceName) {
			// XXX generate a valid name here
			logger.Errorf("invalid space name %v", spaceName)
			return errors.Errorf("invalid space name: %q", spaceName)
		}
		spaceTag := names.NewSpaceTag(spaceName)
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			// XXX skip spaces with no subnets(?)
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				logger.Errorf("invalid creating space %v", err)
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		for _, subnet := range space.Subnets {
			if !names.IsValidSubnet(subnet.CIDR) {
				// XXX generate a valid name here
				logger.Errorf("invalid  subnet %v", subnet.CIDR)
				return errors.Errorf("invalid subnet: %q", subnet.CIDR)
			}
			subnetTag := names.NewSubnetTag(subnet.CIDR)
			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetTag:        subnetTag.String(),
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			// XXX check the error result too.
			_, err = dw.api.AddSubnets(args)
			if err != nil {
				logger.Errorf("invalid creating subnet %v", err)
				return errors.Trace(err)
			}
		}
	}
	return nil
}
<commit_msg>Error handling<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package discoverspaces

import (
	"strings"

	"github.com/juju/errors"
	"github.com/juju/juju/api/discoverspaces"
	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/environs"
	"github.com/juju/juju/worker"
	"github.com/juju/loggo"
	"github.com/juju/names"
	"github.com/juju/utils/set"
	"launchpad.net/tomb"
)

var logger = loggo.GetLogger("juju.discoverspaces")

type discoverspacesWorker struct {
	api      *discoverspaces.API
	tomb     tomb.Tomb
	observer *worker.EnvironObserver
}

// NewWorker returns a worker
func NewWorker(api *discoverspaces.API) worker.Worker {
	dw := &discoverspacesWorker{
		api: api,
	}
	go func() {
		defer dw.tomb.Done()
		dw.tomb.Kill(dw.loop())
	}()
	return dw
}

func (dw *discoverspacesWorker) Kill() {
	dw.tomb.Kill(nil)
}

func (dw *discoverspacesWorker) Wait() error {
	return dw.tomb.Wait()
}

func (dw *discoverspacesWorker) loop() (err error) {
	dw.observer, err = worker.NewEnvironObserver(dw.api)
	if err != nil {
		return err
	}
	defer func() {
		obsErr := worker.Stop(dw.observer)
		if err == nil {
			err = obsErr
		}
	}()
	environ := dw.observer.Environ()
	networkingEnviron, ok := environs.SupportsNetworking(environ)

	if ok {
		err = dw.handleSubnets(networkingEnviron)
		if err != nil {
			return errors.Trace(err)
		}
	}

	// TODO(mfoord): we'll have a watcher here checking if we need to
	// update the spaces/subnets definition.
	dying := dw.tomb.Dying()
	for {
		select {
		case <-dying:
			return nil
		}
	}
	return nil
}

func (dw *discoverspacesWorker) handleSubnets(env environs.NetworkingEnviron) error {
	ok, err := env.SupportsSpaceDiscovery()
	if err != nil {
		return errors.Trace(err)
	}
	if !ok {
		// Nothing to do.
		return nil
	}
	providerSpaces, err := env.Spaces()
	if err != nil {
		return errors.Trace(err)
	}
	listSpacesResult, err := dw.api.ListSpaces()
	if err != nil {
		return errors.Trace(err)
	}
	stateSpaceMap := make(map[string]params.ProviderSpace)
	spaceNames := make(set.Strings)
	for _, space := range listSpacesResult.Results {
		stateSpaceMap[space.ProviderId] = space
		spaceNames.Add(space.Name)
	}

	// TODO(mfoord): we need to delete spaces and subnets that no longer
	// exist, so long as they're not in use.
	for _, space := range providerSpaces {
		spaceName := string(space.ProviderId)
		spaceName = strings.Replace(spaceName, " ", "-", -1)
		spaceName = strings.ToLower(spaceName)
		if !names.IsValidSpace(spaceName) {
			// XXX generate a valid name here
			logger.Errorf("invalid space name %v", spaceName)
			return errors.Errorf("invalid space name: %q", spaceName)
		}
		spaceTag := names.NewSpaceTag(spaceName)
		_, ok := stateSpaceMap[string(space.ProviderId)]
		if !ok {
			// XXX skip spaces with no subnets(?)
			// We need to create the space.
			args := params.CreateSpacesParams{
				Spaces: []params.CreateSpaceParams{{
					Public:   false,
					SpaceTag: spaceTag.String(),
				}}}
			// XXX check the error result too.
			_, err = dw.api.CreateSpaces(args)
			if err != nil {
				logger.Errorf("invalid creating space %v", err)
				return errors.Trace(err)
			}
		}
		// TODO(mfoord): currently no way of removing subnets, or
		// changing the space they're in, so we can only add ones we
		// don't already know about.
		logger.Errorf("Created space %v with %v subnets", spaceName, len(space.Subnets))
		for _, subnet := range space.Subnets {
			args := params.AddSubnetsParams{
				Subnets: []params.AddSubnetParams{{
					SubnetProviderId: string(subnet.ProviderId),
					SpaceTag:         spaceTag.String(),
					Zones:            subnet.AvailabilityZones,
				}}}
			// XXX check the error result too.
			logger.Errorf("Adding subnet %v", subnet.CIDR)
			result, err := dw.api.AddSubnets(args)
			if err != nil {
				logger.Errorf("invalid creating subnet %v", err)
				return errors.Trace(err)
			}
			if result.Results[0].Error != nil {
				logger.Errorf("error creating subnet %v", result.Results[0].Error)
				return errors.Errorf("error creating subnet %v", result.Results[0].Error)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/serf/coordinate"
)

func TestCoordinate_Disabled_Response(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), `
		disable_coordinates = true
`)
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	tests := []func(resp http.ResponseWriter, req *http.Request) (interface{}, error){
		a.srv.CoordinateDatacenters,
		a.srv.CoordinateNodes,
		a.srv.CoordinateNode,
		a.srv.CoordinateUpdate,
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
			req, _ := http.NewRequest("PUT", "/should/not/care", nil)
			resp := httptest.NewRecorder()
			obj, err := tt(resp, req)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if obj != nil {
				t.Fatalf("bad: %#v", obj)
			}
			if got, want := resp.Code, http.StatusUnauthorized; got != want {
				t.Fatalf("got %d want %d", got, want)
			}
			if !strings.Contains(resp.Body.String(), "Coordinate support disabled") {
				t.Fatalf("bad: %#v", resp)
			}
		})
	}
}

func TestCoordinate_Datacenters(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	req, _ := http.NewRequest("GET", "/v1/coordinate/datacenters", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateDatacenters(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	maps := obj.([]structs.DatacenterMap)
	if len(maps) != 1 ||
		maps[0].Datacenter != "dc1" ||
		len(maps[0].Coordinates) != 1 ||
		maps[0].Coordinates[0].Node != a.Config.NodeName {
		t.Fatalf("bad: %v", maps)
	}
}

func TestCoordinate_Nodes(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure an empty list is non-nil.
	req, _ := http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates := obj.(structs.Coordinates)
	if coordinates == nil || len(coordinates) != 0 {
		t.Fatalf("bad: %v", coordinates)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present and sorted correctly.
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 2 ||
		coordinates[0].Node != "bar" ||
		coordinates[1].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 0 {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "bar" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Node(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure we get a 404 with no coordinates.
	req, _ := http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present.
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates := obj.(structs.Coordinates)
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}
}

func TestCoordinate_Update(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Register the node.
	reg := structs.RegisterRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var reply struct{}
	if err := a.RPC("Catalog.Register", &reg, &reply); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Update the coordinates and wait for it to complete.
	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}
	req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
	resp := httptest.NewRecorder()
	_, err := a.srv.CoordinateUpdate(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the coordinates are present.
	args := structs.NodeSpecificRequest{Node: "foo", Datacenter: "dc1"}
	var coords structs.IndexedCoordinates
	if err := a.RPC("Coordinate.Node", &args, &coords); err != nil {
		t.Fatalf("err: %s", err)
	}

	coordinates := coords.Coordinates
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Update_ACLDeny(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), TestACLConfig())
	defer a.Shutdown()
	testrpc.WaitForLeader(t, a.RPC, "dc1")

	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}

	t.Run("no token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); !acl.IsErrPermissionDenied(err) {
			t.Fatalf("err: %v", err)
		}
	})

	t.Run("valid token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update?token=root", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); err != nil {
			t.Fatalf("err: %v", err)
		}
	})
}
<commit_msg>Update type assertion logic<commit_after>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/serf/coordinate"
)

func TestCoordinate_Disabled_Response(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), `
		disable_coordinates = true
`)
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	tests := []func(resp http.ResponseWriter, req *http.Request) (interface{}, error){
		a.srv.CoordinateDatacenters,
		a.srv.CoordinateNodes,
		a.srv.CoordinateNode,
		a.srv.CoordinateUpdate,
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
			req, _ := http.NewRequest("PUT", "/should/not/care", nil)
			resp := httptest.NewRecorder()
			obj, err := tt(resp, req)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if obj != nil {
				t.Fatalf("bad: %#v", obj)
			}
			if got, want := resp.Code, http.StatusUnauthorized; got != want {
				t.Fatalf("got %d want %d", got, want)
			}
			if !strings.Contains(resp.Body.String(), "Coordinate support disabled") {
				t.Fatalf("bad: %#v", resp)
			}
		})
	}
}

func TestCoordinate_Datacenters(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	req, _ := http.NewRequest("GET", "/v1/coordinate/datacenters", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateDatacenters(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	maps := obj.([]structs.DatacenterMap)
	if len(maps) != 1 ||
		maps[0].Datacenter != "dc1" ||
		len(maps[0].Coordinates) != 1 ||
		maps[0].Coordinates[0].Node != a.Config.NodeName {
		t.Fatalf("bad: %v", maps)
	}
}

func TestCoordinate_Nodes(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure an empty list is non-nil.
	req, _ := http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that coordinates are empty before registering a node
	coordinates, ok := obj.(structs.Coordinates)
	if !ok {
		t.Fatalf("expected: structs.Coordinates, received: %v", coordinates)
	}

	if len(coordinates) != 0 {
		t.Fatalf("coordinates should be empty, received: %v", coordinates)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present and sorted correctly.
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 2 ||
		coordinates[0].Node != "bar" ||
		coordinates[1].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 0 {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "bar" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Node(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure we get a 404 with no coordinates.
	req, _ := http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present.
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates := obj.(structs.Coordinates)
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}
}

func TestCoordinate_Update(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Register the node.
	reg := structs.RegisterRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var reply struct{}
	if err := a.RPC("Catalog.Register", &reg, &reply); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Update the coordinates and wait for it to complete.
	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}
	req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
	resp := httptest.NewRecorder()
	_, err := a.srv.CoordinateUpdate(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the coordinates are present.
	args := structs.NodeSpecificRequest{Node: "foo", Datacenter: "dc1"}
	var coords structs.IndexedCoordinates
	if err := a.RPC("Coordinate.Node", &args, &coords); err != nil {
		t.Fatalf("err: %s", err)
	}

	coordinates := coords.Coordinates
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Update_ACLDeny(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), TestACLConfig())
	defer a.Shutdown()
	testrpc.WaitForLeader(t, a.RPC, "dc1")

	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}

	t.Run("no token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); !acl.IsErrPermissionDenied(err) {
			t.Fatalf("err: %v", err)
		}
	})

	t.Run("valid token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update?token=root", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); err != nil {
			t.Fatalf("err: %v", err)
		}
	})
}
<|endoftext|>


<commit_before>package log

import (
	"io"
	golog "log"
	"os"
)

const (
	ERROR = iota
	INFO
	WARN
	DEBUG
)

type Logger struct {
	golog.Logger
	Level  int
	Prefix string
}

func New(out io.Writer, prefix string, level int) *Logger {
	l := &Logger{
		Level:  level,
		Prefix: prefix,
	}
	l.Logger = *(golog.New(out, prefix, golog.LstdFlags))
	return l
}

var DefaultLogger = New(os.Stderr, "", INFO)

func (l *Logger) Debug(v ...interface{}) {
	if l.Level < DEBUG {
		return
	}
	l.Println(v...)
}

func (l *Logger) Debugf(fmt string, v ...interface{}) {
	if l.Level < DEBUG {
		return
	}
	l.Printf(fmt, v...)
}

func (l *Logger) Write(p []byte) (n int, err error) {
	l.Print(string(p))
	return len(p), nil
}

func Debug(v ...interface{})                 { DefaultLogger.Debug(v...) }
func Debugf(format string, v ...interface{}) { DefaultLogger.Debugf(format, v...) }
func Fatal(v ...interface{})                 { DefaultLogger.Fatal(v...) }
func Fatalf(format string, v ...interface{}) { DefaultLogger.Fatalf(format, v...) }
func Fatalln(v ...interface{})               { DefaultLogger.Fatalln(v...) }
func Panic(v ...interface{})                 { DefaultLogger.Panic(v...) }
func Panicf(format string, v ...interface{}) { DefaultLogger.Panicf(format, v...) }
func Panicln(v ...interface{})               { DefaultLogger.Panicln(v...) }
func Print(v ...interface{})                 { DefaultLogger.Print(v...) }
func Printf(format string, v ...interface{}) { DefaultLogger.Printf(format, v...) }
func Println(v ...interface{})               { DefaultLogger.Println(v...) }
<commit_msg>Only print PullImage progress info if -debug is enabled<commit_after>package log

import (
	"io"
	golog "log"
	"os"
)

const (
	ERROR = iota
	INFO
	WARN
	DEBUG
)

type Logger struct {
	golog.Logger
	Level  int
	Prefix string
}

func New(out io.Writer, prefix string, level int) *Logger {
	l := &Logger{
		Level:  level,
		Prefix: prefix,
	}
	l.Logger = *(golog.New(out, prefix, golog.LstdFlags))
	return l
}

var DefaultLogger = New(os.Stderr, "", INFO)

func (l *Logger) Debug(v ...interface{}) {
	if l.Level < DEBUG {
		return
	}
	l.Println(v...)
}

func (l *Logger) Debugf(fmt string, v ...interface{}) {
	if l.Level < DEBUG {
		return
	}
	l.Printf(fmt, v...)
}

func (l *Logger) Write(p []byte) (n int, err error) {
	if l.Level < DEBUG {
		return
	}
	l.Print(string(p))
	return len(p), nil
}

func Debug(v ...interface{})                 { DefaultLogger.Debug(v...) }
func Debugf(format string, v ...interface{}) { DefaultLogger.Debugf(format, v...) }
func Fatal(v ...interface{})                 { DefaultLogger.Fatal(v...) }
func Fatalf(format string, v ...interface{}) { DefaultLogger.Fatalf(format, v...) }
func Fatalln(v ...interface{})               { DefaultLogger.Fatalln(v...) }
func Panic(v ...interface{})                 { DefaultLogger.Panic(v...) }
func Panicf(format string, v ...interface{}) { DefaultLogger.Panicf(format, v...) }
func Panicln(v ...interface{})               { DefaultLogger.Panicln(v...) }
func Print(v ...interface{})                 { DefaultLogger.Print(v...) }
func Printf(format string, v ...interface{}) { DefaultLogger.Printf(format, v...) }
func Println(v ...interface{})               { DefaultLogger.Println(v...) }
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/serf/coordinate"
)

func TestCoordinate_Disabled_Response(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), `
		disable_coordinates = true
`)
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	tests := []func(resp http.ResponseWriter, req *http.Request) (interface{}, error){
		a.srv.CoordinateDatacenters,
		a.srv.CoordinateNodes,
		a.srv.CoordinateNode,
		a.srv.CoordinateUpdate,
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
			req, _ := http.NewRequest("PUT", "/should/not/care", nil)
			resp := httptest.NewRecorder()
			obj, err := tt(resp, req)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if obj != nil {
				t.Fatalf("bad: %#v", obj)
			}
			if got, want := resp.Code, http.StatusUnauthorized; got != want {
				t.Fatalf("got %d want %d", got, want)
			}
			if !strings.Contains(resp.Body.String(), "Coordinate support disabled") {
				t.Fatalf("bad: %#v", resp)
			}
		})
	}
}

func TestCoordinate_Datacenters(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	req, _ := http.NewRequest("GET", "/v1/coordinate/datacenters", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateDatacenters(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	maps := obj.([]structs.DatacenterMap)
	if len(maps) != 1 ||
		maps[0].Datacenter != "dc1" ||
		len(maps[0].Coordinates) != 1 ||
		maps[0].Coordinates[0].Node != a.Config.NodeName {
		t.Fatalf("bad: %v", maps)
	}
}

func TestCoordinate_Nodes(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure an empty list is non-nil.
	req, _ := http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that coordinates are empty before registering a node
	coordinates, ok := obj.(structs.Coordinates)
	if !ok {
		t.Fatalf("expected: structs.Coordinates, received: %v", coordinates)
	}

	if len(coordinates) != 0 {
		t.Fatalf("coordinates should be empty, received: %v", coordinates)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present and sorted correctly.
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 2 ||
		coordinates[0].Node != "bar" ||
		coordinates[1].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 0 {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "bar" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Node(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure we get a 404 with no coordinates.
	req, _ := http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present.
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates := obj.(structs.Coordinates)
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}
}

func TestCoordinate_Update(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Register the node.
	reg := structs.RegisterRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var reply struct{}
	if err := a.RPC("Catalog.Register", &reg, &reply); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Update the coordinates and wait for it to complete.
	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}
	req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
	resp := httptest.NewRecorder()
	_, err := a.srv.CoordinateUpdate(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the coordinates are present.
	args := structs.NodeSpecificRequest{Node: "foo", Datacenter: "dc1"}
	var coords structs.IndexedCoordinates
	if err := a.RPC("Coordinate.Node", &args, &coords); err != nil {
		t.Fatalf("err: %s", err)
	}

	coordinates := coords.Coordinates
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Update_ACLDeny(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), TestACLConfig())
	defer a.Shutdown()
	testrpc.WaitForLeader(t, a.RPC, "dc1")

	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}

	t.Run("no token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); !acl.IsErrPermissionDenied(err) {
			t.Fatalf("err: %v", err)
		}
	})

	t.Run("valid token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update?token=root", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); err != nil {
			t.Fatalf("err: %v", err)
		}
	})
}
<commit_msg>Update agent/coordinate_endpoint_test.go<commit_after>package agent

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/structs"
	"github.com/hashicorp/consul/testrpc"
	"github.com/hashicorp/serf/coordinate"
)

func TestCoordinate_Disabled_Response(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), `
		disable_coordinates = true
`)
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	tests := []func(resp http.ResponseWriter, req *http.Request) (interface{}, error){
		a.srv.CoordinateDatacenters,
		a.srv.CoordinateNodes,
		a.srv.CoordinateNode,
		a.srv.CoordinateUpdate,
	}
	for i, tt := range tests {
		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
			req, _ := http.NewRequest("PUT", "/should/not/care", nil)
			resp := httptest.NewRecorder()
			obj, err := tt(resp, req)
			if err != nil {
				t.Fatalf("err: %v", err)
			}
			if obj != nil {
				t.Fatalf("bad: %#v", obj)
			}
			if got, want := resp.Code, http.StatusUnauthorized; got != want {
				t.Fatalf("got %d want %d", got, want)
			}
			if !strings.Contains(resp.Body.String(), "Coordinate support disabled") {
				t.Fatalf("bad: %#v", resp)
			}
		})
	}
}

func TestCoordinate_Datacenters(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	req, _ := http.NewRequest("GET", "/v1/coordinate/datacenters", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateDatacenters(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	maps := obj.([]structs.DatacenterMap)
	if len(maps) != 1 ||
		maps[0].Datacenter != "dc1" ||
		len(maps[0].Coordinates) != 1 ||
		maps[0].Coordinates[0].Node != a.Config.NodeName {
		t.Fatalf("bad: %v", maps)
	}
}

func TestCoordinate_Nodes(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure an empty list is non-nil.
	req, _ := http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Check that coordinates are empty before registering a node
	coordinates, ok := obj.(structs.Coordinates)
	if !ok {
		t.Fatalf("expected: structs.Coordinates, received: %+v", obj)
	}

	if len(coordinates) != 0 {
		t.Fatalf("coordinates should be empty, received: %v", coordinates)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present and sorted correctly.
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 2 ||
		coordinates[0].Node != "bar" ||
		coordinates[1].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 0 {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/nodes?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNodes(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "bar" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Node(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Make sure we get a 404 with no coordinates.
	req, _ := http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp := httptest.NewRecorder()
	obj, err := a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Register the nodes.
	nodes := []string{"foo", "bar"}
	for _, node := range nodes {
		req := structs.RegisterRequest{
			Datacenter: "dc1",
			Node:       node,
			Address:    "127.0.0.1",
		}
		var reply struct{}
		if err := a.RPC("Catalog.Register", &req, &reply); err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	// Send some coordinates for a few nodes, waiting a little while for the
	// batch update to run.
	arg1 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Segment:    "alpha",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	var out struct{}
	if err := a.RPC("Coordinate.Update", &arg1, &out); err != nil {
		t.Fatalf("err: %v", err)
	}

	arg2 := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "bar",
		Coord:      coordinate.NewCoordinate(coordinate.DefaultConfig()),
	}
	if err := a.RPC("Coordinate.Update", &arg2, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the nodes are present.
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?dc=dc1", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates := obj.(structs.Coordinates)
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Filter on a nonexistent node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=nope", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}

	// Filter on a real node segment
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=alpha", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	coordinates = obj.(structs.Coordinates)
	if len(coordinates) != 1 || coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}

	// Make sure the empty filter works
	req, _ = http.NewRequest("GET", "/v1/coordinate/node/foo?segment=", nil)
	resp = httptest.NewRecorder()
	obj, err = a.srv.CoordinateNode(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if resp.Code != http.StatusNotFound {
		t.Fatalf("bad: %v", resp.Code)
	}
}

func TestCoordinate_Update(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), "")
	defer a.Shutdown()
	testrpc.WaitForTestAgent(t, a.RPC, "dc1")

	// Register the node.
	reg := structs.RegisterRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Address:    "127.0.0.1",
	}
	var reply struct{}
	if err := a.RPC("Catalog.Register", &reg, &reply); err != nil {
		t.Fatalf("err: %s", err)
	}

	// Update the coordinates and wait for it to complete.
	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}
	req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
	resp := httptest.NewRecorder()
	_, err := a.srv.CoordinateUpdate(resp, req)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	time.Sleep(300 * time.Millisecond)

	// Query back and check the coordinates are present.
	args := structs.NodeSpecificRequest{Node: "foo", Datacenter: "dc1"}
	var coords structs.IndexedCoordinates
	if err := a.RPC("Coordinate.Node", &args, &coords); err != nil {
		t.Fatalf("err: %s", err)
	}

	coordinates := coords.Coordinates
	if len(coordinates) != 1 ||
		coordinates[0].Node != "foo" {
		t.Fatalf("bad: %v", coordinates)
	}
}

func TestCoordinate_Update_ACLDeny(t *testing.T) {
	t.Parallel()
	a := NewTestAgent(t, t.Name(), TestACLConfig())
	defer a.Shutdown()
	testrpc.WaitForLeader(t, a.RPC, "dc1")

	coord := coordinate.NewCoordinate(coordinate.DefaultConfig())
	coord.Height = -5.0
	body := structs.CoordinateUpdateRequest{
		Datacenter: "dc1",
		Node:       "foo",
		Coord:      coord,
	}

	t.Run("no token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); !acl.IsErrPermissionDenied(err) {
			t.Fatalf("err: %v", err)
		}
	})

	t.Run("valid token", func(t *testing.T) {
		req, _ := http.NewRequest("PUT", "/v1/coordinate/update?token=root", jsonReader(body))
		if _, err := a.srv.CoordinateUpdate(nil, req); err != nil {
			t.Fatalf("err: %v", err)
		}
	})
}
<|endoftext|>


<commit_before>package slack

import (
	"html"
	"testing"
)

func TestUnslackify(t *testing.T) {
	sc := setup(t)
	// all of these should be known
	raw1 := "<@U11A2B8C1> foobars < > <http://www.test.com> or <http://www.test.com|test.com> in <#C03JAPEHJ> for <mailto:test@example.com|test@example.com>"
	want1 := "testorizor1 foobars < > http://www.test.com or http://www.test.com in #dev for test@example.com"

	// contains unknown IDs
	raw2 := "<@U11A2B8C3> foobars < > <http://www.test.com> or <http://www.test.com|test.com> in <#C03JAPAAA> for <mailto:test@example.com|test@example.com>"
	want2 := "@U11A2B8C3 foobars < > http://www.test.com or http://www.test.com in #C03JAPAAA for test@example.com"

	// contains angle brackets but nothing special
	raw3 := "<fooBarBaz < >>"
	want3 := "<fooBarBaz < >>"

	// contains mailto with and without pipe
	raw4 := "<mailto:test@example.org|test-mailto> <mailto:nopipe@example.org>"
	want4 := "test@example.org nopipe@example.org"

	got1 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw1, sc.unSlackify))
	got2 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw2, sc.unSlackify))
	got3 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw3, sc.unSlackify))
	got4 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw4, sc.unSlackify))

	if got1 != want1 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got1)
		t.Logf("Want: %v", want1)
		t.Fail()
	}

	if got2 != want2 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got2)
		t.Logf("Want: %v", want2)
		t.Fail()
	}
	if got3 != want3 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got3)
		t.Logf("Want: %v", want3)
		t.Fail()
	}

	if got4 != want4 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got4)
		t.Logf("Want: %v", want4)
		t.Fail()
	}
}
<commit_msg>change test to reflect new highlight behaviour<commit_after>package slack

import (
	"html"
	"testing"
)

func TestUnslackify(t *testing.T) {
	sc := setup(t)
	// all of these should be known
	raw1 := "<@U11A2B8C1> foobars < > <http://www.test.com> or <http://www.test.com|test.com> in <#C03JAPEHJ> for <mailto:test@example.com|test@example.com>"
	want1 := "@testorizor1 foobars < > http://www.test.com or http://www.test.com in #dev for test@example.com"

	// contains unknown IDs
	raw2 := "<@U11A2B8C3> foobars < > <http://www.test.com> or <http://www.test.com|test.com> in <#C03JAPAAA> for <mailto:test@example.com|test@example.com>"
	want2 := "@U11A2B8C3 foobars < > http://www.test.com or http://www.test.com in #C03JAPAAA for test@example.com"

	// contains angle brackets but nothing special
	raw3 := "<fooBarBaz < >>"
	want3 := "<fooBarBaz < >>"

	// contains mailto with and without pipe
	raw4 := "<mailto:test@example.org|test-mailto> <mailto:nopipe@example.org>"
	want4 := "test@example.org nopipe@example.org"

	got1 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw1, sc.unSlackify))
	got2 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw2, sc.unSlackify))
	got3 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw3, sc.unSlackify))
	got4 := html.UnescapeString(bracketRe.ReplaceAllStringFunc(raw4, sc.unSlackify))

	if got1 != want1 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got1)
		t.Logf("Want: %v", want1)
		t.Fail()
	}

	if got2 != want2 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got2)
		t.Logf("Want: %v", want2)
		t.Fail()
	}
	if got3 != want3 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got3)
		t.Logf("Want: %v", want3)
		t.Fail()
	}

	if got4 != want4 {
		t.Log("Unslackify failed:")
		t.Logf("Got: %v", got4)
		t.Logf("Want: %v", want4)
		t.Fail()
	}
}
<|endoftext|>


<commit_before>package subchal

import (
    "testing"
)

func Test_LoadStoptimes(t *testing.T) {
    t.Parallel()
    SetTestLogger(t)

    stops, _, _ := LoadStops("../test-data/stops.txt")
    stoptimes, err := LoadStoptimes("../test-data/stop_times.txt", stops)

    switch false {
    case err == nil:
        t.Log("Got error from LoadStoptimes():", err)
        t.FailNow()
    case len(stoptimes) == len(stops):
        t.Log("Didn't load stoptimes for enough stops")
        t.FailNow()
    }

    for stopID, stSlice := range stoptimes {
        for _, st := range stSlice {
            if st.Stop.StopID != stopID {
                t.Log("Stoptime at stop", st.Stop.StopID, "doesn't match stoptimes key", stopID)
            }
            if st.Stop.IsStation() {
                t.Log("Stoptime at stop", st.Stop.StopID, "is a station rather than a stop")
            }
        }
    }
}
<commit_msg>Wrote test for Stoptime.DayLater()<commit_after>package subchal

import (
    "testing"
    "time"
)

func Test_Stoptime_DayLater(t *testing.T) {
    t.Parallel()
    SetTestLogger(t)

    initialTime, _ := time.Parse("15:04:05", "11:32:19")
    st := Stoptime{
        "A20131215WKD_WOLF_NS_01",
        initialTime,
        initialTime,
        nil,
        1,
        "",
    }
    newSt := st.DayLater()

    switch false {
    case newSt.TripID == st.TripID:
        t.Log("Stoptime.DayLater() didn't preserve TripID")
        t.FailNow()
    case 24 * time.Hour == newSt.ArrivalTime.Sub(st.ArrivalTime):
        t.Log("Stoptime.DayLater() didn't transform ArrivalTime")
        t.FailNow()
    case 24 * time.Hour == newSt.DepartureTime.Sub(st.DepartureTime):
        t.Log("Stoptime.DayLater() didn't transform DepartureTime")
        t.FailNow()
    }
}

func Test_LoadStoptimes(t *testing.T) {
    t.Parallel()
    SetTestLogger(t)

    stops, _, _ := LoadStops("../test-data/stops.txt")
    stoptimes, err := LoadStoptimes("../test-data/stop_times.txt", stops)

    switch false {
    case err == nil:
        t.Log("Got error from LoadStoptimes():", err)
        t.FailNow()
    case len(stoptimes) == len(stops):
        t.Log("Didn't load stoptimes for enough stops")
        t.FailNow()
    }

    for stopID, stSlice := range stoptimes {
        for _, st := range stSlice {
            if st.Stop.StopID != stopID {
                t.Log("Stoptime at stop", st.Stop.StopID, "doesn't match stoptimes key", stopID)
            }
            if st.Stop.IsStation() {
                t.Log("Stoptime at stop", st.Stop.StopID, "is a station rather than a stop")
            }
        }
    }
}
<|endoftext|>


<commit_before>package rpc

import (
	"bytes"
	"io"
	"log"
	"time"

	"golang.org/x/net/context"
	"zombiezen.com/go/capnproto"
	"zombiezen.com/go/capnproto/rpc/rpccapnp"
)

// Transport is the interface that abstracts sending and receiving
// individual messages of the Cap'n Proto RPC protocol.
type Transport interface {
	// SendMessage sends msg.
	SendMessage(ctx context.Context, msg rpccapnp.Message) error

	// RecvMessage waits to receive a message and returns it.
	// Implementations may re-use buffers between calls, so the message is
	// only valid until the next call to RecvMessage.
	RecvMessage(ctx context.Context) (rpccapnp.Message, error)

	// Close releases any resources associated with the transport.
	Close() error
}

type streamTransport struct {
	rwc      io.ReadWriteCloser
	deadline writeDeadlineSetter
	rbuf     bytes.Buffer
	wbuf     bytes.Buffer
}

// StreamTransport creates a transport that sends and receives messages
// by serializing and deserializing unpacked Cap'n Proto messages.
// Closing the transport will close the underlying ReadWriteCloser.
func StreamTransport(rwc io.ReadWriteCloser) Transport {
	d, _ := rwc.(writeDeadlineSetter)
	s := &streamTransport{rwc: rwc, deadline: d}
	s.rbuf.Grow(4096)
	s.wbuf.Grow(4096)
	return s
}

func (s *streamTransport) SendMessage(ctx context.Context, msg rpccapnp.Message) error {
	s.wbuf.Reset()
	if _, err := msg.Segment.WriteTo(&s.wbuf); err != nil {
		return err
	}
	// TODO(light): clear deadlines
	if d, ok := ctx.Deadline(); ok && s.deadline != nil {
		// TODO(light): log error
		s.deadline.SetWriteDeadline(d)
	}
	_, err := s.rwc.Write(s.wbuf.Bytes())
	return err
}

func (s *streamTransport) RecvMessage(ctx context.Context) (rpccapnp.Message, error) {
	seg, err := capnp.ReadFromStream(s.rwc, &s.rbuf)
	if err != nil {
		return rpccapnp.Message{}, err
	}
	msg := rpccapnp.ReadRootMessage(seg)
	return msg, nil
}

func (s *streamTransport) Close() error {
	return s.rwc.Close()
}

type writeDeadlineSetter interface {
	SetWriteDeadline(t time.Time) error
}

// dispatchSend runs in its own goroutine and sends messages on a transport.
func dispatchSend(m *manager, transport Transport, msgs <-chan outgoingMessage) {
	for {
		select {
		case m := <-msgs:
			// TODO(light): allow custom context
			err := transport.SendMessage(m.ctx, m.msg)
			if err != nil {
				log.Println("rpc: write error:", err)
			}
		case <-m.finish:
			return
		}
	}
}

// An outgoingMessage holds a packet to be sent on a transport and its context.
type outgoingMessage struct {
	ctx context.Context
	msg rpccapnp.Message
}

// dispatchRecv runs in its own goroutine and receives messages from a transport.
func dispatchRecv(m *manager, transport Transport, msgs chan<- rpccapnp.Message) {
	for {
		msg, err := transport.RecvMessage(m.context())
		if err != nil {
			if isTemporaryError(err) {
				log.Println("rpc: read temporary error:", err)
				continue
			}
			m.shutdown(err)
			return
		}
		select {
		case msgs <- copyRPCMessage(msg):
		case <-m.finish:
			return
		}
	}
}

// copyMessage clones a Cap'n Proto buffer.
func copyMessage(msg capnp.Message) capnp.Message {
	n := 0
	for {
		if _, err := msg.Lookup(uint32(n)); err != nil {
			break
		}
		n++
	}
	segments := make([][]byte, n)
	for i := range segments {
		s, err := msg.Lookup(uint32(i))
		if err != nil {
			panic(err)
		}
		segments[i] = make([]byte, len(s.Data))
		copy(segments[i], s.Data)
	}
	return capnp.NewMultiBuffer(segments).Message
}

// copyRPCMessage clones an RPC packet.
func copyRPCMessage(m rpccapnp.Message) rpccapnp.Message {
	mm := copyMessage(m.Segment.Message)
	seg, err := mm.Lookup(0)
	if err != nil {
		panic(err)
	}
	return rpccapnp.ReadRootMessage(seg)
}

// isTemporaryError reports whether e has a Temporary() method that
// returns true.
func isTemporaryError(e error) bool {
	type temp interface {
		Temporary() bool
	}
	t, ok := e.(temp)
	return ok && t.Temporary()
}
<commit_msg>rpc: log message type on transport write failure<commit_after>package rpc

import (
	"bytes"
	"io"
	"log"
	"time"

	"golang.org/x/net/context"
	"zombiezen.com/go/capnproto"
	"zombiezen.com/go/capnproto/rpc/rpccapnp"
)

// Transport is the interface that abstracts sending and receiving
// individual messages of the Cap'n Proto RPC protocol.
type Transport interface {
	// SendMessage sends msg.
	SendMessage(ctx context.Context, msg rpccapnp.Message) error

	// RecvMessage waits to receive a message and returns it.
	// Implementations may re-use buffers between calls, so the message is
	// only valid until the next call to RecvMessage.
	RecvMessage(ctx context.Context) (rpccapnp.Message, error)

	// Close releases any resources associated with the transport.
	Close() error
}

type streamTransport struct {
	rwc      io.ReadWriteCloser
	deadline writeDeadlineSetter
	rbuf     bytes.Buffer
	wbuf     bytes.Buffer
}

// StreamTransport creates a transport that sends and receives messages
// by serializing and deserializing unpacked Cap'n Proto messages.
// Closing the transport will close the underlying ReadWriteCloser.
func StreamTransport(rwc io.ReadWriteCloser) Transport {
	d, _ := rwc.(writeDeadlineSetter)
	s := &streamTransport{rwc: rwc, deadline: d}
	s.rbuf.Grow(4096)
	s.wbuf.Grow(4096)
	return s
}

func (s *streamTransport) SendMessage(ctx context.Context, msg rpccapnp.Message) error {
	s.wbuf.Reset()
	if _, err := msg.Segment.WriteTo(&s.wbuf); err != nil {
		return err
	}
	// TODO(light): clear deadlines
	if d, ok := ctx.Deadline(); ok && s.deadline != nil {
		// TODO(light): log error
		s.deadline.SetWriteDeadline(d)
	}
	_, err := s.rwc.Write(s.wbuf.Bytes())
	return err
}

func (s *streamTransport) RecvMessage(ctx context.Context) (rpccapnp.Message, error) {
	seg, err := capnp.ReadFromStream(s.rwc, &s.rbuf)
	if err != nil {
		return rpccapnp.Message{}, err
	}
	msg := rpccapnp.ReadRootMessage(seg)
	return msg, nil
}

func (s *streamTransport) Close() error {
	return s.rwc.Close()
}

type writeDeadlineSetter interface {
	SetWriteDeadline(t time.Time) error
}

// dispatchSend runs in its own goroutine and sends messages on a transport.
func dispatchSend(m *manager, transport Transport, msgs <-chan outgoingMessage) {
	for {
		select {
		case m := <-msgs:
			// TODO(light): allow custom context
			err := transport.SendMessage(m.ctx, m.msg)
			if err != nil {
				log.Printf("rpc: writing %v: %v", m.msg.Which(), err)
			}
		case <-m.finish:
			return
		}
	}
}

// An outgoingMessage holds a packet to be sent on a transport and its context.
type outgoingMessage struct {
	ctx context.Context
	msg rpccapnp.Message
}

// dispatchRecv runs in its own goroutine and receives messages from a transport.
func dispatchRecv(m *manager, transport Transport, msgs chan<- rpccapnp.Message) {
	for {
		msg, err := transport.RecvMessage(m.context())
		if err != nil {
			if isTemporaryError(err) {
				log.Println("rpc: read temporary error:", err)
				continue
			}
			m.shutdown(err)
			return
		}
		select {
		case msgs <- copyRPCMessage(msg):
		case <-m.finish:
			return
		}
	}
}

// copyMessage clones a Cap'n Proto buffer.
func copyMessage(msg capnp.Message) capnp.Message {
	n := 0
	for {
		if _, err := msg.Lookup(uint32(n)); err != nil {
			break
		}
		n++
	}
	segments := make([][]byte, n)
	for i := range segments {
		s, err := msg.Lookup(uint32(i))
		if err != nil {
			panic(err)
		}
		segments[i] = make([]byte, len(s.Data))
		copy(segments[i], s.Data)
	}
	return capnp.NewMultiBuffer(segments).Message
}

// copyRPCMessage clones an RPC packet.
func copyRPCMessage(m rpccapnp.Message) rpccapnp.Message {
	mm := copyMessage(m.Segment.Message)
	seg, err := mm.Lookup(0)
	if err != nil {
		panic(err)
	}
	return rpccapnp.ReadRootMessage(seg)
}

// isTemporaryError reports whether e has a Temporary() method that
// returns true.
func isTemporaryError(e error) bool {
	type temp interface {
		Temporary() bool
	}
	t, ok := e.(temp)
	return ok && t.Temporary()
}
<|endoftext|>


<commit_before>package main

import "fmt"
import "io"
import "os"
import "os/exec"
import "strings"

import "./alias"
import "./conio"
import "./dos"
import "./history"
import "./interpreter"
import "./lua"
import "bytes"

import "github.com/shiena/ansicolor"

func cmdAlias(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	key := strings.ToLower(name)
	switch L.GetType(2) {
	case lua.LUA_TSTRING:
		value, err := L.ToString(2)
		if err == nil {
			alias.Table[key] = alias.New(value)
		} else {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
	case lua.LUA_TFUNCTION:
		regkey := "nyagos.alias." + key
		L.SetField(lua.REGISTORYINDEX, regkey)
		alias.Table[key] = LuaFunction{L, regkey}
	}
	L.PushBool(true)
	return 1
}

func cmdSetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	value, valueErr := L.ToString(2)
	if valueErr != nil {
		L.PushNil()
		L.PushString(valueErr.Error())
		return 2
	}
	os.Setenv(name, value)
	L.PushBool(true)
	return 1
}

func cmdGetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		return 1
	}
	value := os.Getenv(name)
	if len(value) > 0 {
		L.PushString(value)
	} else {
		L.PushNil()
	}
	return 1
}

func cmdExec(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	_, err := interpreter.New().Interpret(statement)

	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	L.PushBool(true)
	return 1
}

func cmdEval(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	r, w, err := os.Pipe()
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	go func(statement string, w *os.File) {
		it := interpreter.New()
		it.Stdout = w
		it.Interpret(statement)
		w.Close()
	}(statement, w)

	var result = []byte{}
	for {
		buffer := make([]byte, 256)
		size, err := r.Read(buffer)
		if err != nil || size <= 0 {
			break
		}
		result = append(result, buffer[0:size]...)
	}
	r.Close()
	L.PushAnsiString(bytes.Trim(result, "\r\n\t "))
	return 1
}

func cmdWrite(L *lua.Lua) int {
	var out io.Writer = os.Stdout
	cmd, cmdOk := LuaInstanceToCmd[L.Id()]
	if cmdOk && cmd != nil && cmd.Stdout != nil {
		out = cmd.Stdout
	}
	switch out.(type) {
	case *os.File:
		out = ansicolor.NewAnsiColorWriter(out)
	}

	n := L.GetTop()
	for i := 1; i <= n; i++ {
		str, err := L.ToString(i)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
		if i > 1 {
			fmt.Fprint(out, "\t")
		}
		fmt.Fprint(out, str)
	}
	L.PushBool(true)
	return 1
}

func cmdGetwd(L *lua.Lua) int {
	wd, err := os.Getwd()
	if err == nil {
		L.PushString(wd)
		return 1
	} else {
		return 0
	}
}

func cmdWhich(L *lua.Lua) int {
	if L.GetType(-1) != lua.LUA_TSTRING {
		return 0
	}
	name, nameErr := L.ToString(-1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	path, err := exec.LookPath(name)
	if err == nil {
		L.PushString(path)
		return 1
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdAtoU(L *lua.Lua) int {
	str, err := dos.AtoU(L.ToAnsiString(1))
	if err == nil {
		L.PushString(str)
		return 1
	} else {
		return 0
	}
}

func cmdUtoA(L *lua.Lua) int {
	utf8, utf8err := L.ToString(1)
	if utf8err != nil {
		L.PushNil()
		L.PushString(utf8err.Error())
		return 2
	}
	str, err := dos.UtoA(utf8)
	if err == nil {
		if len(str) >= 1 {
			L.PushAnsiString(str[:len(str)-1])
		} else {
			L.PushString("")
		}
		L.PushNil()
		return 2
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdGlob(L *lua.Lua) int {
	if !L.IsString(-1) {
		return 0
	}
	wildcard, wildcardErr := L.ToString(-1)
	if wildcardErr != nil {
		L.PushNil()
		L.PushString(wildcardErr.Error())
		return 2
	}
	list, err := dos.Glob(wildcard)
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	} else {
		L.NewTable()
		for i := 0; i < len(list); i++ {
			L.PushString(list[i])
			L.RawSetI(-2, i+1)
		}
		return 1
	}
}

func cmdGetHistory(this *lua.Lua) int {
	if this.GetType(-1) == lua.LUA_TNUMBER {
		val, err := this.ToInteger(-1)
		if err != nil {
			this.PushNil()
			this.PushString(err.Error())
		}
		this.PushString(history.Get(val))
	} else {
		this.PushInteger(history.Len())
	}
	return 1
}

func cmdSetRuneWidth(this *lua.Lua) int {
	char, charErr := this.ToInteger(1)
	if charErr != nil {
		this.PushNil()
		this.PushString(charErr.Error())
		return 2
	}
	width, widthErr := this.ToInteger(2)
	if widthErr != nil {
		this.PushNil()
		this.PushString(widthErr.Error())
		return 2
	}
	conio.SetCharWidth(rune(char), width)
	this.PushBool(true)
	return 1
}

func cmdShellExecute(this *lua.Lua) int {
	action, actionErr := this.ToString(1)
	if actionErr != nil {
		this.PushNil()
		this.PushString(actionErr.Error())
		return 2
	}
	path, pathErr := this.ToString(2)
	if pathErr != nil {
		this.PushNil()
		this.PushString(pathErr.Error())
		return 2
	}
	param, paramErr := this.ToString(3)
	if paramErr != nil {
		param = ""
	}
	dir, dirErr := this.ToString(4)
	if dirErr != nil {
		dir = ""
	}
	err := dos.ShellExecute(action, path, param, dir)
	if err != nil {
		this.PushNil()
		this.PushString(err.Error())
	} else {
		this.PushBool(true)
	}
	return 1
}

func cmdAccess(L *lua.Lua) int {
	path, pathErr := L.ToString(1)
	if pathErr != nil {
		L.PushNil()
		L.PushString(pathErr.Error())
		return 2
	}
	mode, modeErr := L.ToInteger(2)
	if modeErr != nil {
		L.PushNil()
		L.PushString(modeErr.Error())
		return 2
	}
	fi, err := os.Stat(path)

	var result bool
	if err != nil {
		result = false
	} else {
		switch {
		case mode == 0:
			result = true
		case mode & 1 != 0: // X_OK
		case mode & 2 != 0: // W_OK
			result = fi.Mode().Perm() & 0200 != 0
		case mode & 4 != 0: // R_OK
			result = fi.Mode().Perm() & 0400 != 0
		}
	}
	L.PushBool(result)
	return 1
}
<commit_msg>Suppress warnings<commit_after>package main

import "fmt"
import "io"
import "os"
import "os/exec"
import "strings"

import "./alias"
import "./conio"
import "./dos"
import "./history"
import "./interpreter"
import "./lua"
import "bytes"

import "github.com/shiena/ansicolor"

func cmdAlias(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	key := strings.ToLower(name)
	switch L.GetType(2) {
	case lua.LUA_TSTRING:
		value, err := L.ToString(2)
		if err == nil {
			alias.Table[key] = alias.New(value)
		} else {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
	case lua.LUA_TFUNCTION:
		regkey := "nyagos.alias." + key
		L.SetField(lua.REGISTORYINDEX, regkey)
		alias.Table[key] = LuaFunction{L, regkey}
	}
	L.PushBool(true)
	return 1
}

func cmdSetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	value, valueErr := L.ToString(2)
	if valueErr != nil {
		L.PushNil()
		L.PushString(valueErr.Error())
		return 2
	}
	os.Setenv(name, value)
	L.PushBool(true)
	return 1
}

func cmdGetEnv(L *lua.Lua) int {
	name, nameErr := L.ToString(1)
	if nameErr != nil {
		L.PushNil()
		return 1
	}
	value := os.Getenv(name)
	if len(value) > 0 {
		L.PushString(value)
	} else {
		L.PushNil()
	}
	return 1
}

func cmdExec(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	_, err := interpreter.New().Interpret(statement)

	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	L.PushBool(true)
	return 1
}

type emptyWriter struct { }

func (e *emptyWriter) Write(b []byte) (int, error) {
	return len(b), nil
}

func cmdEval(L *lua.Lua) int {
	statement, statementErr := L.ToString(1)
	if statementErr != nil {
		L.PushNil()
		L.PushString(statementErr.Error())
		return 2
	}
	r, w, err := os.Pipe()
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
	go func(statement string, w *os.File) {
		it := interpreter.New()
		it.Stdout = w
		it.Stderr = &emptyWriter{}
		it.Interpret(statement)
		w.Close()
	}(statement, w)

	var result = []byte{}
	for {
		buffer := make([]byte, 256)
		size, err := r.Read(buffer)
		if err != nil || size <= 0 {
			break
		}
		result = append(result, buffer[0:size]...)
	}
	r.Close()
	L.PushAnsiString(bytes.Trim(result, "\r\n\t "))
	return 1
}

func cmdWrite(L *lua.Lua) int {
	var out io.Writer = os.Stdout
	cmd, cmdOk := LuaInstanceToCmd[L.Id()]
	if cmdOk && cmd != nil && cmd.Stdout != nil {
		out = cmd.Stdout
	}
	switch out.(type) {
	case *os.File:
		out = ansicolor.NewAnsiColorWriter(out)
	}

	n := L.GetTop()
	for i := 1; i <= n; i++ {
		str, err := L.ToString(i)
		if err != nil {
			L.PushNil()
			L.PushString(err.Error())
			return 2
		}
		if i > 1 {
			fmt.Fprint(out, "\t")
		}
		fmt.Fprint(out, str)
	}
	L.PushBool(true)
	return 1
}

func cmdGetwd(L *lua.Lua) int {
	wd, err := os.Getwd()
	if err == nil {
		L.PushString(wd)
		return 1
	} else {
		return 0
	}
}

func cmdWhich(L *lua.Lua) int {
	if L.GetType(-1) != lua.LUA_TSTRING {
		return 0
	}
	name, nameErr := L.ToString(-1)
	if nameErr != nil {
		L.PushNil()
		L.PushString(nameErr.Error())
		return 2
	}
	path, err := exec.LookPath(name)
	if err == nil {
		L.PushString(path)
		return 1
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdAtoU(L *lua.Lua) int {
	str, err := dos.AtoU(L.ToAnsiString(1))
	if err == nil {
		L.PushString(str)
		return 1
	} else {
		return 0
	}
}

func cmdUtoA(L *lua.Lua) int {
	utf8, utf8err := L.ToString(1)
	if utf8err != nil {
		L.PushNil()
		L.PushString(utf8err.Error())
		return 2
	}
	str, err := dos.UtoA(utf8)
	if err == nil {
		if len(str) >= 1 {
			L.PushAnsiString(str[:len(str)-1])
		} else {
			L.PushString("")
		}
		L.PushNil()
		return 2
	} else {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	}
}

func cmdGlob(L *lua.Lua) int {
	if !L.IsString(-1) {
		return 0
	}
	wildcard, wildcardErr := L.ToString(-1)
	if wildcardErr != nil {
		L.PushNil()
		L.PushString(wildcardErr.Error())
		return 2
	}
	list, err := dos.Glob(wildcard)
	if err != nil {
		L.PushNil()
		L.PushString(err.Error())
		return 2
	} else {
		L.NewTable()
		for i := 0; i < len(list); i++ {
			L.PushString(list[i])
			L.RawSetI(-2, i+1)
		}
		return 1
	}
}

func cmdGetHistory(this *lua.Lua) int {
	if this.GetType(-1) == lua.LUA_TNUMBER {
		val, err := this.ToInteger(-1)
		if err != nil {
			this.PushNil()
			this.PushString(err.Error())
		}
		this.PushString(history.Get(val))
	} else {
		this.PushInteger(history.Len())
	}
	return 1
}

func cmdSetRuneWidth(this *lua.Lua) int {
	char, charErr := this.ToInteger(1)
	if charErr != nil {
		this.PushNil()
		this.PushString(charErr.Error())
		return 2
	}
	width, widthErr := this.ToInteger(2)
	if widthErr != nil {
		this.PushNil()
		this.PushString(widthErr.Error())
		return 2
	}
	conio.SetCharWidth(rune(char), width)
	this.PushBool(true)
	return 1
}

func cmdShellExecute(this *lua.Lua) int {
	action, actionErr := this.ToString(1)
	if actionErr != nil {
		this.PushNil()
		this.PushString(actionErr.Error())
		return 2
	}
	path, pathErr := this.ToString(2)
	if pathErr != nil {
		this.PushNil()
		this.PushString(pathErr.Error())
		return 2
	}
	param, paramErr := this.ToString(3)
	if paramErr != nil {
		param = ""
	}
	dir, dirErr := this.ToString(4)
	if dirErr != nil {
		dir = ""
	}
	err := dos.ShellExecute(action, path, param, dir)
	if err != nil {
		this.PushNil()
		this.PushString(err.Error())
	} else {
		this.PushBool(true)
	}
	return 1
}

func cmdAccess(L *lua.Lua) int {
	path, pathErr := L.ToString(1)
	if pathErr != nil {
		L.PushNil()
		L.PushString(pathErr.Error())
		return 2
	}
	mode, modeErr := L.ToInteger(2)
	if modeErr != nil {
		L.PushNil()
		L.PushString(modeErr.Error())
		return 2
	}
	fi, err := os.Stat(path)

	var result bool
	if err != nil {
		result = false
	} else {
		switch {
		case mode == 0:
			result = true
		case mode & 1 != 0: // X_OK
		case mode & 2 != 0: // W_OK
			result = fi.Mode().Perm() & 0200 != 0
		case mode & 4 != 0: // R_OK
			result = fi.Mode().Perm() & 0400 != 0
		}
	}
	L.PushBool(result)
	return 1
}
<|endoftext|>


<commit_before>package rrs

import (
	"sync"

	"golang.org/x/net/context"

	"github.com/teh-cmc/seq"
)

// -----------------------------------------------------------------------------

// RRSeq implements a buffered `Sequencer` backed by a cluster of `RRServer`s: a
// distributed system that guarantees sequential `ID` generation by using RW
// quorums and read-repair conflict resolution strategies.
//
// You can find more information about the ideas behind such a system in the
// `README.md` file at the root of this repository.
type RRSeq struct {
	cp  *rrAPIPool
	ids chan seq.ID

	stop chan struct{}
	wg   *sync.WaitGroup
}

// NewRRSeq returns a new `RRSeq` with the specified buffer size.
//
// `RRSeq` will constantly try to keep this buffer as full as possible
// using a dedicated background routine.
//
// `bufSize` will default to 0 in case it is < 0.
func NewRRSeq(name string, bufSize int, addrs ...string) (*RRSeq, error) {
	if bufSize < 0 {
		bufSize = 0
	}

	cp, err := newRRAPIPool(addrs...)
	if err != nil {
		return nil, err
	}

	ids := make(chan seq.ID, bufSize)
	stop := make(chan struct{}, 0)
	wg := &sync.WaitGroup{}
	rrseq := &RRSeq{cp: cp, ids: ids, stop: stop, wg: wg}

	var curRange, nextRange [2]seq.ID
	wg.Add(1)
	go func() { // background buffering routine
		defer wg.Done()
		for {

			// if the current range has already been half (or more) consumed:
			//  I: replace current range with the (already stored) next range
			// II: fetch the next available range from the cluster
			if curRange[0] >= curRange[1]/2 {
				curRange = nextRange
				nextRange[0], nextRange[1] = rrseq.getNextRange(name, bufSize)
			}

			select {
			case <-stop: // stream has been closed, kill routine
				return
			default:
				if curRange[0] < curRange[1] {
					select {
					case ids <- curRange[0]: // blocks once `seq` is full
						curRange[0]++
					case <-stop: // stream has been closed, kill routine
						return
					}
				}
			}

		}
	}()

	return rrseq, nil
}

// -----------------------------------------------------------------------------

// getNextRange fetches the next available range of `ID`s from the cluster
// of `RRServer`s.
func (ss RRSeq) getNextRange(name string, rangeSize int) (seq.ID, seq.ID) {
	idReply, err := ss.cp.Client().GRPCNextID(
		//          ^^^^^^^^^^^^^^
		//                 ^--- transparent round-robin
		context.TODO(), // TODO: handle timeouts
		&NextIDRequest{Name: name, RangeSize: int64(rangeSize)},
	)
	if err != nil {
		// NOTE: if things go wrong, an empty `ID` range is returned
		// this will be retried on the next iteration, using the next client (round-robin)
		return 0, 0
	}
	return seq.ID(idReply.FromId), seq.ID(idReply.ToId)
}

// GetStream returns a pre-buffered, range-able stream of sequential `ID`s,
// starting at `1`.
func (ss RRSeq) GetStream() seq.IDStream { return ss.ids }

// Close closes the associated `IDStream`.
//
// It returns the first error met.
func (ss *RRSeq) Close() error {
	close(ss.stop)       // kill background routine
	ss.wg.Wait()         // wait for it to be fully stopped
	close(ss.ids)        // close `ID` stream
	return ss.cp.Close() // close gRPC clients
}
<commit_msg>RRSeq: fixed cur/next range algorithm<commit_after>package rrs

import (
	"sync"

	"golang.org/x/net/context"

	"github.com/teh-cmc/seq"
)

// -----------------------------------------------------------------------------

// RRSeq implements a buffered `Sequencer` backed by a cluster of `RRServer`s: a
// distributed system that guarantees sequential `ID` generation by using RW
// quorums and read-repair conflict resolution strategies.
//
// You can find more information about the ideas behind such a system in the
// `README.md` file at the root of this repository.
type RRSeq struct {
	cp  *rrAPIPool
	ids chan seq.ID

	stop chan struct{}
	wg   *sync.WaitGroup
}

// NewRRSeq returns a new `RRSeq` with the specified buffer size.
//
// `RRSeq` will constantly try to keep this buffer as full as possible
// using a dedicated background routine.
//
// `bufSize` will default to 0 in case it is < 0.
func NewRRSeq(name string, bufSize int, addrs ...string) (*RRSeq, error) {
	if bufSize < 0 {
		bufSize = 0
	}

	cp, err := newRRAPIPool(addrs...)
	if err != nil {
		return nil, err
	}

	ids := make(chan seq.ID, bufSize)
	stop := make(chan struct{}, 0)
	wg := &sync.WaitGroup{}
	rrseq := &RRSeq{cp: cp, ids: ids, stop: stop, wg: wg}

	var curRange, nextRange [2]seq.ID
	wg.Add(1)
	go func() { // background buffering routine
		defer wg.Done()
		for {

			// if the current range has already been half (or more) consumed,
			// fetch and store the next available range from the cluster
			if curRange[0] >= curRange[1]/2 {
				nextRange[0], nextRange[1] = rrseq.getNextRange(name, bufSize)
				// if the current range has been entirely consumed,
				// replace the current range by the already stored next range
				if curRange[0] >= curRange[1] {
					curRange = nextRange
				}
			}

			select {
			case <-stop: // stream has been closed, kill routine
				return
			default:
				if curRange[0] < curRange[1] {
					select {
					case ids <- curRange[0]: // blocks once `seq` is full
						curRange[0]++
					case <-stop: // stream has been closed, kill routine
						return
					}
				}
			}

		}
	}()

	return rrseq, nil
}

// -----------------------------------------------------------------------------

// getNextRange fetches the next available range of `ID`s from the cluster
// of `RRServer`s.
func (ss RRSeq) getNextRange(name string, rangeSize int) (seq.ID, seq.ID) {
	idReply, err := ss.cp.Client().GRPCNextID(
		//          ^^^^^^^^^^^^^^
		//                 ^--- transparent round-robin
		context.TODO(), // TODO: handle timeouts
		&NextIDRequest{Name: name, RangeSize: int64(rangeSize)},
	)
	if err != nil {
		// NOTE: if things go wrong, an empty `ID` range is returned
		// this will be retried on the next iteration, using the next client (round-robin)
		return 0, 0
	}
	return seq.ID(idReply.FromId), seq.ID(idReply.ToId)
}

// GetStream returns a pre-buffered, range-able stream of sequential `ID`s,
// starting at `1`.
func (ss RRSeq) GetStream() seq.IDStream { return ss.ids }

// Close closes the associated `IDStream`.
//
// It returns the first error met.
func (ss *RRSeq) Close() error {
	close(ss.stop)       // kill background routine
	ss.wg.Wait()         // wait for it to be fully stopped
	close(ss.ids)        // close `ID` stream
	return ss.cp.Close() // close gRPC clients
}
<|endoftext|>


<commit_before>package storage

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/mailhog/data"
)

// Maildir is a maildir storage backend
type Maildir struct {
	Path string
}

// CreateMaildir creates a new maildir storage backend
func CreateMaildir(path string) *Maildir {
	if len(path) == 0 {
		dir, err := ioutil.TempDir("", "mailhog")
		if err != nil {
			panic(err)
		}
		path = dir
	}
	if _, err := os.Stat(path); err != nil {
		err := os.MkdirAll(path, 0770)
		if err != nil {
			panic(err)
		}
	}
	log.Println("Maildir path is", path)
	return &Maildir{
		Path: path,
	}
}

// Store stores a message and returns its storage ID
func (maildir *Maildir) Store(m *data.Message) (string, error) {
	b, err := ioutil.ReadAll(m.Raw.Bytes())
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(filepath.Join(maildir.Path, string(m.ID)), b, 0660)
	return string(m.ID), err
}

// Count returns the number of stored messages
func (maildir *Maildir) Count() int {
	// FIXME may be wrong, ../. ?
	// and handle error?
	dir, err := os.Open(maildir.Path)
	if err != nil {
		panic(err)
	}
	defer dir.Close()
	n, _ := dir.Readdirnames(0)
	return len(n)
}

// Search finds messages matching the query
func (maildir *Maildir) Search(kind, query string, start, limit int) (*data.Messages, int, error) {
	query = strings.ToLower(query)
	var filteredMessages = make([]data.Message, 0)

	// for _, m := range memory.Messages {
	// 	switch kind {
	// 	case "to":
	// 	case "from":
	// 	case "containing":
	// 	}
	// }
	msgs := data.Messages(filteredMessages)
	return &msgs, len(filteredMessages), nil
}

// List lists stored messages by index
func (maildir *Maildir) List(start, limit int) (*data.Messages, error) {
	log.Println("Listing messages in", maildir.Path)
	var messages = make([]data.Message, 0)

	dir, err := os.Open(maildir.Path)
	if err != nil {
		return nil, err
	}
	defer dir.Close()

	n, err := dir.Readdirnames(0)
	if err != nil {
		return nil, err
	}

	for _, path := range n {
		b, err := ioutil.ReadFile(filepath.Join(maildir.Path, path))
		if err != nil {
			return nil, err
		}

		msg := data.FromBytes(b)
		// FIXME domain
		m := *msg.Parse("mailhog.example")
		m.ID = data.MessageID(path)
		messages = append(messages, m)
	}

	log.Printf("Found %d messages", len(messages))
	msgs := data.Messages(messages)
	return &msgs, nil
}

// DeleteOne deletes an individual message by storage ID
func (maildir *Maildir) DeleteOne(id string) error {
	return os.Remove(filepath.Join(maildir.Path, id))
}

// DeleteAll deletes all in memory messages
func (maildir *Maildir) DeleteAll() error {
	err := os.RemoveAll(maildir.Path)
	if err != nil {
		return err
	}
	return os.Mkdir(maildir.Path, 0770)
}

// Load returns an individual message by storage ID
func (maildir *Maildir) Load(id string) (*data.Message, error) {
	b, err := ioutil.ReadFile(filepath.Join(maildir.Path, id))
	if err != nil {
		return nil, err
	}
	// FIXME domain
	m := data.FromBytes(b).Parse("mailhog.example")
	m.ID = data.MessageID(id)
	return m, nil
}
<commit_msg>implement maildir search<commit_after>package storage

import (
	"errors"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/mailhog/data"
)

// Maildir is a maildir storage backend
type Maildir struct {
	Path string
}

// CreateMaildir creates a new maildir storage backend
func CreateMaildir(path string) *Maildir {
	if len(path) == 0 {
		dir, err := ioutil.TempDir("", "mailhog")
		if err != nil {
			panic(err)
		}
		path = dir
	}
	if _, err := os.Stat(path); err != nil {
		err := os.MkdirAll(path, 0770)
		if err != nil {
			panic(err)
		}
	}
	log.Println("Maildir path is", path)
	return &Maildir{
		Path: path,
	}
}

// Store stores a message and returns its storage ID
func (maildir *Maildir) Store(m *data.Message) (string, error) {
	b, err := ioutil.ReadAll(m.Raw.Bytes())
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(filepath.Join(maildir.Path, string(m.ID)), b, 0660)
	return string(m.ID), err
}

// Count returns the number of stored messages
func (maildir *Maildir) Count() int {
	// FIXME may be wrong, ../. ?
	// and handle error?
	dir, err := os.Open(maildir.Path)
	if err != nil {
		panic(err)
	}
	defer dir.Close()
	n, _ := dir.Readdirnames(0)
	return len(n)
}

// Search finds messages matching the query
func (maildir *Maildir) Search(kind, query string, start, limit int) (*data.Messages, int, error) {
	query = strings.ToLower(query)
	var filteredMessages = make([]data.Message, 0)

	var matched int

	err := filepath.Walk(maildir.Path, func(path string, info os.FileInfo, err error) error {
		if limit > 0 && len(filteredMessages) >= limit {
			return errors.New("reached limit")
		}

		if info.IsDir() {
			return nil
		}

		msg, err := maildir.Load(info.Name())
		if err != nil {
			log.Println(err)
			return nil
		}

		switch kind {
		case "to":
			for _, t := range msg.To {
				if strings.Contains(t.Mailbox+"@"+t.Domain, query) {
					if start > matched {
						matched++
						break
					}
					filteredMessages = append(filteredMessages, *msg)
					break
				}
			}
		case "from":
			if strings.Contains(msg.From.Mailbox+"@"+msg.From.Domain, query) {
				if start > matched {
					matched++
					break
				}
				filteredMessages = append(filteredMessages, *msg)
			}
		case "containing":
			if strings.Contains(msg.Raw.Data, query) {
				if start > matched {
					matched++
					break
				}
				filteredMessages = append(filteredMessages, *msg)
			}
		}

		return nil
	})

	if err != nil {
		log.Println(err)
	}

	msgs := data.Messages(filteredMessages)
	return &msgs, len(filteredMessages), nil
}

// List lists stored messages by index
func (maildir *Maildir) List(start, limit int) (*data.Messages, error) {
	log.Println("Listing messages in", maildir.Path)
	var messages = make([]data.Message, 0)

	dir, err := os.Open(maildir.Path)
	if err != nil {
		return nil, err
	}
	defer dir.Close()

	n, err := dir.Readdirnames(0)
	if err != nil {
		return nil, err
	}

	for _, path := range n {
		b, err := ioutil.ReadFile(filepath.Join(maildir.Path, path))
		if err != nil {
			return nil, err
		}

		msg := data.FromBytes(b)
		// FIXME domain
		m := *msg.Parse("mailhog.example")
		m.ID = data.MessageID(path)
		messages = append(messages, m)
	}

	log.Printf("Found %d messages", len(messages))
	msgs := data.Messages(messages)
	return &msgs, nil
}

// DeleteOne deletes an individual message by storage ID
func (maildir *Maildir) DeleteOne(id string) error {
	return os.Remove(filepath.Join(maildir.Path, id))
}

// DeleteAll deletes all in memory messages
func (maildir *Maildir) DeleteAll() error {
	err := os.RemoveAll(maildir.Path)
	if err != nil {
		return err
	}
	return os.Mkdir(maildir.Path, 0770)
}

// Load returns an individual message by storage ID
func (maildir *Maildir) Load(id string) (*data.Message, error) {
	b, err := ioutil.ReadFile(filepath.Join(maildir.Path, id))
	if err != nil {
		return nil, err
	}
	// FIXME domain
	m := data.FromBytes(b).Parse("mailhog.example")
	m.ID = data.MessageID(id)
	return m, nil
}
<|endoftext|>


<commit_before>package rules

import (
	"regexp"
	"strings"
)

type keyMatcher interface {
	getPrefix() string
	match(string) (keyMatch, bool)
	getPattern() string
}

type regexKeyMatcher struct {
	regex    *regexp.Regexp
	fieldMap map[string]int
	pattern  string
}

type keyMatch interface {
	GetAttribute(name string) *string
	Format(pattern string) string
}

func (rkm *regexKeyMatcher) getPrefix() string {
	end := strings.Index(rkm.pattern, ":")
	if end == -1 {
		end = len(rkm.pattern)
	}
	return rkm.pattern[0:end]
}

func (rkm *regexKeyMatcher) getPattern() string {
	return rkm.pattern
}

type regexKeyMatch struct {
	matchStrings []string
	fieldMap     map[string]int
}

func newKeyMatch(path string, kmr *regexKeyMatcher) *regexKeyMatch {
	results := kmr.regex.FindStringSubmatch(path)
	if results == nil {
		return nil
	}
	km := &regexKeyMatch{
		matchStrings: results,
		fieldMap:     kmr.fieldMap,
	}
	return km
}

func (m *regexKeyMatch) GetAttribute(name string) *string {
	index, ok := m.fieldMap[name]
	if !ok {
		return nil
	}
	result := m.matchStrings[index]
	return &result
}

func (m *regexKeyMatch) Format(pattern string) string {
	return formatWithAttributes(pattern, m)
}
func formatWithAttributes(pattern string, m Attributes) string {
	paths := strings.Split(pattern, "/")
	colon := ":"[0]
	result := ""
	for _, path := range paths {
		if len(path) == 0 {
			continue
		}
		result = result + "/"
		if path[0] == colon {
			attr := m.GetAttribute(path[1:])
			if attr == nil {
				s := "XXX"
				attr = &s
			}
			result = result + *attr
		} else {
			result = result + path
		}
	}
	return result
}

// Keep the bool return value, because it's tricky to check for null
// references when dealing with interfaces
func (rkm *regexKeyMatcher) match(path string) (keyMatch, bool) {
	m := newKeyMatch(path, rkm)
	if m == nil {
		return nil, false
	}
	return m, true
}

func newRegexKeyMatcher(pattern string) (*regexKeyMatcher, error) {
	fields, regexString := parsePath(pattern)
	regex, err := regexp.Compile(regexString)
	if err != nil {
		return nil, err
	}
	return &regexKeyMatcher{
		regex:    regex,
		fieldMap: fields,
		pattern:  pattern,
	}, nil
}

func parsePath(pattern string) (map[string]int, string) {
	paths := strings.Split(pattern, "/")
	colon := ":"[0]
	regex := ""
	fields := make(map[string]int)
	fieldIndex := 1
	for _, path := range paths {
		if len(path) == 0 {
			continue
		}
		regex = regex + "/"
		if path[0] == colon {
			regex = regex + "([^\\/:]+)"
			fields[path[1:]] = fieldIndex
			fieldIndex++
		} else {
			regex = regex + path
		}
	}
	return fields, regex
}
<commit_msg>Fix matcher to use strings.HasPrefix to find leading colon<commit_after>package rules

import (
	"regexp"
	"strings"
)

type keyMatcher interface {
	getPrefix() string
	match(string) (keyMatch, bool)
	getPattern() string
}

type regexKeyMatcher struct {
	regex    *regexp.Regexp
	fieldMap map[string]int
	pattern  string
}

type keyMatch interface {
	GetAttribute(name string) *string
	Format(pattern string) string
}

func (rkm *regexKeyMatcher) getPrefix() string {
	end := strings.Index(rkm.pattern, ":")
	if end == -1 {
		end = len(rkm.pattern)
	}
	return rkm.pattern[0:end]
}

func (rkm *regexKeyMatcher) getPattern() string {
	return rkm.pattern
}

type regexKeyMatch struct {
	matchStrings []string
	fieldMap     map[string]int
}

func newKeyMatch(path string, kmr *regexKeyMatcher) *regexKeyMatch {
	results := kmr.regex.FindStringSubmatch(path)
	if results == nil {
		return nil
	}
	km := &regexKeyMatch{
		matchStrings: results,
		fieldMap:     kmr.fieldMap,
	}
	return km
}

func (m *regexKeyMatch) GetAttribute(name string) *string {
	index, ok := m.fieldMap[name]
	if !ok {
		return nil
	}
	result := m.matchStrings[index]
	return &result
}

func (m *regexKeyMatch) Format(pattern string) string {
	return formatWithAttributes(pattern, m)
}
func formatWithAttributes(pattern string, m Attributes) string {
	paths := strings.Split(pattern, "/")
	result := ""
	for _, path := range paths {
		if len(path) == 0 {
			continue
		}
		result = result + "/"
		if strings.HasPrefix(path, ":") {
			attr := m.GetAttribute(path[1:])
			if attr == nil {
				s := "XXX"
				attr = &s
			}
			result = result + *attr
		} else {
			result = result + path
		}
	}
	return result
}

// Keep the bool return value, because it's tricky to check for null
// references when dealing with interfaces
func (rkm *regexKeyMatcher) match(path string) (keyMatch, bool) {
	m := newKeyMatch(path, rkm)
	if m == nil {
		return nil, false
	}
	return m, true
}

func newRegexKeyMatcher(pattern string) (*regexKeyMatcher, error) {
	fields, regexString := parsePath(pattern)
	regex, err := regexp.Compile(regexString)
	if err != nil {
		return nil, err
	}
	return &regexKeyMatcher{
		regex:    regex,
		fieldMap: fields,
		pattern:  pattern,
	}, nil
}

func parsePath(pattern string) (map[string]int, string) {
	paths := strings.Split(pattern, "/")
	regex := ""
	fields := make(map[string]int)
	fieldIndex := 1
	for _, path := range paths {
		if len(path) == 0 {
			continue
		}
		regex = regex + "/"
		if strings.HasPrefix(path, ":") {
			regex = regex + "([^\\/:]+)"
			fields[path[1:]] = fieldIndex
			fieldIndex++
		} else {
			regex = regex + path
		}
	}
	return fields, regex
}
<|endoftext|>