<commit_before>package linux

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/containerd"
	"github.com/docker/containerd/api/services/shim"
	"github.com/docker/containerd/api/types/container"
	"github.com/docker/containerd/api/types/mount"
	"github.com/docker/containerd/log"

	"golang.org/x/net/context"
)

const (
	runtimeName    = "linux"
	configFilename = "config.json"
)

func init() {
	containerd.Register(runtimeName, &containerd.Registration{
		Type: containerd.RuntimePlugin,
		Init: New,
	})
}

func New(ic *containerd.InitContext) (interface{}, error) {
	path := filepath.Join(ic.State, runtimeName)
	if err := os.MkdirAll(path, 0700); err != nil {
		return nil, err
	}
	c, cancel := context.WithCancel(ic.Context)
	return &Runtime{
		root:          path,
		events:        make(chan *containerd.Event, 2048),
		eventsContext: c,
		eventsCancel:  cancel,
	}, nil
}

type Runtime struct {
	root string

	events        chan *containerd.Event
	eventsContext context.Context
	eventsCancel  func()
}

func (r *Runtime) Create(ctx context.Context, id string, opts containerd.CreateOpts) (containerd.Container, error) {
	path, err := r.newBundle(id, opts.Spec)
	if err != nil {
		return nil, err
	}
	s, err := newShim(path)
	if err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	if err := r.handleEvents(s); err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	sopts := &shim.CreateRequest{
		ID:       id,
		Bundle:   path,
		Runtime:  "runc",
		Stdin:    opts.IO.Stdin,
		Stdout:   opts.IO.Stdout,
		Stderr:   opts.IO.Stderr,
		Terminal: opts.IO.Terminal,
	}
	for _, m := range opts.Rootfs {
		sopts.Rootfs = append(sopts.Rootfs, &mount.Mount{
			Type:    m.Type,
			Source:  m.Source,
			Options: m.Options,
		})
	}
	if _, err := s.Create(ctx, sopts); err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	return &Container{
		id:   id,
		shim: s,
	}, nil
}

func (r *Runtime) Delete(ctx context.Context, c containerd.Container) error {
	lc, ok := c.(*Container)
	if !ok {
		return fmt.Errorf("container cannot be cast as *linux.Container")
	}
	if _, err := lc.shim.Delete(ctx, &shim.DeleteRequest{}); err != nil {
		return err
	}
	lc.shim.Exit(ctx, &shim.ExitRequest{})
	return r.deleteBundle(lc.id)
}

func (r *Runtime) Containers() ([]containerd.Container, error) {
	dir, err := ioutil.ReadDir(r.root)
	if err != nil {
		return nil, err
	}
	var o []containerd.Container
	for _, fi := range dir {
		if !fi.IsDir() {
			continue
		}
		c, err := r.loadContainer(filepath.Join(r.root, fi.Name()))
		if err != nil {
			return nil, err
		}
		o = append(o, c)
	}
	return o, nil
}

func (r *Runtime) Events(ctx context.Context) <-chan *containerd.Event {
	return r.events
}

func (r *Runtime) handleEvents(s shim.ShimClient) error {
	events, err := s.Events(r.eventsContext, &shim.EventsRequest{})
	if err != nil {
		return err
	}
	go r.forward(events)
	return nil
}

func (r *Runtime) forward(events shim.Shim_EventsClient) {
	for {
		e, err := events.Recv()
		if err != nil {
			log.G(r.eventsContext).WithError(err).Error("get event from shim")
			return
		}
		var et containerd.EventType
		switch e.Type {
		case container.Event_CREATE:
			et = containerd.CreateEvent
		case container.Event_EXEC_ADDED:
			et = containerd.ExecAddEvent
		case container.Event_EXIT:
			et = containerd.ExitEvent
		case container.Event_OOM:
			et = containerd.OOMEvent
		case container.Event_START:
			et = containerd.StartEvent
		}
		r.events <- &containerd.Event{
			Timestamp:  time.Now(),
			Runtime:    runtimeName,
			Type:       et,
			Pid:        e.Pid,
			ID:         e.ID,
			ExitStatus: e.ExitStatus,
		}
	}
}

func (r *Runtime) newBundle(id string, spec []byte) (string, error) {
	path := filepath.Join(r.root, id)
	if err := os.Mkdir(path, 0700); err != nil {
		return "", err
	}
	if err := os.Mkdir(filepath.Join(path, "rootfs"), 0700); err != nil {
		return "", err
	}
	f, err := os.Create(filepath.Join(path, configFilename))
	if err != nil {
		return "", err
	}
	_, err = io.Copy(f, bytes.NewReader(spec))
	return path, err
}

func (r *Runtime) deleteBundle(id string) error {
	return os.RemoveAll(filepath.Join(r.root, id))
}

func (r *Runtime) loadContainer(path string) (*Container, error) {
	id := filepath.Base(path)
	s, err := loadShim(path)
	if err != nil {
		return nil, err
	}
	return &Container{
		id:   id,
		shim: s,
	}, nil
}
<commit_msg>Do not leak open fd to config.json in newBundle<commit_after>package linux

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/docker/containerd"
	"github.com/docker/containerd/api/services/shim"
	"github.com/docker/containerd/api/types/container"
	"github.com/docker/containerd/api/types/mount"
	"github.com/docker/containerd/log"

	"golang.org/x/net/context"
)

const (
	runtimeName    = "linux"
	configFilename = "config.json"
)

func init() {
	containerd.Register(runtimeName, &containerd.Registration{
		Type: containerd.RuntimePlugin,
		Init: New,
	})
}

func New(ic *containerd.InitContext) (interface{}, error) {
	path := filepath.Join(ic.State, runtimeName)
	if err := os.MkdirAll(path, 0700); err != nil {
		return nil, err
	}
	c, cancel := context.WithCancel(ic.Context)
	return &Runtime{
		root:          path,
		events:        make(chan *containerd.Event, 2048),
		eventsContext: c,
		eventsCancel:  cancel,
	}, nil
}

type Runtime struct {
	root string

	events        chan *containerd.Event
	eventsContext context.Context
	eventsCancel  func()
}

func (r *Runtime) Create(ctx context.Context, id string, opts containerd.CreateOpts) (containerd.Container, error) {
	path, err := r.newBundle(id, opts.Spec)
	if err != nil {
		return nil, err
	}
	s, err := newShim(path)
	if err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	if err := r.handleEvents(s); err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	sopts := &shim.CreateRequest{
		ID:       id,
		Bundle:   path,
		Runtime:  "runc",
		Stdin:    opts.IO.Stdin,
		Stdout:   opts.IO.Stdout,
		Stderr:   opts.IO.Stderr,
		Terminal: opts.IO.Terminal,
	}
	for _, m := range opts.Rootfs {
		sopts.Rootfs = append(sopts.Rootfs, &mount.Mount{
			Type:    m.Type,
			Source:  m.Source,
			Options: m.Options,
		})
	}
	if _, err := s.Create(ctx, sopts); err != nil {
		os.RemoveAll(path)
		return nil, err
	}
	return &Container{
		id:   id,
		shim: s,
	}, nil
}

func (r *Runtime) Delete(ctx context.Context, c containerd.Container) error {
	lc, ok := c.(*Container)
	if !ok {
		return fmt.Errorf("container cannot be cast as *linux.Container")
	}
	if _, err := lc.shim.Delete(ctx, &shim.DeleteRequest{}); err != nil {
		return err
	}
	lc.shim.Exit(ctx, &shim.ExitRequest{})
	return r.deleteBundle(lc.id)
}

func (r *Runtime) Containers() ([]containerd.Container, error) {
	dir, err := ioutil.ReadDir(r.root)
	if err != nil {
		return nil, err
	}
	var o []containerd.Container
	for _, fi := range dir {
		if !fi.IsDir() {
			continue
		}
		c, err := r.loadContainer(filepath.Join(r.root, fi.Name()))
		if err != nil {
			return nil, err
		}
		o = append(o, c)
	}
	return o, nil
}

func (r *Runtime) Events(ctx context.Context) <-chan *containerd.Event {
	return r.events
}

func (r *Runtime) handleEvents(s shim.ShimClient) error {
	events, err := s.Events(r.eventsContext, &shim.EventsRequest{})
	if err != nil {
		return err
	}
	go r.forward(events)
	return nil
}

func (r *Runtime) forward(events shim.Shim_EventsClient) {
	for {
		e, err := events.Recv()
		if err != nil {
			log.G(r.eventsContext).WithError(err).Error("get event from shim")
			return
		}
		var et containerd.EventType
		switch e.Type {
		case container.Event_CREATE:
			et = containerd.CreateEvent
		case container.Event_EXEC_ADDED:
			et = containerd.ExecAddEvent
		case container.Event_EXIT:
			et = containerd.ExitEvent
		case container.Event_OOM:
			et = containerd.OOMEvent
		case container.Event_START:
			et = containerd.StartEvent
		}
		r.events <- &containerd.Event{
			Timestamp:  time.Now(),
			Runtime:    runtimeName,
			Type:       et,
			Pid:        e.Pid,
			ID:         e.ID,
			ExitStatus: e.ExitStatus,
		}
	}
}

func (r *Runtime) newBundle(id string, spec []byte) (string, error) {
	path := filepath.Join(r.root, id)
	if err := os.Mkdir(path, 0700); err != nil {
		return "", err
	}
	if err := os.Mkdir(filepath.Join(path, "rootfs"), 0700); err != nil {
		return "", err
	}
	f, err := os.Create(filepath.Join(path, configFilename))
	if err != nil {
		return "", err
	}
	defer f.Close()
	_, err = io.Copy(f, bytes.NewReader(spec))
	return path, err
}

func (r *Runtime) deleteBundle(id string) error {
	return os.RemoveAll(filepath.Join(r.root, id))
}

func (r *Runtime) loadContainer(path string) (*Container, error) {
	id := filepath.Base(path)
	s, err := loadShim(path)
	if err != nil {
		return nil, err
	}
	return &Container{
		id:   id,
		shim: s,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Fredrik Ehnbom
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package text

import (
	"sort"
	"sync"
)

// The RegionSet manages multiple regions,
// merging any regions overlapping.
//
// Note that regions that are right next to each other
// are not merged into a single region. This is because
// otherwise it would not be possible to have multiple
// cursors right next to each other.
type RegionSet struct {
	regions []Region
	lock    sync.Mutex
}

// Adjusts all the regions in the set
func (r *RegionSet) Adjust(position, delta int) {
	r.lock.Lock()
	defer r.lock.Unlock()
	for i := range r.regions {
		r.regions[i].Adjust(position, delta)
	}
	r.flush()
}

// TODO(q): There should be a on modified callback on the RegionSet
func (r *RegionSet) flush() {
	sort.Sort(r)
	for i := 1; i < len(r.regions); {
		if r.regions[i-1] == r.regions[i] || r.regions[i-1].Intersects(r.regions[i]) || r.regions[i].Covers(r.regions[i-1]) {
			r.regions[i-1] = r.regions[i-1].Cover(r.regions[i])
			copy(r.regions[i:], r.regions[i+1:])
			r.regions = r.regions[:len(r.regions)-1]
		} else {
			i++
		}
	}
}

// Removes the given region from the set
func (r *RegionSet) Substract(r2 Region) {
	r.Adjust(r2.Begin(), r2.Size())
}

// Adds the given region to the set
func (r *RegionSet) Add(r2 Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = append(r.regions, r2)
	r.flush()
}

// Clears the set
func (r *RegionSet) Clear() {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = r.regions[0:0]
	r.flush()
}

// Gets the region at index i
func (r *RegionSet) Get(i int) Region {
	r.lock.Lock()
	defer r.lock.Unlock()
	return r.regions[i]
}

// Compares two regions given by their indices
func (r *RegionSet) Less(i, j int) bool {
	if bi, bj := r.regions[i].Begin(), r.regions[j].Begin(); bi < bj {
		return true
	} else if bi == bj {
		return r.regions[i].End() < r.regions[j].End()
	}
	return false
}

// Swaps two regions ath the given indices
func (r *RegionSet) Swap(i, j int) {
	r.regions[i], r.regions[j] = r.regions[j], r.regions[i]
}

// Returns the number of regions contained in the set
func (r *RegionSet) Len() int {
	return len(r.regions)
}

// Adds all regions in the array to the set
func (r *RegionSet) AddAll(rs []Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = append(r.regions, rs...)
	r.flush()
}

// Returns whether the specified region is part of the set
func (r *RegionSet) Contains(r2 Region) bool {
	for i := range r.regions {
		if r.regions[i] == r2 || (r.regions[i].Contains(r2.Begin()) && r.regions[i].Contains(r2.End())) {
			return true
		}
	}
	return false
}

// Returns a copy of the regions in the set
func (r *RegionSet) Regions() (ret []Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	ret = make([]Region, len(r.regions))
	copy(ret, r.regions)
	return
}

// Returns whether the set contains at least one
// region that isn't empty.
func (r *RegionSet) HasNonEmpty() bool {
	r.lock.Lock()
	defer r.lock.Unlock()
	for _, r := range r.regions {
		if !r.Empty() {
			return true
		}
	}
	return false
}

// Opposite of #HasNonEmpty
func (r *RegionSet) HasEmpty() bool {
	r.lock.Lock()
	defer r.lock.Unlock()
	for _, r := range r.regions {
		if r.Empty() {
			return true
		}
	}
	return false
}

// Cuts away the provided region from the set, and returns
// the new set
func (r *RegionSet) Cut(r2 Region) (ret RegionSet) {
	for i := 0; i < len(r.regions); i++ {
		for _, xor := range r.regions[i].Cut(r2) {
			if !xor.Empty() {
				ret.Add(xor)
			}
		}
	}
	return
}
<commit_msg>text/RegionSet: Also lock Contains and Cut<commit_after>// Copyright 2013 Fredrik Ehnbom
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package text

import (
	"sort"
	"sync"
)

// The RegionSet manages multiple regions,
// merging any regions overlapping.
//
// Note that regions that are right next to each other
// are not merged into a single region. This is because
// otherwise it would not be possible to have multiple
// cursors right next to each other.
type RegionSet struct {
	regions []Region
	lock    sync.Mutex
}

// Adjusts all the regions in the set
func (r *RegionSet) Adjust(position, delta int) {
	r.lock.Lock()
	defer r.lock.Unlock()
	for i := range r.regions {
		r.regions[i].Adjust(position, delta)
	}
	r.flush()
}

// TODO(q): There should be a on modified callback on the RegionSet
func (r *RegionSet) flush() {
	sort.Sort(r)
	for i := 1; i < len(r.regions); {
		if r.regions[i-1] == r.regions[i] || r.regions[i-1].Intersects(r.regions[i]) || r.regions[i].Covers(r.regions[i-1]) {
			r.regions[i-1] = r.regions[i-1].Cover(r.regions[i])
			copy(r.regions[i:], r.regions[i+1:])
			r.regions = r.regions[:len(r.regions)-1]
		} else {
			i++
		}
	}
}

// Removes the given region from the set
func (r *RegionSet) Substract(r2 Region) {
	r.Adjust(r2.Begin(), r2.Size())
}

// Adds the given region to the set
func (r *RegionSet) Add(r2 Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = append(r.regions, r2)
	r.flush()
}

// Clears the set
func (r *RegionSet) Clear() {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = r.regions[0:0]
	r.flush()
}

// Gets the region at index i
func (r *RegionSet) Get(i int) Region {
	r.lock.Lock()
	defer r.lock.Unlock()
	return r.regions[i]
}

// Compares two regions given by their indices
func (r *RegionSet) Less(i, j int) bool {
	if bi, bj := r.regions[i].Begin(), r.regions[j].Begin(); bi < bj {
		return true
	} else if bi == bj {
		return r.regions[i].End() < r.regions[j].End()
	}
	return false
}

// Swaps two regions ath the given indices
func (r *RegionSet) Swap(i, j int) {
	r.regions[i], r.regions[j] = r.regions[j], r.regions[i]
}

// Returns the number of regions contained in the set
func (r *RegionSet) Len() int {
	return len(r.regions)
}

// Adds all regions in the array to the set
func (r *RegionSet) AddAll(rs []Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	r.regions = append(r.regions, rs...)
	r.flush()
}

// Returns whether the specified region is part of the set
func (r *RegionSet) Contains(r2 Region) bool {
	r.lock.Lock()
	defer r.lock.Unlock()

	for i := range r.regions {
		if r.regions[i] == r2 || (r.regions[i].Contains(r2.Begin()) && r.regions[i].Contains(r2.End())) {
			return true
		}
	}
	return false
}

// Returns a copy of the regions in the set
func (r *RegionSet) Regions() (ret []Region) {
	r.lock.Lock()
	defer r.lock.Unlock()
	ret = make([]Region, len(r.regions))
	copy(ret, r.regions)
	return
}

// Returns whether the set contains at least one
// region that isn't empty.
func (r *RegionSet) HasNonEmpty() bool {
	r.lock.Lock()
	defer r.lock.Unlock()
	for _, r := range r.regions {
		if !r.Empty() {
			return true
		}
	}
	return false
}

// Opposite of #HasNonEmpty
func (r *RegionSet) HasEmpty() bool {
	r.lock.Lock()
	defer r.lock.Unlock()
	for _, r := range r.regions {
		if r.Empty() {
			return true
		}
	}
	return false
}

// Cuts away the provided region from the set, and returns
// the new set
func (r *RegionSet) Cut(r2 Region) (ret RegionSet) {
	r.lock.Lock()
	defer r.lock.Unlock()

	for i := 0; i < len(r.regions); i++ {
		for _, xor := range r.regions[i].Cut(r2) {
			if !xor.Empty() {
				ret.Add(xor)
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package kinetic

import (
	"errors"
	"sync"
	"syscall"
	"testing"
	"time"

	. "github.com/smartystreets/goconvey/convey"
)

const testEndpoint = "http://127.0.0.1:4567"

func TestListenerStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			listener.interrupts <- syscall.SIGINT

			// Let it finish stopping
			<-time.After(3 * time.Second)

			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerError(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should handle errors successfully", func() {
			listener.errors <- errors.New("All your base are belong to us!")

			// Let the error propagate
			<-time.After(3 * time.Second)

			So(listener.getErrCount(), ShouldNotEqual, 0)
			So(listener.IsListening(), ShouldEqual, true)
		})
	})

	listener.Close()
}

func TestListenerMessage(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
		wg.Done()
	})

	<-time.After(3 * time.Second)

	for _, c := range cases {
		Convey("Given a running listener", t, func() {
			listener.addMessage(new(Message).Init(c.message, "test"))

			Convey("It should handle messages successfully", func() {
				So(listener.IsListening(), ShouldEqual, true)
				So(listener.Errors(), ShouldNotResemble, nil)
			})
		})
	}

	listener.Close()
}

func TestRetrieveMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
	producer, _ := new(Producer).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")

	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				if err != nil {
					t.Fatalf(err.Error())
				}

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}

	producer.Close()
	listener.Close()
}

var cases = []struct {
	message []byte
}{
	{
		message: []byte(`{"foo":"bar"}`),
	},
	{
		message: []byte(`{"bar":"baz"}`),
	},
	{
		message: []byte(`{"baz":"qux"}`),
	},
}
<commit_msg>Add a Gosched to TestListenerStop<commit_after>package kinetic

import (
	"errors"
	. "github.com/smartystreets/goconvey/convey"
	"runtime"
	"sync"
	"syscall"
	"testing"
	"time"
)

const testEndpoint = "http://127.0.0.1:4567"

func TestListenerStop(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should stop listening if sent an interrupt signal", func() {
			listener.interrupts <- syscall.SIGINT
			runtime.Gosched()
			// Let it finish stopping
			<-time.After(3 * time.Second)

			So(listener.IsListening(), ShouldEqual, false)
		})
	})

	listener.Close()
}

func TestListenerError(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	Convey("Given a running listener", t, func() {
		go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
			wg.Done()
		})

		Convey("It should handle errors successfully", func() {
			listener.errors <- errors.New("All your base are belong to us!")

			// Let the error propagate
			<-time.After(3 * time.Second)

			So(listener.getErrCount(), ShouldNotEqual, 0)
			So(listener.IsListening(), ShouldEqual, true)
		})
	})

	listener.Close()
}

func TestListenerMessage(t *testing.T) {
	listener, _ := new(Listener).Init()
	listener.NewEndpoint(testEndpoint, "stream-name")

	go listener.Listen(func(msg []byte, wg *sync.WaitGroup) {
		wg.Done()
	})

	<-time.After(3 * time.Second)

	for _, c := range cases {
		Convey("Given a running listener", t, func() {
			listener.addMessage(new(Message).Init(c.message, "test"))

			Convey("It should handle messages successfully", func() {
				So(listener.IsListening(), ShouldEqual, true)
				So(listener.Errors(), ShouldNotResemble, nil)
			})
		})
	}

	listener.Close()
}

func TestRetrieveMessage(t *testing.T) {
	listener, _ := new(Listener).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)
	producer, _ := new(Producer).InitC("your-stream", "0", ShardIterTypes[3], "accesskey", "secretkey", "us-east-1", 10)

	listener.NewEndpoint(testEndpoint, "your-stream")
	producer.NewEndpoint(testEndpoint, "your-stream")

	for _, c := range cases {
		Convey("Given a valid message", t, func() {
			producer.Send(new(Message).Init(c.message, "test"))

			Convey("It should be passed on the queue without error", func() {
				msg, err := listener.Retrieve()
				if err != nil {
					t.Fatalf(err.Error())
				}

				So(string(msg.Value()), ShouldResemble, string(c.message))
			})
		})
	}

	producer.Close()
	listener.Close()
}

var cases = []struct {
	message []byte
}{
	{
		message: []byte(`{"foo":"bar"}`),
	},
	{
		message: []byte(`{"bar":"baz"}`),
	},
	{
		message: []byte(`{"baz":"qux"}`),
	},
}
<|endoftext|>


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

// Package log provides logging utility.
//
// It abstracts the logger from the standard log package, allowing the
// developer to patck the logging target, changing this to a file, or syslog,
// for example.
package log

import (
	"io"
	"log"
	"sync"

	"github.com/tsuru/config"
)

type Logger interface {
	Error(string)
	Errorf(string, ...interface{})
	Fatal(string)
	Fatalf(string, ...interface{})
	Debug(string)
	Debugf(string, ...interface{})
	GetStdLogger() *log.Logger
}

func Init() {
	debug, err := config.GetBool("debug")
	if err != nil {
		debug = false
	}
	logFileName, err := config.GetString("log:file")
	var logger Logger
	if err != nil {
		logger = NewSyslogLogger("tsr", debug)
	} else {
		logger = NewFileLogger(logFileName, debug)
	}
	SetLogger(logger)
}

// Target is the current target for the log package.
type Target struct {
	logger Logger
	mut    sync.RWMutex
}

// SetLogger defines a new logger for the current target.
//
// See the builtin log package for more details.
func (t *Target) SetLogger(l Logger) {
	t.mut.Lock()
	defer t.mut.Unlock()
	t.logger = l
}

// Error writes the given values to the Target
// logger.
func (t *Target) Error(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Error(v)
	}
}

// Errorf writes the formatted string to the Target
// logger.
func (t *Target) Errorf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Errorf(format, v...)
	}
}

// Fatal writes the given values to the Target
// logger.
func (t *Target) Fatal(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Fatal(v)
	}
}

// Fatalf writes the formatted string to the Target
// logger.
func (t *Target) Fatalf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Fatalf(format, v...)
	}
}

// Debug writes the value to the Target
// logger.
func (t *Target) Debug(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Debug(v)
	}
}

// Debugf writes the formatted string to the Target
// logger.
func (t *Target) Debugf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Debugf(format, v...)
	}
}

// GetStdLogger returns a standard Logger instance
// useful for configuring log in external packages.
func (t *Target) GetStdLogger() *log.Logger {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		return t.logger.GetStdLogger()
	}
	return nil
}

var DefaultTarget = new(Target)

// Error is a wrapper for DefaultTarget.Error.
func Error(v string) {
	DefaultTarget.Error(v)
}

// Errorf is a wrapper for DefaultTarget.Errorf.
func Errorf(format string, v ...interface{}) {
	DefaultTarget.Errorf(format, v...)
}

// Fatal is a wrapper for DefaultTarget.Fatal.
func Fatal(v string) {
	DefaultTarget.Fatal(v)
}

// Fatalf is a wrapper for DefaultTarget.Errorf.
func Fatalf(format string, v ...interface{}) {
	DefaultTarget.Fatalf(format, v...)
}

// Debug is a wrapper for DefaultTarget.Debug.
func Debug(v string) {
	DefaultTarget.Debug(v)
}

// Debugf is a wrapper for DefaultTarget.Debugf.
func Debugf(format string, v ...interface{}) {
	DefaultTarget.Debugf(format, v...)
}

// GetStdLogger is a wrapper for DefaultTarget.GetStdLogger.
func GetStdLogger() *log.Logger {
	return DefaultTarget.GetStdLogger()
}

// SetLogger is a wrapper for DefaultTarget.SetLogger.
func SetLogger(logger Logger) {
	DefaultTarget.SetLogger(logger)
}

func Write(w io.Writer, content []byte) error {
	n, err := w.Write(content)
	if err != nil {
		return err
	}
	if n != len(content) {
		return io.ErrShortWrite
	}
	return nil
}
<commit_msg>Fix typo: patck => pick<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package log provides logging utility.
//
// It abstracts the logger from the standard log package, allowing the
// developer to pick the logging target, changing this to a file, or syslog,
// for example.
package log

import (
	"io"
	"log"
	"sync"

	"github.com/tsuru/config"
)

type Logger interface {
	Error(string)
	Errorf(string, ...interface{})
	Fatal(string)
	Fatalf(string, ...interface{})
	Debug(string)
	Debugf(string, ...interface{})
	GetStdLogger() *log.Logger
}

func Init() {
	debug, err := config.GetBool("debug")
	if err != nil {
		debug = false
	}
	logFileName, err := config.GetString("log:file")
	var logger Logger
	if err != nil {
		logger = NewSyslogLogger("tsr", debug)
	} else {
		logger = NewFileLogger(logFileName, debug)
	}
	SetLogger(logger)
}

// Target is the current target for the log package.
type Target struct {
	logger Logger
	mut    sync.RWMutex
}

// SetLogger defines a new logger for the current target.
//
// See the builtin log package for more details.
func (t *Target) SetLogger(l Logger) {
	t.mut.Lock()
	defer t.mut.Unlock()
	t.logger = l
}

// Error writes the given values to the Target
// logger.
func (t *Target) Error(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Error(v)
	}
}

// Errorf writes the formatted string to the Target
// logger.
func (t *Target) Errorf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Errorf(format, v...)
	}
}

// Fatal writes the given values to the Target
// logger.
func (t *Target) Fatal(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Fatal(v)
	}
}

// Fatalf writes the formatted string to the Target
// logger.
func (t *Target) Fatalf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Fatalf(format, v...)
	}
}

// Debug writes the value to the Target
// logger.
func (t *Target) Debug(v string) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Debug(v)
	}
}

// Debugf writes the formatted string to the Target
// logger.
func (t *Target) Debugf(format string, v ...interface{}) {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		t.logger.Debugf(format, v...)
	}
}

// GetStdLogger returns a standard Logger instance
// useful for configuring log in external packages.
func (t *Target) GetStdLogger() *log.Logger {
	t.mut.RLock()
	defer t.mut.RUnlock()
	if t.logger != nil {
		return t.logger.GetStdLogger()
	}
	return nil
}

var DefaultTarget = new(Target)

// Error is a wrapper for DefaultTarget.Error.
func Error(v string) {
	DefaultTarget.Error(v)
}

// Errorf is a wrapper for DefaultTarget.Errorf.
func Errorf(format string, v ...interface{}) {
	DefaultTarget.Errorf(format, v...)
}

// Fatal is a wrapper for DefaultTarget.Fatal.
func Fatal(v string) {
	DefaultTarget.Fatal(v)
}

// Fatalf is a wrapper for DefaultTarget.Errorf.
func Fatalf(format string, v ...interface{}) {
	DefaultTarget.Fatalf(format, v...)
}

// Debug is a wrapper for DefaultTarget.Debug.
func Debug(v string) {
	DefaultTarget.Debug(v)
}

// Debugf is a wrapper for DefaultTarget.Debugf.
func Debugf(format string, v ...interface{}) {
	DefaultTarget.Debugf(format, v...)
}

// GetStdLogger is a wrapper for DefaultTarget.GetStdLogger.
func GetStdLogger() *log.Logger {
	return DefaultTarget.GetStdLogger()
}

// SetLogger is a wrapper for DefaultTarget.SetLogger.
func SetLogger(logger Logger) {
	DefaultTarget.SetLogger(logger)
}

func Write(w io.Writer, content []byte) error {
	n, err := w.Write(content)
	if err != nil {
		return err
	}
	if n != len(content) {
		return io.ErrShortWrite
	}
	return nil
}
<|endoftext|>


<commit_before>package log

import (
	"fmt"

	"github.com/davecgh/go-spew/spew"
	"github.com/eyecuelab/kit/goenv"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

var (
	//Print is an alias for log.Print in logrus
	Print = log.Print

	//Fatalf is an alias for log.Fatalf in logrus
	Fatalf = log.Fatalf

	//Fatal is an alias for log.Fatal in logrus
	Fatal = log.Fatal

	//Println is an alias for log.Println in logrus
	Println = log.Println

	//Info is an alias for log.Info in logrus
	Info = log.Info

	//Warn is an alias for log.Warn in logrus
	Warn = log.Warn

	//Warnf is an alias for log.Warnf in logrus
	Warnf = log.Warnf
)

func FatalWrap(err error, msg string) {
	log.Fatalf("%+v", errors.Wrap(err, msg))
}

func Infof(format string, args ...interface{}) {
	s := spew.Sprintf(format, args...)
	log.Infof(s)
}

func Infofln(format string, args ...interface{}) {

	Infof(format, args...)
}

func ErrorWrap(err error, msg string) {
	log.Errorf("%+v", errors.Wrap(err, msg))
}

//Check calls Fatal on an error if it is non-nil.
func Check(err error) {
	if err != nil {
		Fatal(err)
	}
}

var Logger = log.StandardLogger()

func init() {
	if goenv.Prod {
		log.SetFormatter(&log.JSONFormatter{})
	}
}
<commit_msg>removed dead import<commit_after>package log

import (
	"github.com/davecgh/go-spew/spew"
	"github.com/eyecuelab/kit/goenv"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

var (
	//Print is an alias for log.Print in logrus
	Print = log.Print

	//Fatalf is an alias for log.Fatalf in logrus
	Fatalf = log.Fatalf

	//Fatal is an alias for log.Fatal in logrus
	Fatal = log.Fatal

	//Println is an alias for log.Println in logrus
	Println = log.Println

	//Info is an alias for log.Info in logrus
	Info = log.Info

	//Warn is an alias for log.Warn in logrus
	Warn = log.Warn

	//Warnf is an alias for log.Warnf in logrus
	Warnf = log.Warnf
)

func FatalWrap(err error, msg string) {
	log.Fatalf("%+v", errors.Wrap(err, msg))
}

func Infof(format string, args ...interface{}) {
	s := spew.Sprintf(format, args...)
	log.Infof(s)
}

func Infofln(format string, args ...interface{}) {

	Infof(format, args...)
}

func ErrorWrap(err error, msg string) {
	log.Errorf("%+v", errors.Wrap(err, msg))
}

//Check calls Fatal on an error if it is non-nil.
func Check(err error) {
	if err != nil {
		Fatal(err)
	}
}

var Logger = log.StandardLogger()

func init() {
	if goenv.Prod {
		log.SetFormatter(&log.JSONFormatter{})
	}
}
<|endoftext|>


<commit_before>package logrus

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"time"
)

const (
	nocolor = 0
	red     = 31
	green   = 32
	yellow  = 33
	blue    = 34
)

var (
	baseTimestamp time.Time
	isTerminal    bool
)

func init() {
	baseTimestamp = time.Now()
	isTerminal = IsTerminal()
}

func miniTS() int {
	return int(time.Since(baseTimestamp) / time.Second)
}

type TextFormatter struct {
	// Set to true to bypass checking for a TTY before outputting colors.
	ForceColors   bool
	DisableColors bool
}

func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {

	var keys []string
	for k := range entry.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	b := &bytes.Buffer{}

	prefixFieldClashes(entry)

	isColored := (f.ForceColors || isTerminal) && !f.DisableColors

	if isColored {
		printColored(b, entry, keys)
	} else {
		f.AppendKeyValue(b, "time", entry.Time.Format(time.RFC3339))
		f.AppendKeyValue(b, "level", entry.Level.String())
		f.AppendKeyValue(b, "msg", entry.Message)
		for _, key := range keys {
			f.AppendKeyValue(b, key, entry.Data[key])
		}
	}

	b.WriteByte('\n')
	return b.Bytes(), nil
}

func printColored(b *bytes.Buffer, entry *Entry, keys []string) {
	var levelColor int
	switch entry.Level {
	case WarnLevel:
		levelColor = yellow
	case ErrorLevel, FatalLevel, PanicLevel:
		levelColor = red
	default:
		levelColor = blue
	}

	levelText := strings.ToUpper(entry.Level.String())[0:4]

	fmt.Fprintf(b, "\x1b[%dm%s\x1b[0m[%04d] %-44s ", levelColor, levelText, miniTS(), entry.Message)
	for _, k := range keys {
		v := entry.Data[k]
		fmt.Fprintf(b, " \x1b[%dm%s\x1b[0m=%v", levelColor, k, v)
	}
}

func (f *TextFormatter) AppendKeyValue(b *bytes.Buffer, key, value interface{}) {
	if _, ok := value.(string); ok {
		fmt.Fprintf(b, "%v=%q ", key, value)
	} else {
		fmt.Fprintf(b, "%v=%v ", key, value)
	}
}
<commit_msg>textformatter: errors should be in quotes like strings are<commit_after>package logrus

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"time"
)

const (
	nocolor = 0
	red     = 31
	green   = 32
	yellow  = 33
	blue    = 34
)

var (
	baseTimestamp time.Time
	isTerminal    bool
)

func init() {
	baseTimestamp = time.Now()
	isTerminal = IsTerminal()
}

func miniTS() int {
	return int(time.Since(baseTimestamp) / time.Second)
}

type TextFormatter struct {
	// Set to true to bypass checking for a TTY before outputting colors.
	ForceColors   bool
	DisableColors bool
}

func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {

	var keys []string
	for k := range entry.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	b := &bytes.Buffer{}

	prefixFieldClashes(entry)

	isColored := (f.ForceColors || isTerminal) && !f.DisableColors

	if isColored {
		printColored(b, entry, keys)
	} else {
		f.appendKeyValue(b, "time", entry.Time.Format(time.RFC3339))
		f.appendKeyValue(b, "level", entry.Level.String())
		f.appendKeyValue(b, "msg", entry.Message)
		for _, key := range keys {
			f.appendKeyValue(b, key, entry.Data[key])
		}
	}

	b.WriteByte('\n')
	return b.Bytes(), nil
}

func printColored(b *bytes.Buffer, entry *Entry, keys []string) {
	var levelColor int
	switch entry.Level {
	case WarnLevel:
		levelColor = yellow
	case ErrorLevel, FatalLevel, PanicLevel:
		levelColor = red
	default:
		levelColor = blue
	}

	levelText := strings.ToUpper(entry.Level.String())[0:4]

	fmt.Fprintf(b, "\x1b[%dm%s\x1b[0m[%04d] %-44s ", levelColor, levelText, miniTS(), entry.Message)
	for _, k := range keys {
		v := entry.Data[k]
		fmt.Fprintf(b, " \x1b[%dm%s\x1b[0m=%v", levelColor, k, v)
	}
}

func (f *TextFormatter) appendKeyValue(b *bytes.Buffer, key, value interface{}) {
	switch value.(type) {
	case string, error:
		fmt.Fprintf(b, "%v=%q ", key, value)
	default:
		fmt.Fprintf(b, "%v=%v ", key, value)
	}
}
<|endoftext|>


<commit_before>package logrus

import (
	"bytes"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"time"
)

const (
	nocolor = 0
	red     = 31
	green   = 32
	yellow  = 33
	blue    = 34
	gray    = 37
)

var (
	baseTimestamp time.Time
	isTerminal    bool
	noQuoteNeeded *regexp.Regexp
)

func init() {
	baseTimestamp = time.Now()
	isTerminal = IsTerminal()
}

func miniTS() int {
	return int(time.Since(baseTimestamp) / time.Second)
}

type TextFormatter struct {
	// Set to true to bypass checking for a TTY before outputting colors.
	ForceColors   bool
	DisableColors bool
	// Set to true to disable timestamp logging (useful when the output
	// is redirected to a logging system already adding a timestamp)
	DisableTimestamp bool
}

func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {

	var keys []string = make([]string, 0, len(entry.Data))
	for k := range entry.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	b := &bytes.Buffer{}

	prefixFieldClashes(entry.Data)

	isColored := (f.ForceColors || isTerminal) && !f.DisableColors

	if isColored {
		printColored(b, entry, keys)
	} else {
		if !f.DisableTimestamp {
			f.appendKeyValue(b, "time", entry.Time.Format(time.RFC3339))
		}
		f.appendKeyValue(b, "level", entry.Level.String())
		f.appendKeyValue(b, "msg", entry.Message)
		for _, key := range keys {
			f.appendKeyValue(b, key, entry.Data[key])
		}
	}

	b.WriteByte('\n')
	return b.Bytes(), nil
}

func printColored(b *bytes.Buffer, entry *Entry, keys []string) {
	var levelColor int
	switch entry.Level {
	case DebugLevel:
		levelColor = gray
	case WarnLevel:
		levelColor = yellow
	case ErrorLevel, FatalLevel, PanicLevel:
		levelColor = red
	default:
		levelColor = blue
	}

	levelText := strings.ToUpper(entry.Level.String())[0:4]

	fmt.Fprintf(b, "\x1b[%dm%s\x1b[0m[%04d] %-44s ", levelColor, levelText, miniTS(), entry.Message)
	for _, k := range keys {
		v := entry.Data[k]
		fmt.Fprintf(b, " \x1b[%dm%s\x1b[0m=%v", levelColor, k, v)
	}
}

func needsQuoting(text string) bool {
	for _, ch := range text {
		if !((ch >= 'a' && ch <= 'z') ||
			(ch >= 'A' && ch <= 'Z') ||
			(ch >= '0' && ch < '9') ||
			ch == '-' || ch == '.') {
			return false
		}
	}
	return true
}

func (f *TextFormatter) appendKeyValue(b *bytes.Buffer, key, value interface{}) {
	switch value.(type) {
	case string:
		if needsQuoting(value.(string)) {
			fmt.Fprintf(b, "%v=%s ", key, value)
		} else {
			fmt.Fprintf(b, "%v=%q ", key, value)
		}
	case error:
		if needsQuoting(value.(error).Error()) {
			fmt.Fprintf(b, "%v=%s ", key, value)
		} else {
			fmt.Fprintf(b, "%v=%q ", key, value)
		}
	default:
		fmt.Fprintf(b, "%v=%v ", key, value)
	}
}
<commit_msg>Add option to show full timestamp in TextFormatter<commit_after>package logrus

import (
	"bytes"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"time"
)

const (
	nocolor = 0
	red     = 31
	green   = 32
	yellow  = 33
	blue    = 34
	gray    = 37
)

var (
	baseTimestamp time.Time
	isTerminal    bool
	noQuoteNeeded *regexp.Regexp
)

func init() {
	baseTimestamp = time.Now()
	isTerminal = IsTerminal()
}

func miniTS() int {
	return int(time.Since(baseTimestamp) / time.Second)
}

type TextFormatter struct {
	// Set to true to bypass checking for a TTY before outputting colors.
	ForceColors   bool
	DisableColors bool
	// Set to true to disable timestamp logging (useful when the output
	// is redirected to a logging system already adding a timestamp)
	DisableTimestamp bool
	FullTimestamp    bool
}

func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {

	var keys []string = make([]string, 0, len(entry.Data))
	for k := range entry.Data {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	b := &bytes.Buffer{}

	prefixFieldClashes(entry.Data)

	isColored := (f.ForceColors || isTerminal) && !f.DisableColors

	if isColored {
		f.printColored(b, entry, keys)
	} else {
		if !f.DisableTimestamp {
			f.appendKeyValue(b, "time", entry.Time.Format(time.RFC3339))
		}
		f.appendKeyValue(b, "level", entry.Level.String())
		f.appendKeyValue(b, "msg", entry.Message)
		for _, key := range keys {
			f.appendKeyValue(b, key, entry.Data[key])
		}
	}

	b.WriteByte('\n')
	return b.Bytes(), nil
}

func (f *TextFormatter) printColored(b *bytes.Buffer, entry *Entry, keys []string) {
	var levelColor int
	switch entry.Level {
	case DebugLevel:
		levelColor = gray
	case WarnLevel:
		levelColor = yellow
	case ErrorLevel, FatalLevel, PanicLevel:
		levelColor = red
	default:
		levelColor = blue
	}

	levelText := strings.ToUpper(entry.Level.String())[0:4]

	if !f.FullTimestamp {
		fmt.Fprintf(b, "\x1b[%dm%s\x1b[0m[%04d] %-44s ", levelColor, levelText, miniTS(), entry.Message)
	} else {
		fmt.Fprintf(b, "\x1b[%dm%s\x1b[0m[%s] %-44s ", levelColor, levelText, entry.Time.Format(time.RFC3339), entry.Message)
	}
	for _, k := range keys {
		v := entry.Data[k]
		fmt.Fprintf(b, " \x1b[%dm%s\x1b[0m=%v", levelColor, k, v)
	}
}

func needsQuoting(text string) bool {
	for _, ch := range text {
		if !((ch >= 'a' && ch <= 'z') ||
			(ch >= 'A' && ch <= 'Z') ||
			(ch >= '0' && ch < '9') ||
			ch == '-' || ch == '.') {
			return false
		}
	}
	return true
}

func (f *TextFormatter) appendKeyValue(b *bytes.Buffer, key, value interface{}) {
	switch value.(type) {
	case string:
		if needsQuoting(value.(string)) {
			fmt.Fprintf(b, "%v=%s ", key, value)
		} else {
			fmt.Fprintf(b, "%v=%q ", key, value)
		}
	case error:
		if needsQuoting(value.(error).Error()) {
			fmt.Fprintf(b, "%v=%s ", key, value)
		} else {
			fmt.Fprintf(b, "%v=%q ", key, value)
		}
	default:
		fmt.Fprintf(b, "%v=%v ", key, value)
	}
}
<|endoftext|>


<commit_before>package llrb

type ItemIterator func(i Item) bool

//func (t *Tree) Ascend(iterator ItemIterator) {
//	t.AscendGreaterOrEqual(Inf(-1), iterator)
//}

func (t *LLRB) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
	t.ascendRange(t.root, greaterOrEqual, lessThan, iterator)
}

func (t *LLRB) ascendRange(h *Node, inf, sup Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !less(h.Item, sup) {
		return t.ascendRange(h.Left, inf, sup, iterator)
	}
	if less(h.Item, inf) {
		return t.ascendRange(h.Right, inf, sup, iterator)
	}

	if !t.ascendRange(h.Left, inf, sup, iterator) {
		return false
	}
	if !iterator(h.Item) {
		return false
	}
	return t.ascendRange(h.Right, inf, sup, iterator)
}

// AscendGreaterOrEqual will call iterator once for each element greater or equal to
// pivot in ascending order. It will stop whenever the iterator returns false.
func (t *LLRB) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
	t.ascendGreaterOrEqual(t.root, pivot, iterator)
}

func (t *LLRB) ascendGreaterOrEqual(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !less(h.Item, pivot) {
		if !t.ascendGreaterOrEqual(h.Left, pivot, iterator) {
			return false
		}
		if !iterator(h.Item) {
			return false
		}
	}
	return t.ascendGreaterOrEqual(h.Right, pivot, iterator)
}

func (t *LLRB) AscendLessThan(pivot Item, iterator ItemIterator) {
	t.ascendLessThan(t.root, pivot, iterator)
}

func (t *LLRB) ascendLessThan(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !t.ascendLessThan(h.Left, pivot, iterator) {
		return false
	}
	if !iterator(h.Item) {
		return false
	}
	if less(h.Item, pivot) {
		return t.ascendLessThan(h.Left, pivot, iterator)
	}
	return true
}

// DescendLessOrEqual will call iterator once for each element less than the
// pivot in descending order. It will stop whenever the iterator returns false.
func (t *LLRB) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
	t.descendLessOrEqual(t.root, pivot, iterator)
}

func (t *LLRB) descendLessOrEqual(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if less(h.Item, pivot) || !less(pivot, h.Item) {
		if !t.descendLessOrEqual(h.Right, pivot, iterator) {
			return false
		}
		if !iterator(h.Item) {
			return false
		}
	}
	return t.descendLessOrEqual(h.Left, pivot, iterator)
}
<commit_msg>Add LLRB.Ascend() method + some refactoring<commit_after>package llrb

type ItemIterator func(i Item) bool

func (t *LLRB) Ascend(iterator ItemIterator) {
	ascend(t.root, iterator)
}

func ascend(h *Node, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !ascend(h.Left, iterator) {
		return false
	}
	if !iterator(h.Item) {
		return false
	}
	return ascend(h.Right, iterator)
}

func (t *LLRB) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
	ascendRange(t.root, greaterOrEqual, lessThan, iterator)
}

func ascendRange(h *Node, inf, sup Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !less(h.Item, sup) {
		return ascendRange(h.Left, inf, sup, iterator)
	}
	if less(h.Item, inf) {
		return ascendRange(h.Right, inf, sup, iterator)
	}

	if !ascendRange(h.Left, inf, sup, iterator) {
		return false
	}
	if !iterator(h.Item) {
		return false
	}
	return ascendRange(h.Right, inf, sup, iterator)
}

// AscendGreaterOrEqual will call iterator once for each element greater or equal to
// pivot in ascending order. It will stop whenever the iterator returns false.
func (t *LLRB) AscendGreaterOrEqual(pivot Item, iterator ItemIterator) {
	ascendGreaterOrEqual(t.root, pivot, iterator)
}

func ascendGreaterOrEqual(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !less(h.Item, pivot) {
		if !ascendGreaterOrEqual(h.Left, pivot, iterator) {
			return false
		}
		if !iterator(h.Item) {
			return false
		}
	}
	return ascendGreaterOrEqual(h.Right, pivot, iterator)
}

func (t *LLRB) AscendLessThan(pivot Item, iterator ItemIterator) {
	ascendLessThan(t.root, pivot, iterator)
}

func ascendLessThan(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if !ascendLessThan(h.Left, pivot, iterator) {
		return false
	}
	if !iterator(h.Item) {
		return false
	}
	if less(h.Item, pivot) {
		return ascendLessThan(h.Left, pivot, iterator)
	}
	return true
}

// DescendLessOrEqual will call iterator once for each element less than the
// pivot in descending order. It will stop whenever the iterator returns false.
func (t *LLRB) DescendLessOrEqual(pivot Item, iterator ItemIterator) {
	descendLessOrEqual(t.root, pivot, iterator)
}

func descendLessOrEqual(h *Node, pivot Item, iterator ItemIterator) bool {
	if h == nil {
		return true
	}
	if less(h.Item, pivot) || !less(pivot, h.Item) {
		if !descendLessOrEqual(h.Right, pivot, iterator) {
			return false
		}
		if !iterator(h.Item) {
			return false
		}
	}
	return descendLessOrEqual(h.Left, pivot, iterator)
}
<|endoftext|>


<commit_before>package lncfg

import (
	"context"
	"crypto/tls"
	"encoding/hex"
	"fmt"
	"net"
	"strconv"
	"strings"

	"github.com/btcsuite/btcd/btcec"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/tor"
)

var (
	loopBackAddrs = []string{"localhost", "127.0.0.1", "[::1]"}
)

// TCPResolver is a function signature that resolves an address on a given
// network.
type TCPResolver = func(network, addr string) (*net.TCPAddr, error)

// NormalizeAddresses returns a new slice with all the passed addresses
// normalized with the given default port and all duplicates removed.
func NormalizeAddresses(addrs []string, defaultPort string,
	tcpResolver TCPResolver) ([]net.Addr, error) {

	result := make([]net.Addr, 0, len(addrs))
	seen := map[string]struct{}{}

	for _, addr := range addrs {
		parsedAddr, err := ParseAddressString(
			addr, defaultPort, tcpResolver,
		)
		if err != nil {
			return nil, err
		}

		if _, ok := seen[parsedAddr.String()]; !ok {
			result = append(result, parsedAddr)
			seen[parsedAddr.String()] = struct{}{}
		}
	}

	return result, nil
}

// EnforceSafeAuthentication enforces "safe" authentication taking into account
// the interfaces that the RPC servers are listening on, and if macaroons and
// TLS is activated or not. To protect users from using dangerous config
// combinations, we'll prevent disabling authentication if the server is
// listening on a public interface.
func EnforceSafeAuthentication(addrs []net.Addr, macaroonsActive,
	tlsActive bool) error {

	// We'll now examine all addresses that this RPC server is listening
	// on. If it's a localhost address or a private address, we'll skip it,
	// otherwise, we'll return an error if macaroons are inactive.
	for _, addr := range addrs {
		if IsLoopback(addr.String()) || IsUnix(addr) || IsPrivate(addr) {
			continue
		}

		if !macaroonsActive {
			return fmt.Errorf("detected RPC server listening on "+
				"publicly reachable interface %v with "+
				"authentication disabled! Refusing to start "+
				"with --no-macaroons specified", addr)
		}

		if !tlsActive {
			return fmt.Errorf("detected RPC server listening on "+
				"publicly reachable interface %v with "+
				"encryption disabled! Refusing to start "+
				"with --notls specified", addr)
		}
	}

	return nil
}

// parseNetwork parses the network type of the given address.
func parseNetwork(addr net.Addr) string {
	switch addr := addr.(type) {
	// TCP addresses resolved through net.ResolveTCPAddr give a default
	// network of "tcp", so we'll map back the correct network for the given
	// address. This ensures that we can listen on the correct interface
	// (IPv4 vs IPv6).
	case *net.TCPAddr:
		if addr.IP.To4() != nil {
			return "tcp4"
		}
		return "tcp6"

	default:
		return addr.Network()
	}
}

// ListenOnAddress creates a listener that listens on the given address.
func ListenOnAddress(addr net.Addr) (net.Listener, error) {
	return net.Listen(parseNetwork(addr), addr.String())
}

// TLSListenOnAddress creates a TLS listener that listens on the given address.
func TLSListenOnAddress(addr net.Addr,
	config *tls.Config) (net.Listener, error) {
	return tls.Listen(parseNetwork(addr), addr.String(), config)
}

// IsLoopback returns true if an address describes a loopback interface.
func IsLoopback(addr string) bool {
	for _, loopback := range loopBackAddrs {
		if strings.Contains(addr, loopback) {
			return true
		}
	}

	return false
}

// isIPv6Host returns true if the host is IPV6 and false otherwise.
func isIPv6Host(host string) bool {
	v6Addr := net.ParseIP(host)
	if v6Addr == nil {
		return false
	}

	// The documentation states that if the IP address is an IPv6 address,
	// then To4() will return nil.
	return v6Addr.To4() == nil
}

// IsUnix returns true if an address describes an Unix socket address.
func IsUnix(addr net.Addr) bool {
	return strings.HasPrefix(addr.Network(), "unix")
}

// IsPrivate returns true if the address is private. The definitions are,
//   https://en.wikipedia.org/wiki/Link-local_address
//   https://en.wikipedia.org/wiki/Multicast_address
//   Local IPv4 addresses, https://tools.ietf.org/html/rfc1918
//   Local IPv6 addresses, https://tools.ietf.org/html/rfc4193
func IsPrivate(addr net.Addr) bool {
	switch addr := addr.(type) {
	case *net.TCPAddr:
		// Check 169.254.0.0/16 and fe80::/10.
		if addr.IP.IsLinkLocalUnicast() {
			return true
		}

		// Check 224.0.0.0/4 and ff00::/8.
		if addr.IP.IsLinkLocalMulticast() {
			return true
		}

		// Check 10.0.0.0/8, 172.16.0.0/12 and 192.168.0.0/16.
		if ip4 := addr.IP.To4(); ip4 != nil {
			return ip4[0] == 10 ||
				(ip4[0] == 172 && ip4[1]&0xf0 == 16) ||
				(ip4[0] == 192 && ip4[1] == 168)
		}

		// Check fc00::/7.
		return len(addr.IP) == net.IPv6len && addr.IP[0]&0xfe == 0xfc

	default:
		return false
	}
}

// ParseAddressString converts an address in string format to a net.Addr that is
// compatible with lnd. UDP is not supported because lnd needs reliable
// connections. We accept a custom function to resolve any TCP addresses so
// that caller is able control exactly how resolution is performed.
func ParseAddressString(strAddress string, defaultPort string,
	tcpResolver TCPResolver) (net.Addr, error) {

	var parsedNetwork, parsedAddr string

	// Addresses can either be in network://address:port format,
	// network:address:port, address:port, or just port. We want to support
	// all possible types.
	if strings.Contains(strAddress, "://") {
		parts := strings.Split(strAddress, "://")
		parsedNetwork, parsedAddr = parts[0], parts[1]
	} else if strings.Contains(strAddress, ":") {
		parts := strings.Split(strAddress, ":")
		parsedNetwork = parts[0]
		parsedAddr = strings.Join(parts[1:], ":")
	}

	// Only TCP and Unix socket addresses are valid. We can't use IP or
	// UDP only connections for anything we do in lnd.
	switch parsedNetwork {
	case "unix", "unixpacket":
		return net.ResolveUnixAddr(parsedNetwork, parsedAddr)

	case "tcp", "tcp4", "tcp6":
		return tcpResolver(
			parsedNetwork, verifyPort(parsedAddr, defaultPort),
		)

	case "ip", "ip4", "ip6", "udp", "udp4", "udp6", "unixgram":
		return nil, fmt.Errorf("only TCP or unix socket "+
			"addresses are supported: %s", parsedAddr)

	default:
		// We'll now possibly apply the default port, use the local
		// host short circuit, or parse out an all interfaces listen.
		addrWithPort := verifyPort(strAddress, defaultPort)
		rawHost, rawPort, _ := net.SplitHostPort(addrWithPort)

		// If we reach this point, then we'll check to see if we have
		// an onion addresses, if so, we can directly pass the raw
		// address and port to create the proper address.
		if tor.IsOnionHost(rawHost) {
			portNum, err := strconv.Atoi(rawPort)
			if err != nil {
				return nil, err
			}

			return &tor.OnionAddr{
				OnionService: rawHost,
				Port:         portNum,
			}, nil
		}

		// Otherwise, we'll attempt the resolve the host. The Tor
		// resolver is unable to resolve local or IPv6 addresses, so
		// we'll use the system resolver instead.
		if rawHost == "" || IsLoopback(rawHost) ||
			isIPv6Host(rawHost) {
			return net.ResolveTCPAddr("tcp", addrWithPort)
		}

		return tcpResolver("tcp", addrWithPort)
	}
}

// ParseLNAddressString converts a string of the form <pubkey>@<addr> into an
// lnwire.NetAddress. The <pubkey> must be presented in hex, and result in a
// 33-byte, compressed public key that lies on the secp256k1 curve. The <addr>
// may be any address supported by ParseAddressString. If no port is specified,
// the defaultPort will be used. Any tcp addresses that need resolving will be
// resolved using the custom TCPResolver.
func ParseLNAddressString(strAddress string, defaultPort string,
	tcpResolver TCPResolver) (*lnwire.NetAddress, error) {

	// Split the address string around the @ sign.
	parts := strings.Split(strAddress, "@")

	// The string is malformed if there are not exactly two parts.
	if len(parts) != 2 {
		return nil, fmt.Errorf("invalid lightning address %s: "+
			"must be of the form <pubkey-hex>@<addr>", strAddress)
	}

	// Now, take the first portion as the hex pubkey, and the latter as the
	// address string.
	parsedPubKey, parsedAddr := parts[0], parts[1]

	// Decode the hex pubkey to get the raw compressed pubkey bytes.
	pubKeyBytes, err := hex.DecodeString(parsedPubKey)
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address pubkey: %v", err)
	}

	// The compressed pubkey should have a length of exactly 33 bytes.
	if len(pubKeyBytes) != 33 {
		return nil, fmt.Errorf("invalid lightning address pubkey: "+
			"length must be 33 bytes, found %d", len(pubKeyBytes))
	}

	// Parse the pubkey bytes to verify that it corresponds to valid public
	// key on the secp256k1 curve.
	pubKey, err := btcec.ParsePubKey(pubKeyBytes, btcec.S256())
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address pubkey: %v", err)
	}

	// Finally, parse the address string using our generic address parser.
	addr, err := ParseAddressString(parsedAddr, defaultPort, tcpResolver)
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address address: %v", err)
	}

	return &lnwire.NetAddress{
		IdentityKey: pubKey,
		Address:     addr,
	}, nil
}

// verifyPort makes sure that an address string has both a host and a port. If
// there is no port found, the default port is appended. If the address is just
// a port, then we'll assume that the user is using the short cut to specify a
// localhost:port address.
func verifyPort(address string, defaultPort string) string {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		// If the address itself is just an integer, then we'll assume
		// that we're mapping this directly to a localhost:port pair.
		// This ensures we maintain the legacy behavior.
		if _, err := strconv.Atoi(address); err == nil {
			return net.JoinHostPort("localhost", address)
		}

		// Otherwise, we'll assume that the address just failed to
		// attach its own port, so we'll use the default port. In the
		// case of IPv6 addresses, if the host is already surrounded by
		// brackets, then we'll avoid using the JoinHostPort function,
		// since it will always add a pair of brackets.
		if strings.HasPrefix(address, "[") {
			return address + ":" + defaultPort
		}
		return net.JoinHostPort(address, defaultPort)
	}

	// In the case that both the host and port are empty, we'll use the
	// default port.
	if host == "" && port == "" {
		return ":" + defaultPort
	}

	return address
}

// ClientAddressDialer creates a gRPC dialer that can also dial unix socket
// addresses instead of just TCP addresses.
func ClientAddressDialer(defaultPort string) func(context.Context,
	string) (net.Conn, error) {

	return func(ctx context.Context, addr string) (net.Conn, error) {
		parsedAddr, err := ParseAddressString(
			addr, defaultPort, net.ResolveTCPAddr,
		)
		if err != nil {
			return nil, err
		}

		d := net.Dialer{}
		return d.DialContext(
			ctx, parsedAddr.Network(), parsedAddr.String(),
		)
	}
}
<commit_msg>lncfg: detect tor resolution error case and fallback to system resolver<commit_after>package lncfg

import (
	"context"
	"crypto/tls"
	"encoding/hex"
	"fmt"
	"net"
	"strconv"
	"strings"

	"github.com/btcsuite/btcd/btcec"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/tor"
)

var (
	loopBackAddrs = []string{"localhost", "127.0.0.1", "[::1]"}
)

// TCPResolver is a function signature that resolves an address on a given
// network.
type TCPResolver = func(network, addr string) (*net.TCPAddr, error)

// NormalizeAddresses returns a new slice with all the passed addresses
// normalized with the given default port and all duplicates removed.
func NormalizeAddresses(addrs []string, defaultPort string,
	tcpResolver TCPResolver) ([]net.Addr, error) {

	result := make([]net.Addr, 0, len(addrs))
	seen := map[string]struct{}{}

	for _, addr := range addrs {
		parsedAddr, err := ParseAddressString(
			addr, defaultPort, tcpResolver,
		)
		if err != nil {
			return nil, err
		}

		if _, ok := seen[parsedAddr.String()]; !ok {
			result = append(result, parsedAddr)
			seen[parsedAddr.String()] = struct{}{}
		}
	}

	return result, nil
}

// EnforceSafeAuthentication enforces "safe" authentication taking into account
// the interfaces that the RPC servers are listening on, and if macaroons and
// TLS is activated or not. To protect users from using dangerous config
// combinations, we'll prevent disabling authentication if the server is
// listening on a public interface.
func EnforceSafeAuthentication(addrs []net.Addr, macaroonsActive,
	tlsActive bool) error {

	// We'll now examine all addresses that this RPC server is listening
	// on. If it's a localhost address or a private address, we'll skip it,
	// otherwise, we'll return an error if macaroons are inactive.
	for _, addr := range addrs {
		if IsLoopback(addr.String()) || IsUnix(addr) || IsPrivate(addr) {
			continue
		}

		if !macaroonsActive {
			return fmt.Errorf("detected RPC server listening on "+
				"publicly reachable interface %v with "+
				"authentication disabled! Refusing to start "+
				"with --no-macaroons specified", addr)
		}

		if !tlsActive {
			return fmt.Errorf("detected RPC server listening on "+
				"publicly reachable interface %v with "+
				"encryption disabled! Refusing to start "+
				"with --notls specified", addr)
		}
	}

	return nil
}

// parseNetwork parses the network type of the given address.
func parseNetwork(addr net.Addr) string {
	switch addr := addr.(type) {
	// TCP addresses resolved through net.ResolveTCPAddr give a default
	// network of "tcp", so we'll map back the correct network for the given
	// address. This ensures that we can listen on the correct interface
	// (IPv4 vs IPv6).
	case *net.TCPAddr:
		if addr.IP.To4() != nil {
			return "tcp4"
		}
		return "tcp6"

	default:
		return addr.Network()
	}
}

// ListenOnAddress creates a listener that listens on the given address.
func ListenOnAddress(addr net.Addr) (net.Listener, error) {
	return net.Listen(parseNetwork(addr), addr.String())
}

// TLSListenOnAddress creates a TLS listener that listens on the given address.
func TLSListenOnAddress(addr net.Addr,
	config *tls.Config) (net.Listener, error) {
	return tls.Listen(parseNetwork(addr), addr.String(), config)
}

// IsLoopback returns true if an address describes a loopback interface.
func IsLoopback(addr string) bool {
	for _, loopback := range loopBackAddrs {
		if strings.Contains(addr, loopback) {
			return true
		}
	}

	return false
}

// isIPv6Host returns true if the host is IPV6 and false otherwise.
func isIPv6Host(host string) bool {
	v6Addr := net.ParseIP(host)
	if v6Addr == nil {
		return false
	}

	// The documentation states that if the IP address is an IPv6 address,
	// then To4() will return nil.
	return v6Addr.To4() == nil
}

// IsUnix returns true if an address describes an Unix socket address.
func IsUnix(addr net.Addr) bool {
	return strings.HasPrefix(addr.Network(), "unix")
}

// IsPrivate returns true if the address is private. The definitions are,
//   https://en.wikipedia.org/wiki/Link-local_address
//   https://en.wikipedia.org/wiki/Multicast_address
//   Local IPv4 addresses, https://tools.ietf.org/html/rfc1918
//   Local IPv6 addresses, https://tools.ietf.org/html/rfc4193
func IsPrivate(addr net.Addr) bool {
	switch addr := addr.(type) {
	case *net.TCPAddr:
		// Check 169.254.0.0/16 and fe80::/10.
		if addr.IP.IsLinkLocalUnicast() {
			return true
		}

		// Check 224.0.0.0/4 and ff00::/8.
		if addr.IP.IsLinkLocalMulticast() {
			return true
		}

		// Check 10.0.0.0/8, 172.16.0.0/12 and 192.168.0.0/16.
		if ip4 := addr.IP.To4(); ip4 != nil {
			return ip4[0] == 10 ||
				(ip4[0] == 172 && ip4[1]&0xf0 == 16) ||
				(ip4[0] == 192 && ip4[1] == 168)
		}

		// Check fc00::/7.
		return len(addr.IP) == net.IPv6len && addr.IP[0]&0xfe == 0xfc

	default:
		return false
	}
}

// ParseAddressString converts an address in string format to a net.Addr that is
// compatible with lnd. UDP is not supported because lnd needs reliable
// connections. We accept a custom function to resolve any TCP addresses so
// that caller is able control exactly how resolution is performed.
func ParseAddressString(strAddress string, defaultPort string,
	tcpResolver TCPResolver) (net.Addr, error) {

	var parsedNetwork, parsedAddr string

	// Addresses can either be in network://address:port format,
	// network:address:port, address:port, or just port. We want to support
	// all possible types.
	if strings.Contains(strAddress, "://") {
		parts := strings.Split(strAddress, "://")
		parsedNetwork, parsedAddr = parts[0], parts[1]
	} else if strings.Contains(strAddress, ":") {
		parts := strings.Split(strAddress, ":")
		parsedNetwork = parts[0]
		parsedAddr = strings.Join(parts[1:], ":")
	}

	// Only TCP and Unix socket addresses are valid. We can't use IP or
	// UDP only connections for anything we do in lnd.
	switch parsedNetwork {
	case "unix", "unixpacket":
		return net.ResolveUnixAddr(parsedNetwork, parsedAddr)

	case "tcp", "tcp4", "tcp6":
		return tcpResolver(
			parsedNetwork, verifyPort(parsedAddr, defaultPort),
		)

	case "ip", "ip4", "ip6", "udp", "udp4", "udp6", "unixgram":
		return nil, fmt.Errorf("only TCP or unix socket "+
			"addresses are supported: %s", parsedAddr)

	default:
		// We'll now possibly apply the default port, use the local
		// host short circuit, or parse out an all interfaces listen.
		addrWithPort := verifyPort(strAddress, defaultPort)
		rawHost, rawPort, _ := net.SplitHostPort(addrWithPort)

		// If we reach this point, then we'll check to see if we have
		// an onion addresses, if so, we can directly pass the raw
		// address and port to create the proper address.
		if tor.IsOnionHost(rawHost) {
			portNum, err := strconv.Atoi(rawPort)
			if err != nil {
				return nil, err
			}

			return &tor.OnionAddr{
				OnionService: rawHost,
				Port:         portNum,
			}, nil
		}

		// Otherwise, we'll attempt the resolve the host. The Tor
		// resolver is unable to resolve local or IPv6 addresses, so
		// we'll use the system resolver instead.
		if rawHost == "" || IsLoopback(rawHost) ||
			isIPv6Host(rawHost) {

			return net.ResolveTCPAddr("tcp", addrWithPort)
		}

		// If we've reached this point, then it's possible that this
		// resolve returns an error if it isn't able to resolve the
		// host. For eaxmple, local entries in /etc/hosts will fail to
		// be resolved by Tor. In order to handle this case, we'll fall
		// back to the normal system resolver if we fail with an
		// identifiable error.
		addr, err := tcpResolver("tcp", addrWithPort)
		if err != nil {
			torErrStr := "tor host is unreachable"
			if strings.Contains(err.Error(), torErrStr) {
				return net.ResolveTCPAddr("tcp", addrWithPort)
			}

			return nil, err
		}

		return addr, nil
	}
}

// ParseLNAddressString converts a string of the form <pubkey>@<addr> into an
// lnwire.NetAddress. The <pubkey> must be presented in hex, and result in a
// 33-byte, compressed public key that lies on the secp256k1 curve. The <addr>
// may be any address supported by ParseAddressString. If no port is specified,
// the defaultPort will be used. Any tcp addresses that need resolving will be
// resolved using the custom TCPResolver.
func ParseLNAddressString(strAddress string, defaultPort string,
	tcpResolver TCPResolver) (*lnwire.NetAddress, error) {

	// Split the address string around the @ sign.
	parts := strings.Split(strAddress, "@")

	// The string is malformed if there are not exactly two parts.
	if len(parts) != 2 {
		return nil, fmt.Errorf("invalid lightning address %s: "+
			"must be of the form <pubkey-hex>@<addr>", strAddress)
	}

	// Now, take the first portion as the hex pubkey, and the latter as the
	// address string.
	parsedPubKey, parsedAddr := parts[0], parts[1]

	// Decode the hex pubkey to get the raw compressed pubkey bytes.
	pubKeyBytes, err := hex.DecodeString(parsedPubKey)
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address pubkey: %v", err)
	}

	// The compressed pubkey should have a length of exactly 33 bytes.
	if len(pubKeyBytes) != 33 {
		return nil, fmt.Errorf("invalid lightning address pubkey: "+
			"length must be 33 bytes, found %d", len(pubKeyBytes))
	}

	// Parse the pubkey bytes to verify that it corresponds to valid public
	// key on the secp256k1 curve.
	pubKey, err := btcec.ParsePubKey(pubKeyBytes, btcec.S256())
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address pubkey: %v", err)
	}

	// Finally, parse the address string using our generic address parser.
	addr, err := ParseAddressString(parsedAddr, defaultPort, tcpResolver)
	if err != nil {
		return nil, fmt.Errorf("invalid lightning address address: %v", err)
	}

	return &lnwire.NetAddress{
		IdentityKey: pubKey,
		Address:     addr,
	}, nil
}

// verifyPort makes sure that an address string has both a host and a port. If
// there is no port found, the default port is appended. If the address is just
// a port, then we'll assume that the user is using the short cut to specify a
// localhost:port address.
func verifyPort(address string, defaultPort string) string {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		// If the address itself is just an integer, then we'll assume
		// that we're mapping this directly to a localhost:port pair.
		// This ensures we maintain the legacy behavior.
		if _, err := strconv.Atoi(address); err == nil {
			return net.JoinHostPort("localhost", address)
		}

		// Otherwise, we'll assume that the address just failed to
		// attach its own port, so we'll use the default port. In the
		// case of IPv6 addresses, if the host is already surrounded by
		// brackets, then we'll avoid using the JoinHostPort function,
		// since it will always add a pair of brackets.
		if strings.HasPrefix(address, "[") {
			return address + ":" + defaultPort
		}
		return net.JoinHostPort(address, defaultPort)
	}

	// In the case that both the host and port are empty, we'll use the
	// default port.
	if host == "" && port == "" {
		return ":" + defaultPort
	}

	return address
}

// ClientAddressDialer creates a gRPC dialer that can also dial unix socket
// addresses instead of just TCP addresses.
func ClientAddressDialer(defaultPort string) func(context.Context,
	string) (net.Conn, error) {

	return func(ctx context.Context, addr string) (net.Conn, error) {
		parsedAddr, err := ParseAddressString(
			addr, defaultPort, net.ResolveTCPAddr,
		)
		if err != nil {
			return nil, err
		}

		d := net.Dialer{}
		return d.DialContext(
			ctx, parsedAddr.Network(), parsedAddr.String(),
		)
	}
}
<|endoftext|>


<commit_before>package resources

import (
	"github.com/hashicorp/terraform/helper/schema"
	"code.cloudfoundry.org/cli/cf/models"
	"github.com/orange-cloudfoundry/terraform-provider-cloudfoundry/cf_client"
	"log"
	"strings"
)

const DEFAULT_ORG_QUOTA_NAME = "default"

type CfOrganizationResource struct {
	CfResource
}

func NewCfOrganizationResource() CfResource {
	return &CfOrganizationResource{}
}
func (c CfOrganizationResource) resourceObject(d *schema.ResourceData) models.Organization {
	quotaDef := models.QuotaFields{
		GUID: d.Get("quota_id").(string),
	}
	orgField := models.OrganizationFields{
		GUID: d.Id(),
		Name: d.Get("name").(string),
		QuotaDefinition: quotaDef,
	}
	repo := models.Organization{
		OrganizationFields: orgField,
	}

	return repo
}
func (c CfOrganizationResource) Create(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	org := c.resourceObject(d)
	if ok, _ := c.Exists(d, meta); ok {
		log.Printf(
			"[INFO] skipping creation of organization %s/%s because it already exists on your Cloud Foundry",
			client.Config().ApiEndpoint,
			org.Name,
		)
	} else {
		err := client.Organizations().Create(org)
		if err != nil {
			return err
		}
		c.Exists(d, meta)
	}
	orgCf, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	err = c.bindQuota(client, org.QuotaDefinition.GUID, orgCf.GUID)
	if err != nil {
		return err
	}
	d.SetId(orgCf.GUID)
	return nil
}
func (c CfOrganizationResource) bindQuota(client cf_client.Client, quotaId, orgId string) error {
	if quotaId == "" {
		return nil
	}
	return client.Quotas().AssignQuotaToOrg(orgId, quotaId)
}
func (c CfOrganizationResource) Read(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	orgName := d.Get("name").(string)
	org, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	if org.GUID == "" {
		log.Printf(
			"[WARN] removing organization %s/%s from state because it no longer exists in your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		d.SetId("")
		return nil
	}
	quota, _ := client.Quotas().FindByName(DEFAULT_ORG_QUOTA_NAME)
	if quota.GUID == "" || org.QuotaDefinition.GUID != quota.GUID {
		d.Set("quota_id", org.QuotaDefinition.GUID)
	}
	d.Set("name", org.Name)
	return nil

}
func (c CfOrganizationResource) getOrgFromCf(client cf_client.Client, orgGuid string) (models.Organization, error) {
	orgs, err := client.Organizations().GetManyOrgsByGUID([]string{orgGuid})
	if err != nil {
		if strings.Contains(err.Error(), "status code: 404") {
			return models.Organization{}, nil
		}
		return models.Organization{}, err
	}
	return client.Organizations().FindByName(orgs[0].Name)
}
func (c CfOrganizationResource) Exists(d *schema.ResourceData, meta interface{}) (bool, error) {
	client := meta.(cf_client.Client)
	orgName := d.Get("name").(string)
	org, err := client.Organizations().FindByName(orgName)
	if err != nil {
		if strings.Contains(err.Error(), "not found") {
			return false, nil
		}
		return false, err
	}
	d.SetId(org.GUID)
	return true, nil
}
func (c CfOrganizationResource) Update(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	quotaId := d.Get("quota_id").(string)
	orgName := d.Get("name").(string)
	org, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	if org.GUID == "" {
		log.Printf(
			"[WARN] removing organization %s/%s from state because it no longer exists in your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		d.SetId("")
		return nil
	}
	if org.QuotaDefinition.GUID != quotaId {
		err = c.bindQuota(client, quotaId, d.Id())
		if err != nil {
			return err
		}
	}
	if org.Name != orgName {
		err = client.Organizations().Rename(d.Id(), d.Get("name").(string))
		if err != nil {
			return err
		}
	}

	return nil
}
func (c CfOrganizationResource) Delete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	isSystemDomain := d.Get("is_system_domain").(bool)
	orgName := d.Get("name").(bool)
	if isSystemDomain {
		log.Printf(
			"[WARN] removing organization %s/%s isn't possible because it's the system_domain organization of your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		return nil
	}
	return client.Organizations().Delete(d.Id())
}
func (c CfOrganizationResource) Schema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
		},
		"quota_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
		},
		"is_system_domain": &schema.Schema{
			Type:     schema.TypeBool,
			Optional: true,
		},
	}
}
<commit_msg>fix deletion on organization<commit_after>package resources

import (
	"github.com/hashicorp/terraform/helper/schema"
	"code.cloudfoundry.org/cli/cf/models"
	"github.com/orange-cloudfoundry/terraform-provider-cloudfoundry/cf_client"
	"log"
	"strings"
)

const DEFAULT_ORG_QUOTA_NAME = "default"

type CfOrganizationResource struct {
	CfResource
}

func NewCfOrganizationResource() CfResource {
	return &CfOrganizationResource{}
}
func (c CfOrganizationResource) resourceObject(d *schema.ResourceData) models.Organization {
	quotaDef := models.QuotaFields{
		GUID: d.Get("quota_id").(string),
	}
	orgField := models.OrganizationFields{
		GUID: d.Id(),
		Name: d.Get("name").(string),
		QuotaDefinition: quotaDef,
	}
	repo := models.Organization{
		OrganizationFields: orgField,
	}

	return repo
}
func (c CfOrganizationResource) Create(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	org := c.resourceObject(d)
	if ok, _ := c.Exists(d, meta); ok {
		log.Printf(
			"[INFO] skipping creation of organization %s/%s because it already exists on your Cloud Foundry",
			client.Config().ApiEndpoint,
			org.Name,
		)
	} else {
		err := client.Organizations().Create(org)
		if err != nil {
			return err
		}
		c.Exists(d, meta)
	}
	orgCf, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	err = c.bindQuota(client, org.QuotaDefinition.GUID, orgCf.GUID)
	if err != nil {
		return err
	}
	d.SetId(orgCf.GUID)
	return nil
}
func (c CfOrganizationResource) bindQuota(client cf_client.Client, quotaId, orgId string) error {
	if quotaId == "" {
		return nil
	}
	return client.Quotas().AssignQuotaToOrg(orgId, quotaId)
}
func (c CfOrganizationResource) Read(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	orgName := d.Get("name").(string)
	org, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	if org.GUID == "" {
		log.Printf(
			"[WARN] removing organization %s/%s from state because it no longer exists in your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		d.SetId("")
		return nil
	}
	quota, _ := client.Quotas().FindByName(DEFAULT_ORG_QUOTA_NAME)
	if quota.GUID == "" || org.QuotaDefinition.GUID != quota.GUID {
		d.Set("quota_id", org.QuotaDefinition.GUID)
	}
	d.Set("name", org.Name)
	return nil

}
func (c CfOrganizationResource) getOrgFromCf(client cf_client.Client, orgGuid string) (models.Organization, error) {
	orgs, err := client.Organizations().GetManyOrgsByGUID([]string{orgGuid})
	if err != nil {
		if strings.Contains(err.Error(), "status code: 404") {
			return models.Organization{}, nil
		}
		return models.Organization{}, err
	}
	return client.Organizations().FindByName(orgs[0].Name)
}
func (c CfOrganizationResource) Exists(d *schema.ResourceData, meta interface{}) (bool, error) {
	client := meta.(cf_client.Client)
	orgName := d.Get("name").(string)
	org, err := client.Organizations().FindByName(orgName)
	if err != nil {
		if strings.Contains(err.Error(), "not found") {
			return false, nil
		}
		return false, err
	}
	d.SetId(org.GUID)
	return true, nil
}
func (c CfOrganizationResource) Update(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	quotaId := d.Get("quota_id").(string)
	orgName := d.Get("name").(string)
	org, err := c.getOrgFromCf(client, d.Id())
	if err != nil {
		return err
	}
	if org.GUID == "" {
		log.Printf(
			"[WARN] removing organization %s/%s from state because it no longer exists in your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		d.SetId("")
		return nil
	}
	if org.QuotaDefinition.GUID != quotaId {
		err = c.bindQuota(client, quotaId, d.Id())
		if err != nil {
			return err
		}
	}
	if org.Name != orgName {
		err = client.Organizations().Rename(d.Id(), d.Get("name").(string))
		if err != nil {
			return err
		}
	}

	return nil
}
func (c CfOrganizationResource) Delete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(cf_client.Client)
	isSystemDomain := d.Get("is_system_domain").(bool)
	orgName := d.Get("name").(string)
	if isSystemDomain {
		log.Printf(
			"[WARN] removing organization %s/%s isn't possible because it's the system_domain organization of your Cloud Foundry",
			client.Config().ApiEndpoint,
			orgName,
		)
		return nil
	}
	return client.Organizations().Delete(d.Id())
}
func (c CfOrganizationResource) Schema() map[string]*schema.Schema {
	return map[string]*schema.Schema{
		"name": &schema.Schema{
			Type:     schema.TypeString,
			Required: true,
		},
		"quota_id": &schema.Schema{
			Type:     schema.TypeString,
			Optional: true,
		},
		"is_system_domain": &schema.Schema{
			Type:     schema.TypeBool,
			Optional: true,
		},
	}
}
<|endoftext|>


<commit_before>// Package BattleEye doco goes here
package BattleEye

import (
	"errors"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"
)

// Config documentation
type Config struct {
	addr     *net.UDPAddr
	Password string
	// time in seconds to wait for a response. defaults to 2.
	ConnTimeout uint32
	// time in seconds to wait for response. defaults to 1.
	ResponseTimeout uint32
	// wait time after first command response to check if multiple packets arrive.
	// defaults. 0.5s
	MultiResponseTimeout uint32
	// Time in seconds between sending a heartbeat when no commands are being sent. defaults 5
	HeartBeatTimer uint32
}

// GetConfig Returns the config to satisfy the interface for setting up a new battleeye connection
func (bec Config) GetConfig() Config {
	return bec
}

// BeConfig is the interface for passing in a configuration for the client.
// this allows other types to be implemented that also contain the type desired
type BeConfig interface {
	GetConfig() Config
}
type transmission struct {
	packet   []byte
	sequence byte
	sent     time.Time
	w        io.Writer
}

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

// BattleEye must do doco soon
type BattleEye struct {

	// Passed in config

	password             string
	addr                 *net.UDPAddr
	connTimeout          uint32
	responseTimeout      uint32
	multiResponseTimeout uint32
	heartbeatTimer       uint32

	// Sequence byte to determine the packet we are up to in the chain.
	sequence struct {
		sync.Locker
		n byte
	}
	chatWriter  io.Writer
	writebuffer []byte

	conn              *net.UDPConn
	lastCommandPacket struct {
		sync.Mutex
		time.Time
	}
	running bool
	wg      sync.WaitGroup

	// use this to unlock before reading.
	// and match reads to waiting confirms to purge this list.
	// or possibly resend
	packetQueue struct {
		sync.Mutex
		queue []transmission
	}
}

// New Creates and Returns a new Client
func New(config BeConfig) *BattleEye {
	// setup all variables
	cfg := config.GetConfig()
	if cfg.ConnTimeout == 0 {
		cfg.ConnTimeout = 2
	}
	if cfg.ResponseTimeout == 0 {
		cfg.ResponseTimeout = 1
	}
	if cfg.HeartBeatTimer == 0 {
		cfg.HeartBeatTimer = 5
	}

	return &BattleEye{
		password:        cfg.Password,
		addr:            cfg.addr,
		connTimeout:     cfg.ConnTimeout,
		responseTimeout: cfg.ResponseTimeout,
		heartbeatTimer:  cfg.HeartBeatTimer,
		writebuffer:     make([]byte, 4096),
	}

}

// SendCommand takes a byte array of a command string i.e 'ban xyz' and a io.Writer, it will
// Make sure the server recieves the command by retrying if needed and write the response to the writer.
// if no response is recieved then a response has not yet been recieved. a empty write
func (be *BattleEye) SendCommand(command []byte, w io.Writer) error {
	be.sequence.Lock()
	sequence := be.sequence.n
	// increment the sending packet.
	if be.sequence.n == 255 {
		be.sequence.n = 0
	} else {
		be.sequence.n++
	}
	be.sequence.Unlock()

	packet := buildCommandPacket(command, sequence)
	be.conn.SetWriteDeadline(time.Now().Add(time.Second * time.Duration(be.responseTimeout)))
	be.conn.Write(packet)
	be.packetQueue.Lock()
	be.packetQueue.queue = append(be.packetQueue.queue, transmission{packet: packet, sequence: sequence, sent: time.Now(), w: w})
	be.lastCommandPacket.Lock()
	be.lastCommandPacket.Time = time.Now()
	be.lastCommandPacket.Unlock()

	return nil
}

// Connect attempts to establish a connection with the BattlEye Rcon server and if it works it then sets up a loop in a goroutine
// to recieve all callbacks.
func (be *BattleEye) Connect() (bool, error) {
	be.wg = sync.WaitGroup{}
	var err error
	// dial the Address
	be.conn, err = net.DialUDP("udp", nil, be.addr)
	if err != nil {
		return false, err
	}
	// make a buffer to read the packet packed with extra space
	packet := make([]byte, 9)

	// set timeout deadline so we dont block forever
	be.conn.SetReadDeadline(time.Now().Add(time.Second * 2))
	// Send a Connection Packet
	be.conn.Write(buildConnectionPacket(be.password))
	// Read connection and hope it doesn't time out and the server responds
	n, err := be.conn.Read(packet)
	// check if this is a timeout error.
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return false, errors.New("Connection Timed Out")
	}
	if err != nil {
		return false, err
	}

	result, err := checkLogin(packet[:n])
	if err != nil {
		return false, err
	}

	if result == packetResponse.LoginFail {
		return false, nil
	}

	// nothing has failed we are good to go :).
	// Spin up a go routine to read back on a connection
	be.wg.Add(1)
	//go
	return true, nil
}

func (be *BattleEye) updateLoop() {
	defer be.wg.Done()
	for {
		if be.conn == nil {
			return
		}
		t := time.Now()

		be.lastCommandPacket.Lock()
		if t.After(be.lastCommandPacket.Add(time.Second * time.Duration(be.heartbeatTimer))) {
			err := be.SendCommand([]byte{}, ioutil.Discard)
			if err != nil {
				return
			}
			be.lastCommandPacket.Time = t
		}
		be.lastCommandPacket.Unlock()

		// do check for new incoming data
		be.conn.SetReadDeadline(time.Now().Add(time.Millisecond * 100))
		n, err := be.conn.Read(be.writebuffer)
		if err != nil {
			continue
		}
		data := be.writebuffer[:n]
		be.processPacket(data)
	}
}
func (be *BattleEye) processPacket(data []byte) error {
	sequence, content, pType, err := verifyPacket(data)

	if err != nil {
		// maybe should log this shit somewhere
		//fmt.Println("")
		//fmt.Errorf(format, ...)
		return err
	}
	// if say command write acknoledge and leave
	if pType == packetType.ServerMessage {
		be.sequence.Lock()
		if sequence >= be.sequence.n {
			// not sure how byte overflow is handled in golang... not sure if it would roll over to 0 or throw and error.
			if sequence == 255 {
				be.sequence.n = 0x00
			} else {
				be.sequence.n = sequence + 1
			}
		}
		be.sequence.Unlock()
		be.chatWriter.Write(content)
		// we must acknoledge we recieved this first
		if be.conn != nil {
			be.conn.Write(buildPacket([]byte{sequence}, packetType.ServerMessage))
		}

		return nil
	}

	// else for command check if we expect more packets and how many.
	if pType != packetType.Command {
		return errors.New("Unknown way to respond to packet type: " + string(pType))
	}
	packetCount, currentPacket, isMultiPacket := checkMultiPacketResponse(content)
	// process the packet if it is not a multipacket
	if !isMultiPacket {
		be.handleResponseToQueue(sequence, content[2:], false)
		return nil
	}
	// loop till we have all the messages and i guess send confirms back.
	for ; packetCount < currentPacket; packetCount++ {
		be.conn.SetReadDeadline(time.Now().Add(time.Second))
		n, err := be.conn.Read(be.writebuffer)
		if err != nil {
			return err
		}
		// lets re verify this entire thing
		p := be.writebuffer[:n]
		seq, cont, _, err := verifyPacket(p)
		if err != nil {
			return err
		}
		if seq != sequence {
			be.processPacket(p)
			packetCount--
		}
		be.handleResponseToQueue(seq, cont[2:], true)

	}
	//closes it as we know we have done our job
	be.handleResponseToQueue(sequence, []byte{}, false)
	// now that we have goten all the packets we are after and writen them to the buffer lets return the result.
	return nil
}

// Disconnect shuts down the infinite loop to recieve packets and closes the connection
func (be *BattleEye) Disconnect() error {
	// maybe also close the main loop and wait for that?
	be.conn.Close()
	be.wg.Wait()
	return nil
}

func checkLogin(packet []byte) (byte, error) {
	var err error
	if len(packet) != 9 {
		return 0, errors.New("Packet Size Invalid for Response")
	}
	// check if we have a valid packet
	if match, err := packetMatchesChecksum(packet); match == false || err != nil {
		return 0, err
	}
	// now check if we got a success or a fail
	// 2 byte prefix. 4 byte checksum. 1 byte terminate header. 1 byte login type. 1 byte result
	return packet[8], err
}

func (be *BattleEye) handleResponseToQueue(sequence byte, response []byte, moreToCome bool) {
	be.packetQueue.Lock()
	for k, v := range be.packetQueue.queue {
		if v.sequence == sequence {
			if v.w != nil {
				v.w.Write(response)
			}

			if !moreToCome {
				//v.w.Close()
				be.packetQueue.queue = append(be.packetQueue.queue[:k], be.packetQueue.queue[k+1:]...)
			}
			break
		}
	}
	be.packetQueue.Unlock()
}

func verifyPacket(data []byte) (sequence byte, content []byte, pType byte, err error) {
	checksum, err := getCheckSumFromBEPacket(data)
	if err != nil {
		return
	}
	if len(data) < 6 {
		err = errors.New("Packet size too small to have data")
		return
	}
	match := dataMatchesCheckSum(data[6:], checksum)
	if !match {
		err = errors.New("Checksum does not match data")
		return
	}
	sequence, err = getSequenceFromPacket(data)
	if err != nil {
		return
	}

	content, err = stripHeader(data)
	if err != nil {
		return
	}
	pType, err = responseType(data)

	return
}
<commit_msg>made variable public that should be.<commit_after>// Package BattleEye doco goes here
package BattleEye

import (
	"errors"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"
)

// Config documentation
type Config struct {
	Addr     *net.UDPAddr
	Password string
	// time in seconds to wait for a response. defaults to 2.
	ConnTimeout uint32
	// time in seconds to wait for response. defaults to 1.
	ResponseTimeout uint32
	// wait time after first command response to check if multiple packets arrive.
	// defaults. 0.5s
	MultiResponseTimeout uint32
	// Time in seconds between sending a heartbeat when no commands are being sent. defaults 5
	HeartBeatTimer uint32
}

// GetConfig Returns the config to satisfy the interface for setting up a new battleeye connection
func (bec Config) GetConfig() Config {
	return bec
}

// BeConfig is the interface for passing in a configuration for the client.
// this allows other types to be implemented that also contain the type desired
type BeConfig interface {
	GetConfig() Config
}
type transmission struct {
	packet   []byte
	sequence byte
	sent     time.Time
	w        io.Writer
}

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

// BattleEye must do doco soon
type BattleEye struct {

	// Passed in config

	password             string
	addr                 *net.UDPAddr
	connTimeout          uint32
	responseTimeout      uint32
	multiResponseTimeout uint32
	heartbeatTimer       uint32

	// Sequence byte to determine the packet we are up to in the chain.
	sequence struct {
		sync.Locker
		n byte
	}
	chatWriter  io.Writer
	writebuffer []byte

	conn              *net.UDPConn
	lastCommandPacket struct {
		sync.Mutex
		time.Time
	}
	running bool
	wg      sync.WaitGroup

	// use this to unlock before reading.
	// and match reads to waiting confirms to purge this list.
	// or possibly resend
	packetQueue struct {
		sync.Mutex
		queue []transmission
	}
}

// New Creates and Returns a new Client
func New(config BeConfig) *BattleEye {
	// setup all variables
	cfg := config.GetConfig()
	if cfg.ConnTimeout == 0 {
		cfg.ConnTimeout = 2
	}
	if cfg.ResponseTimeout == 0 {
		cfg.ResponseTimeout = 1
	}
	if cfg.HeartBeatTimer == 0 {
		cfg.HeartBeatTimer = 5
	}

	return &BattleEye{
		password:        cfg.Password,
		addr:            cfg.Addr,
		connTimeout:     cfg.ConnTimeout,
		responseTimeout: cfg.ResponseTimeout,
		heartbeatTimer:  cfg.HeartBeatTimer,
		writebuffer:     make([]byte, 4096),
	}

}

// SendCommand takes a byte array of a command string i.e 'ban xyz' and a io.Writer, it will
// Make sure the server recieves the command by retrying if needed and write the response to the writer.
// if no response is recieved then a response has not yet been recieved. a empty write
func (be *BattleEye) SendCommand(command []byte, w io.Writer) error {
	be.sequence.Lock()
	sequence := be.sequence.n
	// increment the sending packet.
	if be.sequence.n == 255 {
		be.sequence.n = 0
	} else {
		be.sequence.n++
	}
	be.sequence.Unlock()

	packet := buildCommandPacket(command, sequence)
	be.conn.SetWriteDeadline(time.Now().Add(time.Second * time.Duration(be.responseTimeout)))
	be.conn.Write(packet)
	be.packetQueue.Lock()
	be.packetQueue.queue = append(be.packetQueue.queue, transmission{packet: packet, sequence: sequence, sent: time.Now(), w: w})
	be.lastCommandPacket.Lock()
	be.lastCommandPacket.Time = time.Now()
	be.lastCommandPacket.Unlock()

	return nil
}

// Connect attempts to establish a connection with the BattlEye Rcon server and if it works it then sets up a loop in a goroutine
// to recieve all callbacks.
func (be *BattleEye) Connect() (bool, error) {
	be.wg = sync.WaitGroup{}
	var err error
	// dial the Address
	be.conn, err = net.DialUDP("udp", nil, be.addr)
	if err != nil {
		return false, err
	}
	// make a buffer to read the packet packed with extra space
	packet := make([]byte, 9)

	// set timeout deadline so we dont block forever
	be.conn.SetReadDeadline(time.Now().Add(time.Second * 2))
	// Send a Connection Packet
	be.conn.Write(buildConnectionPacket(be.password))
	// Read connection and hope it doesn't time out and the server responds
	n, err := be.conn.Read(packet)
	// check if this is a timeout error.
	if err, ok := err.(net.Error); ok && err.Timeout() {
		return false, errors.New("Connection Timed Out")
	}
	if err != nil {
		return false, err
	}

	result, err := checkLogin(packet[:n])
	if err != nil {
		return false, err
	}

	if result == packetResponse.LoginFail {
		return false, nil
	}

	// nothing has failed we are good to go :).
	// Spin up a go routine to read back on a connection
	be.wg.Add(1)
	//go
	return true, nil
}

func (be *BattleEye) updateLoop() {
	defer be.wg.Done()
	for {
		if be.conn == nil {
			return
		}
		t := time.Now()

		be.lastCommandPacket.Lock()
		if t.After(be.lastCommandPacket.Add(time.Second * time.Duration(be.heartbeatTimer))) {
			err := be.SendCommand([]byte{}, ioutil.Discard)
			if err != nil {
				return
			}
			be.lastCommandPacket.Time = t
		}
		be.lastCommandPacket.Unlock()

		// do check for new incoming data
		be.conn.SetReadDeadline(time.Now().Add(time.Millisecond * 100))
		n, err := be.conn.Read(be.writebuffer)
		if err != nil {
			continue
		}
		data := be.writebuffer[:n]
		be.processPacket(data)
	}
}
func (be *BattleEye) processPacket(data []byte) error {
	sequence, content, pType, err := verifyPacket(data)

	if err != nil {
		// maybe should log this shit somewhere
		//fmt.Println("")
		//fmt.Errorf(format, ...)
		return err
	}
	// if say command write acknoledge and leave
	if pType == packetType.ServerMessage {
		be.sequence.Lock()
		if sequence >= be.sequence.n {
			// not sure how byte overflow is handled in golang... not sure if it would roll over to 0 or throw and error.
			if sequence == 255 {
				be.sequence.n = 0x00
			} else {
				be.sequence.n = sequence + 1
			}
		}
		be.sequence.Unlock()
		be.chatWriter.Write(content)
		// we must acknoledge we recieved this first
		if be.conn != nil {
			be.conn.Write(buildPacket([]byte{sequence}, packetType.ServerMessage))
		}

		return nil
	}

	// else for command check if we expect more packets and how many.
	if pType != packetType.Command {
		return errors.New("Unknown way to respond to packet type: " + string(pType))
	}
	packetCount, currentPacket, isMultiPacket := checkMultiPacketResponse(content)
	// process the packet if it is not a multipacket
	if !isMultiPacket {
		be.handleResponseToQueue(sequence, content[2:], false)
		return nil
	}
	// loop till we have all the messages and i guess send confirms back.
	for ; packetCount < currentPacket; packetCount++ {
		be.conn.SetReadDeadline(time.Now().Add(time.Second))
		n, err := be.conn.Read(be.writebuffer)
		if err != nil {
			return err
		}
		// lets re verify this entire thing
		p := be.writebuffer[:n]
		seq, cont, _, err := verifyPacket(p)
		if err != nil {
			return err
		}
		if seq != sequence {
			be.processPacket(p)
			packetCount--
		}
		be.handleResponseToQueue(seq, cont[2:], true)

	}
	//closes it as we know we have done our job
	be.handleResponseToQueue(sequence, []byte{}, false)
	// now that we have goten all the packets we are after and writen them to the buffer lets return the result.
	return nil
}

// Disconnect shuts down the infinite loop to recieve packets and closes the connection
func (be *BattleEye) Disconnect() error {
	// maybe also close the main loop and wait for that?
	be.conn.Close()
	be.wg.Wait()
	return nil
}

func checkLogin(packet []byte) (byte, error) {
	var err error
	if len(packet) != 9 {
		return 0, errors.New("Packet Size Invalid for Response")
	}
	// check if we have a valid packet
	if match, err := packetMatchesChecksum(packet); match == false || err != nil {
		return 0, err
	}
	// now check if we got a success or a fail
	// 2 byte prefix. 4 byte checksum. 1 byte terminate header. 1 byte login type. 1 byte result
	return packet[8], err
}

func (be *BattleEye) handleResponseToQueue(sequence byte, response []byte, moreToCome bool) {
	be.packetQueue.Lock()
	for k, v := range be.packetQueue.queue {
		if v.sequence == sequence {
			if v.w != nil {
				v.w.Write(response)
			}

			if !moreToCome {
				//v.w.Close()
				be.packetQueue.queue = append(be.packetQueue.queue[:k], be.packetQueue.queue[k+1:]...)
			}
			break
		}
	}
	be.packetQueue.Unlock()
}

func verifyPacket(data []byte) (sequence byte, content []byte, pType byte, err error) {
	checksum, err := getCheckSumFromBEPacket(data)
	if err != nil {
		return
	}
	if len(data) < 6 {
		err = errors.New("Packet size too small to have data")
		return
	}
	match := dataMatchesCheckSum(data[6:], checksum)
	if !match {
		err = errors.New("Checksum does not match data")
		return
	}
	sequence, err = getSequenceFromPacket(data)
	if err != nil {
		return
	}

	content, err = stripHeader(data)
	if err != nil {
		return
	}
	pType, err = responseType(data)

	return
}
<|endoftext|>


<commit_before>package Common

import (
	"fmt"
	"log"
	"net"
	"net/rpc"
)

func ListenRpc(addr string, obj interface{}) error {
	defer CheckPanic()

	if obj == nil {
		return fmt.Errorf("rpc object is nil")
	}

	server := rpc.NewServer()
	server.Register(obj)

	if listener, err := net.Listen("tcp", addr); err != nil {
		return fmt.Errorf("rpc listen error : %v : %v", addr, err)
	} else {
		log.Println("rpc running @", addr)
		server.Accept(listener)
	}

	return fmt.Errorf("rpc server quit : %v", addr)
}

func ListenSocket(addr string, keepalive bool, reactiver func(net.Conn)) error {
	defer CheckPanic()

	if reactiver == nil {
		return fmt.Errorf("socket reactiver is nil")
	}

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return fmt.Errorf("can't resolve addr : %v : %v", addr, err)
	}

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return fmt.Errorf("can't listen tcp : %v : %v", addr, err)
	}

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			log.Println("accept ", err)
			break
		}

		conn.SetNoDelay(true)
		conn.SetKeepAlive(keepalive)
		conn.SetLinger(-1)
		go reactiver(conn)
	}

	return fmt.Errorf("socket server quit : %v ", addr)
}
<commit_msg>add listen udp<commit_after>package Common

import (
	"fmt"
	"log"
	"net"
	"net/rpc"
)

func ListenRpc(addr string, obj interface{}) error {
	defer CheckPanic()

	if obj == nil {
		return fmt.Errorf("rpc object is nil")
	}

	server := rpc.NewServer()
	server.Register(obj)

	if listener, err := net.Listen("tcp", addr); err != nil {
		return fmt.Errorf("rpc listen error : %v : %v", addr, err)
	} else {
		log.Println("rpc running @", addr)
		server.Accept(listener)
	}

	return fmt.Errorf("rpc server quit : %v", addr)
}

func ListenSocket(addr string, keepalive bool, reactiver func(net.Conn)) error {
	defer CheckPanic()

	if reactiver == nil {
		return fmt.Errorf("socket reactiver is nil")
	}

	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return fmt.Errorf("can't resolve addr : %v : %v", addr, err)
	}

	listener, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return fmt.Errorf("can't listen tcp : %v : %v", addr, err)
	}

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			log.Println("accept ", err)
			break
		}

		conn.SetNoDelay(true)
		conn.SetKeepAlive(keepalive)
		conn.SetLinger(-1)
		go reactiver(conn)
	}

	return fmt.Errorf("socket server quit : %v ", addr)
}

func ListenUdp(addr string, bufsize int, reactiver func(*net.UDPAddr, []byte) []byte) error {
	defer CheckPanic()

	if reactiver == nil {
		return fmt.Errorf("udp reactiver is nil")
	}

	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return fmt.Errorf("can't resolve addr : %v : %v", addr, err)
	}

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return fmt.Errorf("can't listen udp : %v : %v", addr, err)
	}
	defer conn.Close()

	for {
		data := make([]byte, bufsize)
		nr, remote, err := conn.ReadFromUDP(data)
		if err != nil {
			log.Println("read from udp ", remote, err)
			continue
		}

		conn.WriteToUDP(reactiver(remote, data[:nr]), remote)
	}

	return fmt.Errorf("udp server quit : %v ", addr)
}
<|endoftext|>