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

// NOTE: If you change this file you must run "./mkbuiltin"
// to update builtin.c.boot.  This is not done automatically
// to avoid depending on having a working compiler binary.

// +build ignore

package PACKAGE

// emitted by compiler, not referred to by go programs

func new(typ *byte) *any
func panicindex()
func panicslice()
func throwreturn()
func throwinit()
func panicwrap(string, string, string)

func panic(interface{})
func recover(*int32) interface{}

func printbool(bool)
func printfloat(float64)
func printint(int64)
func printuint(uint64)
func printcomplex(complex128)
func printstring(string)
func printpointer(any)
func printiface(any)
func printeface(any)
func printslice(any)
func printnl()
func printsp()
func goprintf()

// filled in by compiler: int n, string, string, ...
func concatstring()

// filled in by compiler: Type*, int n, Slice, ...
func append()
func appendslice(typ *byte, x any, y []any) any
func appendstr(typ *byte, x []byte, y string) []byte

func cmpstring(string, string) int
func slicestring(string, int, int) string
func slicestring1(string, int) string
func intstring(int64) string
func slicebytetostring([]byte) string
func slicerunetostring([]rune) string
func stringtoslicebyte(string) []byte
func stringtoslicerune(string) []rune
func stringiter(string, int) int
func stringiter2(string, int) (retk int, retv rune)
func copy(to any, fr any, wid uint32) int
func slicestringcopy(to any, fr any) int

// interface conversions
func convI2E(elem any) (ret any)
func convI2I(typ *byte, elem any) (ret any)
func convT2E(typ *byte, elem any) (ret any)
func convT2I(typ *byte, typ2 *byte, elem any) (ret any)

// interface type assertions  x.(T)
func assertE2E(typ *byte, iface any) (ret any)
func assertE2E2(typ *byte, iface any) (ret any, ok bool)
func assertE2I(typ *byte, iface any) (ret any)
func assertE2I2(typ *byte, iface any) (ret any, ok bool)
func assertE2T(typ *byte, iface any) (ret any)
func assertE2T2(typ *byte, iface any) (ret any, ok bool)
func assertI2E(typ *byte, iface any) (ret any)
func assertI2E2(typ *byte, iface any) (ret any, ok bool)
func assertI2I(typ *byte, iface any) (ret any)
func assertI2I2(typ *byte, iface any) (ret any, ok bool)
func assertI2T(typ *byte, iface any) (ret any)
func assertI2T2(typ *byte, iface any) (ret any, ok bool)

func ifaceeq(i1 any, i2 any) (ret bool)
func efaceeq(i1 any, i2 any) (ret bool)
func ifacethash(i1 any) (ret uint32)
func efacethash(i1 any) (ret uint32)

func equal(typ *byte, x1, x2 any) (ret bool)

// *byte is really *runtime.Type
func makemap(mapType *byte, hint int64) (hmap map[any]any)
func mapaccess1(mapType *byte, hmap map[any]any, key any) (val any)
func mapaccess2(mapType *byte, hmap map[any]any, key any) (val any, pres bool)
func mapassign1(mapType *byte, hmap map[any]any, key any, val any)
func mapassign2(mapType *byte, hmap map[any]any, key any, val any, pres bool)
func mapiterinit(mapType *byte, hmap map[any]any, hiter *any)
func mapdelete(mapType *byte, hmap map[any]any, key any)
func mapiternext(hiter *any)
func mapiter1(hiter *any) (key any)
func mapiter2(hiter *any) (key any, val any)

// *byte is really *runtime.Type
func makechan(chanType *byte, hint int64) (hchan chan any)
func chanrecv1(chanType *byte, hchan <-chan any) (elem any)
func chanrecv2(chanType *byte, hchan <-chan any) (elem any, received bool)
func chansend1(chanType *byte, hchan chan<- any, elem any)
func closechan(hchan any)

func selectnbsend(chanType *byte, hchan chan<- any, elem any) bool
func selectnbrecv(chanType *byte, elem *any, hchan <-chan any) bool
func selectnbrecv2(chanType *byte, elem *any, received *bool, hchan <-chan any) bool

func newselect(size int) (sel *byte)
func selectsend(sel *byte, hchan chan<- any, elem *any) (selected bool)
func selectrecv(sel *byte, hchan <-chan any, elem *any) (selected bool)
func selectrecv2(sel *byte, hchan <-chan any, elem *any, received *bool) (selected bool)
func selectdefault(sel *byte) (selected bool)
func selectgo(sel *byte)
func block()

func makeslice(typ *byte, nel int64, cap int64) (ary []any)
func growslice(typ *byte, old []any, n int64) (ary []any)
func sliceslice1(old []any, lb uint64, width uint64) (ary []any)
func sliceslice(old []any, lb uint64, hb uint64, width uint64) (ary []any)
func slicearray(old *any, nel uint64, lb uint64, hb uint64, width uint64) (ary []any)

func closure() // has args, but compiler fills in

func memequal(eq *bool, size uintptr, x, y *any)
func memequal8(eq *bool, size uintptr, x, y *any)
func memequal16(eq *bool, size uintptr, x, y *any)
func memequal32(eq *bool, size uintptr, x, y *any)
func memequal64(eq *bool, size uintptr, x, y *any)
func memequal128(eq *bool, size uintptr, x, y *any)

// only used on 32-bit
func int64div(int64, int64) int64
func uint64div(uint64, uint64) uint64
func int64mod(int64, int64) int64
func uint64mod(uint64, uint64) uint64
func float64toint64(float64) int64
func float64touint64(float64) uint64
func int64tofloat64(int64) float64
func uint64tofloat64(uint64) float64

func complex128div(num complex128, den complex128) (quo complex128)
<commit_msg>gc: correct comment in runtime.go<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// NOTE: If you change this file you must run "./mkbuiltin"
// to update builtin.c.  This is not done automatically
// to avoid depending on having a working compiler binary.

// +build ignore

package PACKAGE

// emitted by compiler, not referred to by go programs

func new(typ *byte) *any
func panicindex()
func panicslice()
func throwreturn()
func throwinit()
func panicwrap(string, string, string)

func panic(interface{})
func recover(*int32) interface{}

func printbool(bool)
func printfloat(float64)
func printint(int64)
func printuint(uint64)
func printcomplex(complex128)
func printstring(string)
func printpointer(any)
func printiface(any)
func printeface(any)
func printslice(any)
func printnl()
func printsp()
func goprintf()

// filled in by compiler: int n, string, string, ...
func concatstring()

// filled in by compiler: Type*, int n, Slice, ...
func append()
func appendslice(typ *byte, x any, y []any) any
func appendstr(typ *byte, x []byte, y string) []byte

func cmpstring(string, string) int
func slicestring(string, int, int) string
func slicestring1(string, int) string
func intstring(int64) string
func slicebytetostring([]byte) string
func slicerunetostring([]rune) string
func stringtoslicebyte(string) []byte
func stringtoslicerune(string) []rune
func stringiter(string, int) int
func stringiter2(string, int) (retk int, retv rune)
func copy(to any, fr any, wid uint32) int
func slicestringcopy(to any, fr any) int

// interface conversions
func convI2E(elem any) (ret any)
func convI2I(typ *byte, elem any) (ret any)
func convT2E(typ *byte, elem any) (ret any)
func convT2I(typ *byte, typ2 *byte, elem any) (ret any)

// interface type assertions  x.(T)
func assertE2E(typ *byte, iface any) (ret any)
func assertE2E2(typ *byte, iface any) (ret any, ok bool)
func assertE2I(typ *byte, iface any) (ret any)
func assertE2I2(typ *byte, iface any) (ret any, ok bool)
func assertE2T(typ *byte, iface any) (ret any)
func assertE2T2(typ *byte, iface any) (ret any, ok bool)
func assertI2E(typ *byte, iface any) (ret any)
func assertI2E2(typ *byte, iface any) (ret any, ok bool)
func assertI2I(typ *byte, iface any) (ret any)
func assertI2I2(typ *byte, iface any) (ret any, ok bool)
func assertI2T(typ *byte, iface any) (ret any)
func assertI2T2(typ *byte, iface any) (ret any, ok bool)

func ifaceeq(i1 any, i2 any) (ret bool)
func efaceeq(i1 any, i2 any) (ret bool)
func ifacethash(i1 any) (ret uint32)
func efacethash(i1 any) (ret uint32)

func equal(typ *byte, x1, x2 any) (ret bool)

// *byte is really *runtime.Type
func makemap(mapType *byte, hint int64) (hmap map[any]any)
func mapaccess1(mapType *byte, hmap map[any]any, key any) (val any)
func mapaccess2(mapType *byte, hmap map[any]any, key any) (val any, pres bool)
func mapassign1(mapType *byte, hmap map[any]any, key any, val any)
func mapassign2(mapType *byte, hmap map[any]any, key any, val any, pres bool)
func mapiterinit(mapType *byte, hmap map[any]any, hiter *any)
func mapdelete(mapType *byte, hmap map[any]any, key any)
func mapiternext(hiter *any)
func mapiter1(hiter *any) (key any)
func mapiter2(hiter *any) (key any, val any)

// *byte is really *runtime.Type
func makechan(chanType *byte, hint int64) (hchan chan any)
func chanrecv1(chanType *byte, hchan <-chan any) (elem any)
func chanrecv2(chanType *byte, hchan <-chan any) (elem any, received bool)
func chansend1(chanType *byte, hchan chan<- any, elem any)
func closechan(hchan any)

func selectnbsend(chanType *byte, hchan chan<- any, elem any) bool
func selectnbrecv(chanType *byte, elem *any, hchan <-chan any) bool
func selectnbrecv2(chanType *byte, elem *any, received *bool, hchan <-chan any) bool

func newselect(size int) (sel *byte)
func selectsend(sel *byte, hchan chan<- any, elem *any) (selected bool)
func selectrecv(sel *byte, hchan <-chan any, elem *any) (selected bool)
func selectrecv2(sel *byte, hchan <-chan any, elem *any, received *bool) (selected bool)
func selectdefault(sel *byte) (selected bool)
func selectgo(sel *byte)
func block()

func makeslice(typ *byte, nel int64, cap int64) (ary []any)
func growslice(typ *byte, old []any, n int64) (ary []any)
func sliceslice1(old []any, lb uint64, width uint64) (ary []any)
func sliceslice(old []any, lb uint64, hb uint64, width uint64) (ary []any)
func slicearray(old *any, nel uint64, lb uint64, hb uint64, width uint64) (ary []any)

func closure() // has args, but compiler fills in

func memequal(eq *bool, size uintptr, x, y *any)
func memequal8(eq *bool, size uintptr, x, y *any)
func memequal16(eq *bool, size uintptr, x, y *any)
func memequal32(eq *bool, size uintptr, x, y *any)
func memequal64(eq *bool, size uintptr, x, y *any)
func memequal128(eq *bool, size uintptr, x, y *any)

// only used on 32-bit
func int64div(int64, int64) int64
func uint64div(uint64, uint64) uint64
func int64mod(int64, int64) int64
func uint64mod(uint64, uint64) uint64
func float64toint64(float64) int64
func float64touint64(float64) uint64
func int64tofloat64(int64) float64
func uint64tofloat64(uint64) float64

func complex128div(num complex128, den complex128) (quo complex128)
<|endoftext|>


<commit_before>package datad

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/coreos/go-etcd/etcd"
)

var NodeMembershipTTL = 5 * time.Second

// A Node ensures that the provider's keys are registered and coordinates
// distribution of data among the other nodes in the cluster.
type Node struct {
	Name     string
	Provider Provider

	backend  Backend
	registry *Registry

	Log *log.Logger

	stopChan chan struct{}
}

// NewNode creates a new node to publish data from a provider to the cluster.
// The name ("host:port") is advertised to the cluster and therefore must be
// accessible by the other clients and nodes in the cluster. The name should be
// the host and port where the data on this machine is accessible.
//
// Call Start on this node to begin publishing its keys to the cluster.
func NewNode(name string, b Backend, p Provider) *Node {
	name = cleanNodeName(name)
	return &Node{
		Name:     name,
		Provider: p,
		backend:  b,
		registry: NewRegistry(b),
		Log:      log.New(os.Stderr, "", log.Ltime|log.Lmicroseconds|log.Lshortfile),
		stopChan: make(chan struct{}),
	}
}

func cleanNodeName(name string) string {
	name = strings.TrimPrefix(name, "http://")
	parseName := name
	if !strings.Contains(parseName, ":") {
		parseName += ":80"
	}
	_, _, err := net.SplitHostPort(parseName)
	if err != nil {
		panic("NewNode: bad name '" + name + "': " + err.Error() + " (name should be 'host:port')")
	}
	return name
}

// Start begins advertising this node's provider's keys to the
// cluster.
func (n *Node) Start() error {
	n.logf("Starting node %s.", n.Name)

	err := n.joinCluster()
	if err != nil {
		n.logf("Failed to join cluster: %s", err)
		return err
	}

	err = n.registerExistingKeys()
	if err != nil {
		n.logf("Failed to register existing keys: %s", err)
		return err
	}

	go n.watchRegisteredKeys()
	go n.balancePeriodically()

	return nil
}

// Stop deregisters this node's keys and stops background processes
// for this node.
func (n *Node) Stop() error {
	close(n.stopChan)
	return nil
}

// joinCluster adds this node's provider to the cluster, making it available to
// receive requests for and be assigned keys. It then periodically re-adds this
// node to the cluster before the TTL on the etcd cluster membership key
// elapses.
func (n *Node) joinCluster() error {
	err := n.refreshClusterMembership()
	if err != nil {
		return err
	}

	if NodeMembershipTTL < time.Second {
		panic("NodeMembershipTTL must be at least 1 second")
	}

	go func() {
		t := time.NewTicker(NodeMembershipTTL)
		for {
			select {
			case <-t.C:
				err := n.refreshClusterMembership()
				if err != nil {
					n.logf("Error refreshing node %s cluster membership: %s.", n.Name, err)
				}
			case <-n.stopChan:
				t.Stop()
				return
			}
		}
	}()

	return nil
}

func (n *Node) refreshClusterMembership() error {
	err := n.backend.SetDir(keyPathJoin(nodesPrefix, n.Name), uint64(NodeMembershipTTL/time.Second))
	if isEtcdErrorCode(err, 102) {
		err = n.backend.UpdateDir(keyPathJoin(nodesPrefix, n.Name), uint64(NodeMembershipTTL/time.Second))
	}
	return err
}

// watchRegisteredKeys watches the registry for changes to the list of keys that
// this node is registered for, or for modifications of existing registrations
// (e.g., updates requested).
func (n *Node) watchRegisteredKeys() error {
	watchKey := keysForNodeDir(n.Name)
	fullKey := n.backend.(*EtcdBackend).fullKey(watchKey)

	recv := make(chan *etcd.Response, 10)
	stopWatch := make(chan bool, 1)

	// Receive watched changes.
	go func() {
		for {
			select {
			case resp, ok := <-recv:
				if !ok {
					return
				}
				key := strings.TrimPrefix(resp.Node.Key, fullKey+"/")
				n.logf("Registry changed: %s on key %q.", resp.Action, key)
				if !strings.Contains(strings.ToLower(resp.Action), "delete") {
					n.logf("Updating key %q in data source (in response to registry %s).", key, resp.Action)
					err := n.Provider.Update(key)
					if err != nil {
						n.logf("Error updating key %q in data source: %s.", key, err)
					}
				}
			case <-n.stopChan:
				n.logf("Stopping registry watcher.")
				stopWatch <- true
				return
			}
		}
	}()

	_, err := n.backend.(*EtcdBackend).etcd.Watch(fullKey, 0, true, recv, stopWatch)
	if err != etcd.ErrWatchStoppedByUser {
		return err
	}
	return nil
}

// registerExistingKeys examines this node's provider's local storage for data
// and registers each data key it finds. This means that when the node starts
// up, it's immediately able to receive requests for the data it already has on
// disk. Without this, the cluster would not know that this node's provider has
// these keys.
func (n *Node) registerExistingKeys() error {
	keys, err := n.Provider.Keys("")
	if err != nil {
		return err
	}

	if len(keys) == 0 {
		return nil
	}

	n.logf("Found %d existing keys in provider: %v. Registering existing keys to this node...", len(keys), keys)
	for _, key := range keys {
		err := n.registry.Add(key, n.Name)
		if err != nil {
			return err
		}
	}
	n.logf("Finished registering existing %d keys to this node.", len(keys))

	return nil
}

// startBalancer starts a periodic process that balances the distribution of
// keys to nodes.
func (n *Node) balancePeriodically() {
	// TODO(sqs): allow tweaking the balance interval
	t := time.NewTicker(time.Minute)
	for {
		select {
		case <-t.C:
			err := n.balance()
			if err != nil {
				n.logf("Error balancing: %s. Will retry next balance interval.", err)
			}
		case <-n.stopChan:
			t.Stop()
			return
		}
	}
}

// balance examines all keys and ensures each key has a registered node. If not,
// it registers a node for the key. This lets the cluster heal itself after a
// node goes down (which causes keys to be orphaned).
func (n *Node) balance() error {
	keyMap, err := n.registry.KeyMap()
	if err != nil {
		return err
	}

	if len(keyMap) == 0 {
		return nil
	}

	c := NewClient(n.backend)
	clusterNodes, err := c.NodesInCluster()
	if err != nil {
		return err
	}

	// TODO(sqs): allow tweaking this parameter
	x := rand.Intn(10)
	start := time.Now()

	n.logf("Balancer: starting on %d keys, with known cluster nodes %v.", len(keyMap), clusterNodes)
	actions := 0
	for key, nodes := range keyMap {
		if len(nodes) == 0 {
			regNode := clusterNodes[keyBucket(key, len(clusterNodes))]

			n.logf("Balancer: found unregistered key %q; registering it to node %s.", key, regNode)

			// TODO(sqs): optimize this by only adding if not exists, and then
			// seeing if it exists (to avoid potentially duplicating work).
			err := c.registry.Add(key, regNode)
			if err != nil {
				return err
			}

			actions++
			continue
		}

		// Check liveness of key on each node.
		for _, node := range nodes {
			t, err := c.transportForKey(key, nil, []string{node})
			if err != nil {
				return err
			}
			resp, err := (&http.Client{Transport: t}).Get(slash(key))
			if err != nil {
				actions++
				n.logf("Balancer: liveness check failed for key %q on node %s: %s. Client deregistered key from node.", key, node, err)
			}
			if resp != nil {
				resp.Body.Close()
			}
		}

		// Update keys on this node (but not each time, to avoid overloading the
		// origin servers).
		if x == 0 {
			for _, node := range nodes {
				if node == n.Name {
					n.logf("Balancer: updating key %q in data source on current node.", key)
					err := n.Provider.Update(key)
					if err != nil {
						n.logf("Balancer: error updating key %q in data source on current node: %s. (Continuing to balance other keys.)", key, err)
					}
					actions++
				}
			}
		}
	}
	n.logf("Balancer: completed in %s for %d keys (%d non-read actions performed).", time.Since(start), len(keyMap), actions)

	return nil
}

func (n *Node) logf(format string, a ...interface{}) {
	if n.Log != nil {
		n.Log.Printf(fmt.Sprintf("Node %s: ", n.Name)+format, a...)
	}
}
<commit_msg>concurrent updates<commit_after>package datad

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/coreos/go-etcd/etcd"
)

var (
	// NodeMembershipTTL is the time-to-live of the etcd key that denotes a
	// node's membership in the cluster.
	NodeMembershipTTL = 5 * time.Second

	// BalanceInterval is the time interval for starting a balancing job on the
	// whole keyspace on each node.
	BalanceInterval = time.Minute
)

// A Node ensures that the provider's keys are registered and coordinates
// distribution of data among the other nodes in the cluster.
type Node struct {
	Name     string
	Provider Provider

	// Updaters is the maximum number of concurrent calls to Provider.Update
	// that may be executing at any given time on this node.
	Updaters  int
	updateQ   chan string
	updateQMu sync.Mutex

	backend  Backend
	registry *Registry

	Log *log.Logger

	stopChan chan struct{}
}

// NewNode creates a new node to publish data from a provider to the cluster.
// The name ("host:port") is advertised to the cluster and therefore must be
// accessible by the other clients and nodes in the cluster. The name should be
// the host and port where the data on this machine is accessible.
//
// Call Start on this node to begin publishing its keys to the cluster.
func NewNode(name string, b Backend, p Provider) *Node {
	name = cleanNodeName(name)
	return &Node{
		Name:     name,
		Provider: p,
		Updaters: 1,
		updateQ:  make(chan string),
		backend:  b,
		registry: NewRegistry(b),
		Log:      log.New(os.Stderr, "", log.Ltime|log.Lmicroseconds|log.Lshortfile),
		stopChan: make(chan struct{}),
	}
}

func cleanNodeName(name string) string {
	name = strings.TrimPrefix(name, "http://")
	parseName := name
	if !strings.Contains(parseName, ":") {
		parseName += ":80"
	}
	_, _, err := net.SplitHostPort(parseName)
	if err != nil {
		panic("NewNode: bad name '" + name + "': " + err.Error() + " (name should be 'host:port')")
	}
	return name
}

// Start begins advertising this node's provider's keys to the
// cluster.
func (n *Node) Start() error {
	n.logf("Starting node %s.", n.Name)

	err := n.joinCluster()
	if err != nil {
		n.logf("Failed to join cluster: %s", err)
		return err
	}

	err = n.registerExistingKeys()
	if err != nil {
		n.logf("Failed to register existing keys: %s", err)
		return err
	}

	go n.watchRegisteredKeys()
	go n.balancePeriodically()
	go n.startUpdater()

	return nil
}

// Stop deregisters this node's keys and stops background processes
// for this node.
func (n *Node) Stop() error {
	close(n.stopChan)
	return nil
}

// joinCluster adds this node's provider to the cluster, making it available to
// receive requests for and be assigned keys. It then periodically re-adds this
// node to the cluster before the TTL on the etcd cluster membership key
// elapses.
func (n *Node) joinCluster() error {
	err := n.refreshClusterMembership()
	if err != nil {
		return err
	}

	if NodeMembershipTTL < time.Second {
		panic("NodeMembershipTTL must be at least 1 second")
	}

	go func() {
		t := time.NewTicker(NodeMembershipTTL)
		for {
			select {
			case <-t.C:
				err := n.refreshClusterMembership()
				if err != nil {
					n.logf("Error refreshing node %s cluster membership: %s.", n.Name, err)
				}
			case <-n.stopChan:
				t.Stop()
				return
			}
		}
	}()

	return nil
}

func (n *Node) refreshClusterMembership() error {
	err := n.backend.SetDir(keyPathJoin(nodesPrefix, n.Name), uint64(NodeMembershipTTL/time.Second))
	if isEtcdErrorCode(err, 102) {
		err = n.backend.UpdateDir(keyPathJoin(nodesPrefix, n.Name), uint64(NodeMembershipTTL/time.Second))
	}
	return err
}

// watchRegisteredKeys watches the registry for changes to the list of keys that
// this node is registered for, or for modifications of existing registrations
// (e.g., updates requested).
func (n *Node) watchRegisteredKeys() error {
	watchKey := keysForNodeDir(n.Name)
	fullKey := n.backend.(*EtcdBackend).fullKey(watchKey)

	recv := make(chan *etcd.Response, 10)
	stopWatch := make(chan bool, 1)

	// Receive watched changes.
	go func() {
		for {
			select {
			case resp, ok := <-recv:
				if !ok {
					return
				}
				key := strings.TrimPrefix(resp.Node.Key, fullKey+"/")
				n.logf("Registry changed: %s on key %q.", resp.Action, key)
				if !strings.Contains(strings.ToLower(resp.Action), "delete") {
					n.logf("Queueing update for key %q in data source (in response to registry %s).", key, resp.Action)
					n.updateQ <- key
				}
			case <-n.stopChan:
				n.logf("Stopping registry watcher.")
				stopWatch <- true
				return
			}
		}
	}()

	_, err := n.backend.(*EtcdBackend).etcd.Watch(fullKey, 0, true, recv, stopWatch)
	if err != etcd.ErrWatchStoppedByUser {
		return err
	}
	return nil
}

// registerExistingKeys examines this node's provider's local storage for data
// and registers each data key it finds. This means that when the node starts
// up, it's immediately able to receive requests for the data it already has on
// disk. Without this, the cluster would not know that this node's provider has
// these keys.
func (n *Node) registerExistingKeys() error {
	keys, err := n.Provider.Keys("")
	if err != nil {
		return err
	}

	if len(keys) == 0 {
		return nil
	}

	n.logf("Found %d existing keys in provider: %v. Registering existing keys to this node...", len(keys), keys)
	for _, key := range keys {
		err := n.registry.Add(key, n.Name)
		if err != nil {
			return err
		}
	}
	n.logf("Finished registering existing %d keys to this node.", len(keys))

	return nil
}

func (n *Node) startUpdater() {
	keyToUpdate := make(chan string)

	// Use a map to avoid updating the same key concurrently. If a key's value
	// is false, it's queued; if true, its update is in progress.
	pending := make(map[string]bool)

	started := make(chan string)
	completed := make(chan string)

	// Consume queue and distribute keys to updaters.
	go func() {
		for {
			select {
			case key := <-completed:
				delete(pending, key)
			case key := <-started:
				pending[key] = true
			case key := <-n.updateQ:
				if _, isPending := pending[key]; !isPending {
					pending[key] = false
					go func() { keyToUpdate <- key }() // TODO(sqs): hacky
					if len(pending) > n.Updaters {
						n.logf("%d key updates pending: %v (false=queued, true=update in progress).", len(pending), pending)
					}
				}
			case <-n.stopChan:
				return
			}
		}
	}()

	if n.Updaters <= 0 {
		panic("invalid Updaters value for node (Updaters <= 0)")
	}

	// Updaters.
	for i := 0; i < n.Updaters; i++ {
		go func() {
			for {
				select {
				case key := <-keyToUpdate:
					started <- key
					err := n.Provider.Update(key)
					if err != nil {
						n.logf("Update failed for key %q: %s.", key, err)
					}
					completed <- key
				case <-n.stopChan:
					return
				}
			}
		}()
	}
}

// startBalancer starts a periodic process that balances the distribution of
// keys to nodes.
func (n *Node) balancePeriodically() {
	// TODO(sqs): allow tweaking the balance interval
	t := time.NewTicker(BalanceInterval)
	for {
		select {
		case <-t.C:
			err := n.balance()
			if err != nil {
				n.logf("Error balancing: %s. Will retry next balance interval.", err)
			}
		case <-n.stopChan:
			t.Stop()
			return
		}
	}
}

// balance examines all keys and ensures each key has a registered node. If not,
// it registers a node for the key. This lets the cluster heal itself after a
// node goes down (which causes keys to be orphaned).
func (n *Node) balance() error {
	keyMap, err := n.registry.KeyMap()
	if err != nil {
		return err
	}

	if len(keyMap) == 0 {
		return nil
	}

	c := NewClient(n.backend)
	clusterNodes, err := c.NodesInCluster()
	if err != nil {
		return err
	}

	// TODO(sqs): allow tweaking this parameter
	x := rand.Intn(10)
	start := time.Now()

	n.logf("Balancer: starting on %d keys, with known cluster nodes %v.", len(keyMap), clusterNodes)
	actions := 0
	iterations := 0
	for key, nodes := range keyMap {
		if maxDuration := BalanceInterval / 2; time.Since(start) > maxDuration {
			// Rough heuristic that the KeyMap is stale, so let's end this
			// balance run. The next balance run won't have to redo the work we
			// did.
			n.logf("Balancer: ending before complete because max balance duration %s elapsed. Checked %d/%d entries in KeyMap, %d non-read ops. Will resume in next balance run.", maxDuration, iterations, len(keyMap), actions)
			return nil
		}

		iterations++

		if len(nodes) == 0 {
			regNode := clusterNodes[keyBucket(key, len(clusterNodes))]

			n.logf("Balancer: found unregistered key %q; registering it to node %s.", key, regNode)

			// TODO(sqs): optimize this by only adding if not exists, and then
			// seeing if it exists (to avoid potentially duplicating work).
			err := c.registry.Add(key, regNode)
			if err != nil {
				return err
			}

			actions++
			continue
		}

		// Check liveness of key on each node.
		for _, node := range nodes {
			t, err := c.transportForKey(key, nil, []string{node})
			if err != nil {
				return err
			}
			resp, err := (&http.Client{Transport: t}).Get(slash(key))
			if err != nil {
				actions++
				n.logf("Balancer: liveness check failed for key %q on node %s: %s. Client deregistered key from node.", key, node, err)
			}
			if resp != nil {
				resp.Body.Close()
			}
		}

		// Update keys on this node (but not each time, to avoid overloading the
		// origin servers).
		if x == 0 {
			for _, node := range nodes {
				if node == n.Name {
					n.logf("Balancer: queueing update for key %q in data source on current node.", key)
					n.updateQ <- key
					actions++
				}
			}
		}
	}
	n.logf("Balancer: completed in %s for %d keys (%d non-read actions performed).", time.Since(start), len(keyMap), actions)

	return nil
}

func (n *Node) logf(format string, a ...interface{}) {
	if n.Log != nil {
		n.Log.Printf(fmt.Sprintf("Node %s: ", n.Name)+format, a...)
	}
}
<|endoftext|>


<commit_before>package networkpolicy

import (
	"fmt"

	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"

	"github.com/rancher/norman/types/convert"
	"github.com/rancher/rancher/pkg/controllers/managementuser/nodesyncer"
	v1 "github.com/rancher/rancher/pkg/generated/norman/core/v1"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/types/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

type clusterHandler struct {
	cluster          *config.UserContext
	pnpLister        v3.ProjectNetworkPolicyLister
	podLister        v1.PodLister
	serviceLister    v1.ServiceLister
	pLister          v3.ProjectLister
	clusters         v3.ClusterInterface
	pnps             v3.ProjectNetworkPolicyInterface
	npmgr            *netpolMgr
	clusterNamespace string
}

/*
clusterHandler enqueues resources for creating/deleting network policies
based on cluster.Annotations[netPolAnnotation] and sets status if successful
*/

func (ch *clusterHandler) Sync(key string, cluster *v3.Cluster) (runtime.Object, error) {
	if cluster == nil || cluster.DeletionTimestamp != nil ||
		cluster.Name != ch.clusterNamespace ||
		!v32.ClusterConditionReady.IsTrue(cluster) {
		return nil, nil
	}

	if cluster.Spec.EnableNetworkPolicy == nil {
		return nil, nil
	}

	toEnable := convert.ToBool(cluster.Annotations[netPolAnnotation])

	if cluster.Status.AppliedEnableNetworkPolicy == toEnable {
		return nil, nil
	}

	if toEnable != *cluster.Spec.EnableNetworkPolicy {
		// allow clusterNetAnnHandler to update first
		return nil, nil
	}

	var err error
	if toEnable {
		logrus.Infof("clusterHandler: calling sync to create network policies for cluster %v", cluster.Name)
		err = ch.createNetworkPolicies(cluster)
	} else {
		logrus.Infof("clusterHandler: deleting network policies for cluster %s", cluster.Name)
		err = ch.deleteNetworkPolicies(cluster)
	}

	if err != nil {
		return nil, err
	}

	cluster.Status.AppliedEnableNetworkPolicy = toEnable

	_, err = ch.clusters.Update(cluster)
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (ch *clusterHandler) createNetworkPolicies(cluster *v3.Cluster) error {
	projects, err := ch.pLister.List(cluster.Name, labels.NewSelector())
	if err != nil {
		return fmt.Errorf("projectLister: %v", err)
	}

	for _, project := range projects {
		ch.npmgr.projects.Controller().Enqueue(project.Namespace, project.Name)
	}

	systemNamespaces, _, err := ch.npmgr.getSystemNSInfo(cluster.Name)
	if err != nil {
		return fmt.Errorf("systemNS: %v", err)
	}

	//hostPort
	pods, err := ch.podLister.List("", labels.NewSelector())
	if err != nil {
		return fmt.Errorf("podLister: %v", err)
	}

	for _, pod := range pods {
		if systemNamespaces[pod.Namespace] {
			continue
		}
		if hostPortPod(pod) {
			ch.cluster.Core.Pods("").Controller().Enqueue(pod.Namespace, pod.Name)
		}
	}

	// nodePort
	svcs, err := ch.serviceLister.List("", labels.NewSelector())
	if err != nil {
		return err
	}
	for _, svc := range svcs {
		if systemNamespaces[svc.Namespace] {
			continue
		}
		if nodePortService(svc) {
			ch.cluster.Core.Services("").Controller().Enqueue(svc.Namespace, svc.Name)
		}
	}

	ch.cluster.Management.Management.Nodes(ch.cluster.ClusterName).Controller().Enqueue(
		cluster.ClusterName, fmt.Sprintf("%s/%s", ch.cluster.ClusterName, nodesyncer.AllNodeKey))

	return nil
	//skipping nssyncer, projectSyncer + nodehandler would result into handling nssyncer as well
}

func (ch *clusterHandler) deleteNetworkPolicies(cluster *v3.Cluster) error {
	nps, err := ch.npmgr.npLister.List("", labels.NewSelector())
	if err != nil {
		return fmt.Errorf("npLister: %v", err)
	}
	for _, np := range nps {
		if err := ch.npmgr.delete(np.Namespace, np.Name); err != nil {
			return fmt.Errorf("npDelete: %v", err)
		}
	}

	projects, err := ch.pLister.List(cluster.Name, labels.NewSelector())
	if err != nil {
		return fmt.Errorf("projectLister: %v", err)
	}

	for _, project := range projects {
		pnps, err := ch.pnpLister.List(project.Name, labels.NewSelector())
		if err != nil {
			return fmt.Errorf("pnpLister: %v", err)
		}

		for _, pnp := range pnps {
			err := ch.pnps.DeleteNamespaced(pnp.Namespace, pnp.Name, &metav1.DeleteOptions{})
			if err != nil {
				return fmt.Errorf("pnpDelete: %v", err)
			}
		}
	}
	return nil
}
<commit_msg>only remove rancher created network policies when disabling PNI<commit_after>package networkpolicy

import (
	"fmt"

	"github.com/rancher/norman/types/convert"
	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/controllers/managementuser/nodesyncer"
	v1 "github.com/rancher/rancher/pkg/generated/norman/core/v1"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/types/config"
	"github.com/sirupsen/logrus"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
)

type clusterHandler struct {
	cluster          *config.UserContext
	pnpLister        v3.ProjectNetworkPolicyLister
	podLister        v1.PodLister
	serviceLister    v1.ServiceLister
	pLister          v3.ProjectLister
	clusters         v3.ClusterInterface
	pnps             v3.ProjectNetworkPolicyInterface
	npmgr            *netpolMgr
	clusterNamespace string
}

/*
clusterHandler enqueues resources for creating/deleting network policies
based on cluster.Annotations[netPolAnnotation] and sets status if successful
*/

func (ch *clusterHandler) Sync(key string, cluster *v3.Cluster) (runtime.Object, error) {
	if cluster == nil || cluster.DeletionTimestamp != nil ||
		cluster.Name != ch.clusterNamespace ||
		!v32.ClusterConditionReady.IsTrue(cluster) {
		return nil, nil
	}

	if cluster.Spec.EnableNetworkPolicy == nil {
		return nil, nil
	}

	toEnable := convert.ToBool(cluster.Annotations[netPolAnnotation])

	if cluster.Status.AppliedEnableNetworkPolicy == toEnable {
		return nil, nil
	}

	if toEnable != *cluster.Spec.EnableNetworkPolicy {
		// allow clusterNetAnnHandler to update first
		return nil, nil
	}

	var err error
	if toEnable {
		logrus.Infof("clusterHandler: calling sync to create network policies for cluster %v", cluster.Name)
		err = ch.createNetworkPolicies(cluster)
	} else {
		logrus.Infof("clusterHandler: deleting network policies for cluster %s", cluster.Name)
		err = ch.deleteNetworkPolicies(cluster)
	}

	if err != nil {
		return nil, err
	}

	cluster.Status.AppliedEnableNetworkPolicy = toEnable

	_, err = ch.clusters.Update(cluster)
	if err != nil {
		return nil, err
	}

	return nil, nil
}

func (ch *clusterHandler) createNetworkPolicies(cluster *v3.Cluster) error {
	projects, err := ch.pLister.List(cluster.Name, labels.NewSelector())
	if err != nil {
		return fmt.Errorf("projectLister: %v", err)
	}

	for _, project := range projects {
		ch.npmgr.projects.Controller().Enqueue(project.Namespace, project.Name)
	}

	systemNamespaces, _, err := ch.npmgr.getSystemNSInfo(cluster.Name)
	if err != nil {
		return fmt.Errorf("systemNS: %v", err)
	}

	//hostPort
	pods, err := ch.podLister.List("", labels.NewSelector())
	if err != nil {
		return fmt.Errorf("podLister: %v", err)
	}

	for _, pod := range pods {
		if systemNamespaces[pod.Namespace] {
			continue
		}
		if hostPortPod(pod) {
			ch.cluster.Core.Pods("").Controller().Enqueue(pod.Namespace, pod.Name)
		}
	}

	// nodePort
	svcs, err := ch.serviceLister.List("", labels.NewSelector())
	if err != nil {
		return err
	}
	for _, svc := range svcs {
		if systemNamespaces[svc.Namespace] {
			continue
		}
		if nodePortService(svc) {
			ch.cluster.Core.Services("").Controller().Enqueue(svc.Namespace, svc.Name)
		}
	}

	ch.cluster.Management.Management.Nodes(ch.cluster.ClusterName).Controller().Enqueue(
		cluster.ClusterName, fmt.Sprintf("%s/%s", ch.cluster.ClusterName, nodesyncer.AllNodeKey))

	return nil
	//skipping nssyncer, projectSyncer + nodehandler would result into handling nssyncer as well
}

const creatorLabel = "cattle.io/creator"

// deleteNetworkPolicies removes Rancher created NetworkPolicy resources from the downstream cluster and
// removes ProjectNetworkPolicy resources from the management cluster
func (ch *clusterHandler) deleteNetworkPolicies(cluster *v3.Cluster) error {
	// consider nps for deletion if they were created by Rancher, i.e. they have a label: "cattle.io/creator": "norman"
	set := labels.Set(map[string]string{creatorLabel: "norman"})
	nps, err := ch.npmgr.npLister.List("", set.AsSelector())
	if err != nil {
		return fmt.Errorf("npLister: %v", err)
	}
	for _, np := range nps {
		if err := ch.npmgr.delete(np.Namespace, np.Name); err != nil {
			return fmt.Errorf("npDelete: %v", err)
		}
	}

	projects, err := ch.pLister.List(cluster.Name, labels.NewSelector())
	if err != nil {
		return fmt.Errorf("projectLister: %v", err)
	}

	for _, project := range projects {
		pnps, err := ch.pnpLister.List(project.Name, labels.NewSelector())
		if err != nil {
			return fmt.Errorf("pnpLister: %v", err)
		}

		for _, pnp := range pnps {
			err := ch.pnps.DeleteNamespaced(pnp.Namespace, pnp.Name, &metav1.DeleteOptions{})
			if err != nil {
				return fmt.Errorf("pnpDelete: %v", err)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package goenocean

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"

	"github.com/tarm/goserial"
)

type Encoder interface {
	Encode() []byte
}

type request struct {
	Encoder
	responseChannel chan Packet
}

type Request interface {
	Encoder
	ResponseChannel() chan Packet
}

func (r *request) ResponseChannel() chan Packet {
	return r.responseChannel
}

func NewRequest(p Encoder) Request {
	return &request{p, make(chan Packet)}
}

func Serial(send chan Request, recv chan Packet) {
	c := &serial.Config{Name: "/dev/ttyUSB0", Baud: 57600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatal(err)
	}

	response := make(chan Packet)

	go readPackets(s, func(data []byte) {
		reciever(data, recv, response)
	})

	go sender(s, send, response)
}

func sender(data io.ReadWriter, send chan Request, response chan Packet) {

	for r := range send {
		_, err := data.Write(r.Encode())
		response <- nil
		r.ResponseChannel() <- <-response //TODO test this. might work :)
		if err != nil {
			log.Fatal(err)
		}
	}

}

func reciever(data []byte, recv chan Packet, response chan Packet) {
	p, err := Decode(data)
	fmt.Printf("%#v\n", p)
	fmt.Printf("%#v\n", p.Header())
	fmt.Printf("Data: %#v\n", p.Data())
	if err != nil {
		fmt.Println(err)
		return
	}

	select {
	case <-response:
		response <- p
	default:
		recv <- p
	}
}

func readPackets(rd io.ReadWriter, f func([]byte)) {

	buf := make([]byte, 1)
	var rawPacket []byte
	state := 0
	var length int

	for {
		readLen, err := rd.Read(buf)
		if err != nil {
			fmt.Println("ERROR reading:", err)
			continue
		}

		//TODO add debug here seelog
		//fmt.Printf("% x ", buf)
		//continue

		if readLen > 0 && buf[0] == 0x55 {
			rawPacket = []byte{}
			state = 1
		}

		rawPacket = append(rawPacket, buf...)

		switch state {
		case 1: //header
			if len(rawPacket) > 5 {
				length = int(binary.BigEndian.Uint16(rawPacket[1:3]))
				state = 2
			}

		case 2: // the rest!
			if len(rawPacket) > 5+length+int(rawPacket[3]) {
				state = 3
			}

		case 3: //data crc
			state = 0
			f(rawPacket)
		}

	}

}
<commit_msg>added a experimental way to get response from a send<commit_after>package goenocean

import (
	"encoding/binary"
	"fmt"
	"io"
	"log"

	"github.com/tarm/goserial"
)

type Encoder interface {
	Encode() []byte
}

type request struct {
	Encoder
	responseChannel chan Packet
}

type Request interface {
	Encoder
	ResponseChannel() chan Packet
}

func (r *request) ResponseChannel() chan Packet {
	return r.responseChannel
}

func NewRequest(p Encoder) Request {
	return &request{p, make(chan Packet)}
}

func Serial(send chan Request, recv chan Packet) {
	c := &serial.Config{Name: "/dev/ttyUSB0", Baud: 57600}
	s, err := serial.OpenPort(c)
	if err != nil {
		log.Fatal(err)
	}

	response := make(chan Packet)
	waitForResponse := make(chan bool)

	go readPackets(s, func(data []byte) {
		reciever(data, recv, response, waitForResponse)
	})

	go sender(s, send, response, waitForResponse)
}

func sender(data io.ReadWriter, send chan Request, response chan Packet, waitForResponse chan bool) {

	for r := range send {
		_, err := data.Write(r.Encode())
		waitForResponse <- true
		r.ResponseChannel() <- <-response //TODO test this. might work :)
		if err != nil {
			log.Fatal(err)
		}
	}

}

func reciever(data []byte, recv chan Packet, response chan Packet, waitForResponse chan bool) {
	p, err := Decode(data)
	fmt.Printf("%#v\n", p)
	fmt.Printf("%#v\n", p.Header())
	fmt.Printf("Data: %#v\n", p.Data())
	if err != nil {
		fmt.Println(err)
		return
	}

	select {
	case <-waitForResponse:
		response <- p
	default:
		recv <- p
	}
}

func readPackets(rd io.ReadWriter, f func([]byte)) {

	buf := make([]byte, 1)
	var rawPacket []byte
	state := 0
	var length int

	for {
		readLen, err := rd.Read(buf)
		if err != nil {
			fmt.Println("ERROR reading:", err)
			continue
		}

		//TODO add debug here seelog
		//fmt.Printf("% x ", buf)
		//continue

		if readLen > 0 && buf[0] == 0x55 {
			rawPacket = []byte{}
			state = 1
		}

		rawPacket = append(rawPacket, buf...)

		switch state {
		case 1: //header
			if len(rawPacket) > 5 {
				length = int(binary.BigEndian.Uint16(rawPacket[1:3]))
				state = 2
			}

		case 2: // the rest!
			if len(rawPacket) > 5+length+int(rawPacket[3]) {
				state = 3
			}

		case 3: //data crc
			state = 0
			f(rawPacket)
		}

	}

}
<|endoftext|>


<commit_before>package manager

import (
	"fmt"

	log "github.com/Dataman-Cloud/swan/src/context_logger"
	"github.com/Dataman-Cloud/swan/src/manager/apiserver"
	"github.com/Dataman-Cloud/swan/src/manager/framework"
	"github.com/Dataman-Cloud/swan/src/manager/ipam"
	"github.com/Dataman-Cloud/swan/src/manager/ns"
	"github.com/Dataman-Cloud/swan/src/manager/raft"
	"github.com/Dataman-Cloud/swan/src/manager/sched"
	"github.com/Dataman-Cloud/swan/src/manager/store"
	"github.com/Dataman-Cloud/swan/src/manager/swancontext"
	"github.com/Dataman-Cloud/swan/src/util"
	"github.com/boltdb/bolt"

	"github.com/Sirupsen/logrus"
	events "github.com/docker/go-events"
	"golang.org/x/net/context"
)

type Manager struct {
	store     *store.Store
	apiserver *apiserver.ApiServer
	//proxyserver

	ipamAdapter *ipam.IpamAdapter
	resolver    *ns.Resolver
	sched       *sched.Sched
	raftNode    *raft.Node
	CancelFunc  context.CancelFunc

	framework *framework.Framework

	swanContext *swancontext.SwanContext
	config      util.SwanConfig
}

func New(config util.SwanConfig, db *bolt.DB) (*Manager, error) {
	manager := &Manager{
		config: config,
	}

	raftNode, err := raft.NewNode(config.Raft, db)
	if err != nil {
		logrus.Errorf("init raft node failed. Error: %s", err.Error())
		return nil, err
	}
	manager.raftNode = raftNode

	store := store.NewManagerStore(db, raftNode)

	manager.swanContext = &swancontext.SwanContext{
		Config: config,
		Store:  store,
		ApiServer: apiserver.NewApiServer(manager.config.HttpListener.TCPAddr,
			manager.config.HttpListener.UnixAddr),
	}

	manager.swanContext.Config.IPAM.StorePath = fmt.Sprintf(manager.config.IPAM.StorePath+"ipam.db.%d", config.Raft.RaftId)
	manager.ipamAdapter, err = ipam.New(manager.swanContext)
	if err != nil {
		logrus.Errorf("init ipam adapter failed. Error: ", err.Error())
		return nil, err
	}

	manager.resolver = ns.New(manager.config.DNS)
	manager.sched = sched.New(manager.config.Scheduler, manager.swanContext)
	manager.framework, err = framework.New(manager.config)
	if err != nil {
		logrus.Errorf("init framework failed. Error: ", err.Error())
		return nil, err
	}

	return manager, nil
}

func (manager *Manager) Stop(cancel context.CancelFunc) {
	cancel()
	return
}

func (manager *Manager) Start(ctx context.Context) error {
	errCh := make(chan error)
	leadershipCh, QueueCancel := manager.raftNode.SubscribeLeadership()
	defer QueueCancel()

	eventCtx, _ := context.WithCancel(ctx)
	go manager.handleLeadershipEvents(eventCtx, leadershipCh)

	raftCtx, _ := context.WithCancel(ctx)
	if err := manager.raftNode.StartRaft(raftCtx); err != nil {
		return err
	}

	if err := manager.raftNode.WaitForLeader(ctx); err != nil {
		return err
	}

	go func() {
		resolverCtx, _ := context.WithCancel(ctx)
		errCh <- manager.resolver.Start(resolverCtx)
	}()

	go func() {
		if err := manager.ipamAdapter.Start(); err != nil {
			errCh <- err
			return
		}

		managerRoute := NewRouter(manager)
		manager.swanContext.ApiServer.AppendRouter(managerRoute)

		errCh <- manager.swanContext.ApiServer.ListenAndServe()
	}()

	return <-errCh
}

func (manager *Manager) handleLeadershipEvents(ctx context.Context, leadershipCh chan events.Event) {
	for {
		select {
		case leadershipEvent := <-leadershipCh:
			// TODO lock it and if manager stop return
			newState := leadershipEvent.(raft.LeadershipState)

			var cancelFunc context.CancelFunc
			ctx = log.WithLogger(ctx, logrus.WithField("raft_id", fmt.Sprintf("%x", manager.config.Raft.RaftId)))
			if newState == raft.IsLeader {
				log.G(ctx).Info("Now i become a leader !!!")
				if manager.config.WithEngine == "sched" {
					sechedCtx, cancel := context.WithCancel(ctx)
					cancelFunc = cancel
					if err := manager.sched.Start(sechedCtx); err != nil {
						log.G(ctx).Error("Scheduler started unsuccessful")
						return
					}
					log.G(ctx).Info("Scheduler has been started")
				} else {
				}

			} else if newState == raft.IsFollower {
				log.G(ctx).Info("Now i become a follower !!!")
				if cancelFunc != nil {
					cancelFunc()
					log.G(ctx).Info("Scheduler has been stopped")
				}
				cancelFunc = nil
			}
		case <-ctx.Done():
			return
		}
	}
}
<commit_msg>enable sched and disable framework<commit_after>package manager

import (
	"fmt"

	log "github.com/Dataman-Cloud/swan/src/context_logger"
	"github.com/Dataman-Cloud/swan/src/manager/apiserver"
	"github.com/Dataman-Cloud/swan/src/manager/framework"
	"github.com/Dataman-Cloud/swan/src/manager/ipam"
	"github.com/Dataman-Cloud/swan/src/manager/ns"
	"github.com/Dataman-Cloud/swan/src/manager/raft"
	"github.com/Dataman-Cloud/swan/src/manager/sched"
	"github.com/Dataman-Cloud/swan/src/manager/store"
	"github.com/Dataman-Cloud/swan/src/manager/swancontext"
	"github.com/Dataman-Cloud/swan/src/util"
	"github.com/boltdb/bolt"

	"github.com/Sirupsen/logrus"
	events "github.com/docker/go-events"
	"golang.org/x/net/context"
)

type Manager struct {
	store     *store.Store
	apiserver *apiserver.ApiServer
	//proxyserver

	ipamAdapter *ipam.IpamAdapter
	resolver    *ns.Resolver
	sched       *sched.Sched
	raftNode    *raft.Node
	CancelFunc  context.CancelFunc

	framework *framework.Framework

	swanContext *swancontext.SwanContext
	config      util.SwanConfig
}

func New(config util.SwanConfig, db *bolt.DB) (*Manager, error) {
	manager := &Manager{
		config: config,
	}

	raftNode, err := raft.NewNode(config.Raft, db)
	if err != nil {
		logrus.Errorf("init raft node failed. Error: %s", err.Error())
		return nil, err
	}
	manager.raftNode = raftNode

	store := store.NewManagerStore(db, raftNode)

	manager.swanContext = &swancontext.SwanContext{
		Config: config,
		Store:  store,
		ApiServer: apiserver.NewApiServer(manager.config.HttpListener.TCPAddr,
			manager.config.HttpListener.UnixAddr),
	}

	manager.swanContext.Config.IPAM.StorePath = fmt.Sprintf(manager.config.IPAM.StorePath+"ipam.db.%d", config.Raft.RaftId)
	manager.ipamAdapter, err = ipam.New(manager.swanContext)
	if err != nil {
		logrus.Errorf("init ipam adapter failed. Error: ", err.Error())
		return nil, err
	}

	manager.resolver = ns.New(manager.config.DNS)
	manager.sched = sched.New(manager.config.Scheduler, manager.swanContext)
	manager.framework, err = framework.New(manager.config)
	if err != nil {
		logrus.Errorf("init framework failed. Error: ", err.Error())
		return nil, err
	}

	return manager, nil
}

func (manager *Manager) Stop(cancel context.CancelFunc) {
	cancel()
	return
}

func (manager *Manager) Start(ctx context.Context) error {
	errCh := make(chan error)
	leadershipCh, QueueCancel := manager.raftNode.SubscribeLeadership()
	defer QueueCancel()

	eventCtx, _ := context.WithCancel(ctx)
	go manager.handleLeadershipEvents(eventCtx, leadershipCh)

	raftCtx, _ := context.WithCancel(ctx)
	if err := manager.raftNode.StartRaft(raftCtx); err != nil {
		return err
	}

	if err := manager.raftNode.WaitForLeader(ctx); err != nil {
		return err
	}

	go func() {
		resolverCtx, _ := context.WithCancel(ctx)
		errCh <- manager.resolver.Start(resolverCtx)
	}()

	go func() {
		if err := manager.ipamAdapter.Start(); err != nil {
			errCh <- err
			return
		}

		managerRoute := NewRouter(manager)
		manager.swanContext.ApiServer.AppendRouter(managerRoute)

		errCh <- manager.swanContext.ApiServer.ListenAndServe()
	}()
	//go func() {
	//manager.framework.Start()
	//}()

	return <-errCh
}

func (manager *Manager) handleLeadershipEvents(ctx context.Context, leadershipCh chan events.Event) {
	for {
		select {
		case leadershipEvent := <-leadershipCh:
			// TODO lock it and if manager stop return
			newState := leadershipEvent.(raft.LeadershipState)

			var cancelFunc context.CancelFunc
			ctx = log.WithLogger(ctx, logrus.WithField("raft_id", fmt.Sprintf("%x", manager.config.Raft.RaftId)))
			if newState == raft.IsLeader {
				log.G(ctx).Info("Now i become a leader !!!")
				if manager.config.WithEngine == "sched" {
					sechedCtx, cancel := context.WithCancel(ctx)
					cancelFunc = cancel
					if err := manager.sched.Start(sechedCtx); err != nil {
						log.G(ctx).Error("Scheduler started unsuccessful")
						return
					}
					log.G(ctx).Info("Scheduler has been started")
				} else {
				}

			} else if newState == raft.IsFollower {
				log.G(ctx).Info("Now i become a follower !!!")
				if cancelFunc != nil {
					cancelFunc()
					log.G(ctx).Info("Scheduler has been stopped")
				}
				cancelFunc = nil
			}
		case <-ctx.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>// +build !no_gzip

package compress

import (
	"bytes"
	"compress/gzip"
	"github.com/xitongsys/parquet-go/parquet"
	"io/ioutil"
	"sync"
)

var gzipWriterPool sync.Pool

func init() {
	gzipWriterPool = sync.Pool{
		New: func() interface{} {
			return gzip.NewWriter(nil)
		},
	}

	compressors[parquet.CompressionCodec_GZIP] = &Compressor{
		Compress: func(buf []byte) []byte {
			res := new(bytes.Buffer)
			gzipWriter := gzipWriterPool.Get().(*gzip.Writer)
			gzipWriter.Reset(res)
			gzipWriter.Write(buf)
			gzipWriter.Close()
			gzipWriterPool.Put(gzipWriter)
			return res.Bytes()
		},
		Uncompress: func(buf []byte) (i []byte, err error) {
			rbuf := bytes.NewReader(buf)
			gzipReader, _ := gzip.NewReader(rbuf)
			res, err := ioutil.ReadAll(gzipReader)
			return res, err
		},
	}
}
<commit_msg>reset writer buffer reference before putting it back on the pool<commit_after>// +build !no_gzip

package compress

import (
	"bytes"
	"compress/gzip"
	"github.com/xitongsys/parquet-go/parquet"
	"io/ioutil"
	"sync"
)

var gzipWriterPool sync.Pool

func init() {
	gzipWriterPool = sync.Pool{
		New: func() interface{} {
			return gzip.NewWriter(nil)
		},
	}

	compressors[parquet.CompressionCodec_GZIP] = &Compressor{
		Compress: func(buf []byte) []byte {
			res := new(bytes.Buffer)
			gzipWriter := gzipWriterPool.Get().(*gzip.Writer)
			gzipWriter.Reset(res)
			gzipWriter.Write(buf)
			gzipWriter.Close()
			gzipWriter.Reset(nil)
			gzipWriterPool.Put(gzipWriter)
			return res.Bytes()
		},
		Uncompress: func(buf []byte) (i []byte, err error) {
			rbuf := bytes.NewReader(buf)
			gzipReader, _ := gzip.NewReader(rbuf)
			res, err := ioutil.ReadAll(gzipReader)
			return res, err
		},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        HoverflyStats `json:"stats"`
	RecordsCount int           `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

func (d *DBClient) startAdminInterface() {
	// starting admin interface
	mux := getBoneRouter(*d)
	n := negroni.Classic()

	loglevel := log.WarnLevel

	if d.cfg.verbose {
		loglevel = log.DebugLevel
	}

	n.Use(negronilogrus.NewCustomMiddleware(loglevel, &log.JSONFormatter{}, "admin"))
	n.UseHandler(mux)

	// admin interface starting message
	log.WithFields(log.Fields{
		"AdminPort": d.cfg.adminPort,
	}).Info("Admin interface is starting...")

	n.Run(fmt.Sprintf(":%s", d.cfg.adminPort))
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// preparing static assets for embedded admin
	statikFS, err := fs.New()
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to load statikFS, admin UI might not work :(")
	}

	mux.Get("/records", http.HandlerFunc(d.AllRecordsHandler))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.cfg.development {
		mux.Handle("/*", http.FileServer(http.Dir("static/dist")))
	} else {
		mux.Handle("/*", http.FileServer(statikFS))
	}

	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	records, err := d.cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.counter.Flush()

	count, err := d.cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	payloads := requests.Data
	if len(payloads) > 0 {
		for _, pl := range payloads {
			bts, err := pl.encode()
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to encode payload")
			} else {
				// recalculating request hash and storing it in database
				r := request{details: pl.Request}
				d.cache.Set([]byte(r.hash()), bts)
			}
		}
		response.Message = fmt.Sprintf("%d requests imported successfully", len(payloads))
	} else {
		response.Message = "Bad request. Nothing to import!"
		w.WriteHeader(400)
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.cache.DeleteBucket(d.cache.requestsBucket)

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.cfg.GetMode()
	resp.Destination = d.cfg.destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.cfg.SetMode(sr.Mode)

	var resp stateRequest
	resp.Mode = d.cfg.GetMode()
	resp.Destination = d.cfg.destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<commit_msg>refactoring to work as a library, also exporting configuration variables<commit_after>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	// starting admin interface
	mux := getBoneRouter(*d)
	n := negroni.Classic()

	logLevel := log.WarnLevel

	if d.Cfg.Verbose {
		logLevel = log.DebugLevel
	}

	n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
	n.UseHandler(mux)

	// admin interface starting message
	log.WithFields(log.Fields{
		"AdminPort": d.Cfg.AdminPort,
	}).Info("Admin interface is starting...")

	n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// preparing static assets for embedded admin
	statikFS, err := fs.New()
	if err != nil {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to load statikFS, admin UI might not work :(")
	}

	mux.Get("/records", http.HandlerFunc(d.AllRecordsHandler))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		mux.Handle("/*", http.FileServer(http.Dir("static/dist")))
	} else {
		mux.Handle("/*", http.FileServer(statikFS))
	}

	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	payloads := requests.Data
	if len(payloads) > 0 {
		for _, pl := range payloads {
			bts, err := pl.encode()
			if err != nil {
				log.WithFields(log.Fields{
					"error": err.Error(),
				}).Error("Failed to encode payload")
			} else {
				// recalculating request hash and storing it in database
				r := RequestContainer{details: pl.Request}
				d.Cache.Set([]byte(r.Hash()), bts)
			}
		}
		response.Message = fmt.Sprintf("%d requests imported successfully", len(payloads))
	} else {
		response.Message = "Bad request. Nothing to import!"
		w.WriteHeader(400)
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteBucket(d.Cache.RequestsBucket)

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<|endoftext|>


<commit_before>package main

import "os"

func main() {
	logger := Logger{"Main", os.Stdout}
	// Check if env vars are set, use defaults if not (or exit if needed)
	// and initialize oauth token.
	err := SetupClientData()
	if err != nil {
		logger.Log(err.Error())
		os.Exit(1)
	}

	err = EnsureGit(logger)
	if err != nil {
		logger.Log(err.Error())
		os.Exit(1)
	}

	jobsChannel := make(chan *TestJob)
	reportsChannel := make(chan *TestJob)

	manager := NewManager(jobsChannel)
	worker := NewWorker(jobsChannel, reportsChannel)
	reporter := NewReporter(reportsChannel)

	go worker.Start()
	go reporter.Start()
	manager.Start()
}
<commit_msg>Print logo on startup<commit_after>package main

import "os"

func main() {
	logger := Logger{"Main", os.Stdout}

	printLogo(logger)

	// Check if env vars are set, use defaults if not (or exit if needed)
	// and initialize oauth token.
	err := SetupClientData()
	if err != nil {
		logger.Log(err.Error())
		os.Exit(1)
	}

	err = EnsureGit(logger)
	if err != nil {
		logger.Log(err.Error())
		os.Exit(1)
	}

	jobsChannel := make(chan *TestJob)
	reportsChannel := make(chan *TestJob)

	manager := NewManager(jobsChannel)
	worker := NewWorker(jobsChannel, reportsChannel)
	reporter := NewReporter(reportsChannel)

	go worker.Start()
	go reporter.Start()
	manager.Start()
}

// Because we can
func printLogo(logger Logger) {
	logger.Log(`
 _______________________  _______  __  ____________  ___
/_  __/ __/ __/_  __/ _ \/  _/ _ )/ / / /_  __/ __ \/ _ \
 / / / _/_\ \  / / / , _// // _  / /_/ / / / / /_/ / , _/
/_/ /___/___/ /_/ /_/|_/___/____/\____/ /_/  \____/_/|_|

                              https://www.testributor.com
`)
}
<|endoftext|>


<commit_before>package diste

import (
	"errors"
	"sync"
)

type Enum string

const (
	OFF      Enum = "OFF"
	START    Enum = "START"
	ON       Enum = "ON"
	SHUTDOWN Enum = "SHUTDOWN"
)

type Agent struct {
	State    Enum
	Hostname string
}

func NewAgent(hostname string) *Agent {
	return &Agent{
		State:    OFF,
		Hostname: hostname,
	}
}

type AgentService struct {
	Agent *Agent
	Mutex *sync.RWMutex
}

func (s *AgentService) Ping(
	args ServiceRequest,
	reply *ServiceResponse) error {
	reply.Result = "Pong"
	return nil
}

func (s *AgentService) State(
	args ServiceRequest,
	reply *ServiceResponse) error {
	s.Mutex.RLock()
	defer s.Mutex.RUnlock()
	reply.Result = string(s.Agent.State)
	return nil
}

func (s *AgentService) Start(
	args ServiceRequest,
	reply *ServiceResponse) error {
	s.Mutex.RLock()
	defer s.Mutex.RUnlock()
	s.Agent.State = START
	reply.Result = string(s.Agent.State)
	return nil
}

func (s *AgentService) FakeError(
	args ServiceRequest,
	reply *ServiceResponse) error {

	return errors.New("System Error")
}
<commit_msg>Fixed write lock issue<commit_after>package diste

import (
	"errors"
	"sync"
)

type Enum string

const (
	OFF      Enum = "OFF"
	START    Enum = "START"
	ON       Enum = "ON"
	SHUTDOWN Enum = "SHUTDOWN"
)

type Agent struct {
	State    Enum
	Hostname string
}

func NewAgent(hostname string) *Agent {
	return &Agent{
		State:    OFF,
		Hostname: hostname,
	}
}

type AgentService struct {
	Agent *Agent
	Mutex *sync.RWMutex
}

func (s *AgentService) Ping(args ServiceRequest,
	reply *ServiceResponse) error {
	reply.Result = "Pong"
	return nil
}

func (s *AgentService) FakeError(args ServiceRequest,
	reply *ServiceResponse) error {

	return errors.New("System Error")
}

func (s *AgentService) State(args ServiceRequest,
	reply *ServiceResponse) error {
	s.Mutex.RLock()
	defer s.Mutex.RUnlock()
	reply.Result = string(s.Agent.State)
	return nil
}

func (s *AgentService) State(args ServiceRequest,
	reply *ServiceResponse) error {
	s.Mutex.RLock()
	defer s.Mutex.RUnlock()
	reply.Result = string(s.Agent.State)
	return nil
}

func (s *AgentService) Start(args ServiceRequest,
	reply *ServiceResponse) error {
	s.Mutex.Lock()
	defer s.Mutex.Unlock()
	s.Agent.State = START
	reply.Result = string(s.Agent.State)
	return nil
}
<|endoftext|>


<commit_before>package monitor

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	_ "io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"encoding/json"
)

type ipdata struct{
	Epoch		string `json:"Epoch"`
	SourceIP	string `json:"SourceIP"`
	SourcePort	string `json:"SourcePort"`
	ReceiveIP	string `json:"ReceiveIP"`
	ReceivePort string `json:"ReceivePort"`

}
/* The Monitor interface defines a series of methods which will be defined on
 * monitor structs. The Start method takes a channel to send messages over,
 * back to the configurator. The Stop method kills the process which is
 * performing the actual monitoring.
 */
type Monitor interface {
	Start(messages chan<- []byte, dockerComposeName string)
	Stop()
}

/* Store the MonitorName, which is the name of the program to execute, the
 * DockerDirs which are the directories to monitor which our Docker containers
 * of interest live in, and a pointer to the exec.Cmd struct which describes
 * the running command.
 */
type FSMonitor struct {
	MonitorName   string
	DockerDirs    []string
	fsWatcherProc *exec.Cmd
}

type NetMonitor struct {
	MonitorName  string
	ContainerIds []string
}

// Memoize this. It's kind of expensive to get.
var dockerContainerIds = []string{}

func runCommandAndSlurpOutput(commandname string, args []string) ([]string, error) {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return nil, err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	output := []string{}
	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return nil, err
			}
		}
		if len(line) > 0 {
			output = append(output, string(line))
		}
	}
	return output, nil
}

func runCommandAndChannelOutput(commandname string, args []string, output chan<- []byte) error {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return err
			}
		}
		if len(line) > 0 {
			output <- line
		}
	}
	return nil
}

func getDockerContainerIds(dockerComposeName string) []string {
	if len(dockerContainerIds) != 0 {
		return dockerContainerIds
	}
	ids, err := runCommandAndSlurpOutput("docker-compose", []string{"ps", "-q"})
	if err != nil {
		panic(err)
	}
	dockerContainerIds = ids
	return ids
}

func (n NetMonitor) Start(messages chan<- []byte, dockerComposeName string) {

	ids := getDockerContainerIds(dockerComposeName)
	arguments := []string{"inspect", "-f", "'{{ .State.Pid }}'"}
	arguments = append(arguments, ids...)
	output, err := runCommandAndSlurpOutput("docker", arguments)
	if err != nil {
		panic(err)
	}
	nmProcessorChan := make(chan []byte)
	go networkMonitorProcessor(messages, nmProcessorChan)
	for _, procId := range output {
		// FIXME: do we want to throw an error instead?
		if procId == "'0'" {
			fmt.Println("Container isn't running")
			continue
		}
		// Replace all quotes
		scrubbedProcId := strings.Replace(procId, "'", "", -1)
		// TODO: Check for ids named 0, which means no proc id
		err := setSymlink(scrubbedProcId, scrubbedProcId)
		if err != nil {
			panic(err)
		}
		go startIPProcess(nmProcessorChan, scrubbedProcId, "tcpdump","-tt","-n","-i","any")
	}
}
func networkMonitorProcessor(sending chan<-[]byte,receiving <-chan[]byte)(error){
	platform := []byte{}
	carlson := []byte{}
	pattern := "(\\d+\\.\\d+) IP (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+) > (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+)"
	compiled_pattern,err := regexp.Compile(pattern)
	if err != nil{
		return err
	}
	for{
		platform = <- receiving
		matches := compiled_pattern.FindSubmatch(platform)
		fmt.Println(string(matches[2]), string(matches[1]))
		carl := ipdata{
			Epoch: string(matches[1]),
			SourceIP: string(matches[2]),
			SourcePort: string(matches[3]),
			ReceiveIP: string(matches[4]),
			ReceivePort: string(matches[5]),
		}
		carlson,err = json.Marshal(carl)
		if err != nil{
			return err
		}
		fmt.Println("Sending json: ", string(carlson))
		sending<-carlson
	}

	return nil
}
func startIPProcess(messages chan<- []byte, procId string, watcherName string,
	watcherArgs ...string) {
	arguments := []string{"netns", "exec", procId, watcherName}
	arguments = append(arguments, watcherArgs...)
	err := runCommandAndChannelOutput("ip", arguments, messages)
	if err != nil {
		panic(err)
	}
}

func setSymlink(procId string, destination string) error {
	nameSpaceDir := "/var/run/netns/"
	target := nameSpaceDir+destination
	source := "/proc/"+procId+"/ns/net"

	err := os.MkdirAll(nameSpaceDir, 0777)
	if err != nil {
		return err
	}
	_, err = os.Stat(target)
	if err == nil {
		err = os.Remove(target)
		if err != nil {
			return err
		}
	}

	_, err = os.Stat(source)
	if os.IsNotExist(err) {
		return errors.New("Could not set symlink: Target directory " +
			target + " does not exist")
	}

	err = os.Symlink(source, target)
	return err
}

/* This function is going to need some comments describing why we chose this
 * approach, because it will be hairy.
 * FIXME: Break out functionality into multiple functions so we can test this
 * easily!
 */
func (m FSMonitor) getDockerFSDirectory(dockerComposeName string) []string {
	ids := getDockerContainerIds(dockerComposeName)
	/*
		dockerInfoCommand := exec.Command("docker", "info")
		infoOutPipe, err := dockerInfoCommand.StdoutPipe()

		if err != nil {
			panic(err)
		}

		dockerInfoCommand.Start()
		defer dockerInfoCommand.Wait()

		infoBuf, err := ioutil.ReadAll(infoOutPipe)

		if err != nil {
			panic(err)
		}

		info := string(infoBuf)

		re := regexp.MustCompile("Docker Root Dir: (.*)")

		submatch := re.FindStringSubmatch(info)

		if len(submatch) < 2 {
			fmt.Println(submatch, info)
			panic("Couldn't find the docker root directory")
		}
		dockerRootPath := submatch[1]
	*/
	dockerRootPath := "/var/lib/docker/aufs"

	for i := 0; i < len(ids); i++ {
		ids[i] = dockerRootPath + "/mnt/" + ids[i]
	}
	return ids
}

/* Start the process running on the honeypot host to monitor the Docker
 * container. The Docker container's filesystem is mounted on the host. Find
 * the location of this filesystem with the getDockerFSDirectory function and
 * store it in the struct. Then create and start the process and forward
 * the output of the process on to the messages channel.
 */
func (m FSMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	m.DockerDirs = m.getDockerFSDirectory(dockerComposeName)
	// FIXME Make arguments configurable
	//arguments := append([]string{"-wte"}, m.DockerDirs...)
	arguments := m.DockerDirs
	fmt.Println(m.MonitorName, arguments)
	m.fsWatcherProc = exec.Command(m.MonitorName, arguments...)
	defer m.fsWatcherProc.Wait()

	outpipe, err := m.fsWatcherProc.StdoutPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stdout pipe")
		panic(err)
	}
	stderr, err := m.fsWatcherProc.StderrPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stderr pipe")
		panic(err)
	}
	go io.Copy(os.Stderr, stderr)

	m.fsWatcherProc.Start()

	stdoutReader := bufio.NewReader(outpipe)

	for {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutReader.ReadLine()
			if err != nil {
				fmt.Println("File monitor stopped")
				panic(err)
			}
			fetch = f
			line = append(line, partial_line...)
		}
		messages <- line

	}
}

/* Stop the filesystem monitor. Kill the process monitoring the Docker
 * container's filesysem.
 */
func (m FSMonitor) Stop() {
	err := m.fsWatcherProc.Process.Kill()
	if err != nil {
		log.Fatal("Could not kill ", m.MonitorName, err)
	}
}
<commit_msg>Add -l option to tcpdump to view data in real time<commit_after>package monitor

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	_ "io/ioutil"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"encoding/json"
)

type ipdata struct{
	Epoch		string `json:"Epoch"`
	SourceIP	string `json:"SourceIP"`
	SourcePort	string `json:"SourcePort"`
	ReceiveIP	string `json:"ReceiveIP"`
	ReceivePort string `json:"ReceivePort"`

}
/* The Monitor interface defines a series of methods which will be defined on
 * monitor structs. The Start method takes a channel to send messages over,
 * back to the configurator. The Stop method kills the process which is
 * performing the actual monitoring.
 */
type Monitor interface {
	Start(messages chan<- []byte, dockerComposeName string)
	Stop()
}

/* Store the MonitorName, which is the name of the program to execute, the
 * DockerDirs which are the directories to monitor which our Docker containers
 * of interest live in, and a pointer to the exec.Cmd struct which describes
 * the running command.
 */
type FSMonitor struct {
	MonitorName   string
	DockerDirs    []string
	fsWatcherProc *exec.Cmd
}

type NetMonitor struct {
	MonitorName  string
	ContainerIds []string
}

// Memoize this. It's kind of expensive to get.
var dockerContainerIds = []string{}

func runCommandAndSlurpOutput(commandname string, args []string) ([]string, error) {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return nil, err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return nil, err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	output := []string{}
	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return nil, err
			}
		}
		if len(line) > 0 {
			output = append(output, string(line))
		}
	}
	return output, nil
}

func runCommandAndChannelOutput(commandname string, args []string, output chan<- []byte) error {
	command := exec.Command(commandname, args...)
	fmt.Print("running the command: ")
	fmt.Println(commandname, args)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, stderr)
	command.Start()
	defer command.Wait()

	stdoutreader := bufio.NewReader(stdout)
	slurp := true
	for slurp {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutreader.ReadLine()
			fetch = f
			line = append(line, partial_line...)
			if err == io.EOF {
				slurp = false
				break
			} else if err != nil {
				return err
			}
		}
		if len(line) > 0 {
			output <- line
		}
	}
	return nil
}

func getDockerContainerIds(dockerComposeName string) []string {
	if len(dockerContainerIds) != 0 {
		return dockerContainerIds
	}
	ids, err := runCommandAndSlurpOutput("docker-compose", []string{"ps", "-q"})
	if err != nil {
		panic(err)
	}
	dockerContainerIds = ids
	return ids
}

func (n NetMonitor) Start(messages chan<- []byte, dockerComposeName string) {

	ids := getDockerContainerIds(dockerComposeName)
	arguments := []string{"inspect", "-f", "'{{ .State.Pid }}'"}
	arguments = append(arguments, ids...)
	output, err := runCommandAndSlurpOutput("docker", arguments)
	if err != nil {
		panic(err)
	}
	nmProcessorChan := make(chan []byte)
	go networkMonitorProcessor(messages, nmProcessorChan)
	for _, procId := range output {
		// FIXME: do we want to throw an error instead?
		if procId == "'0'" {
			fmt.Println("Container isn't running")
			continue
		}
		// Replace all quotes
		scrubbedProcId := strings.Replace(procId, "'", "", -1)
		// TODO: Check for ids named 0, which means no proc id
		err := setSymlink(scrubbedProcId, scrubbedProcId)
		if err != nil {
			panic(err)
		}
		go startIPProcess(nmProcessorChan, scrubbedProcId, "tcpdump", "-tt", "-n", "-i", "any", "-l")
	}
}
func networkMonitorProcessor(sending chan<-[]byte,receiving <-chan[]byte)(error){
	platform := []byte{}
	carlson := []byte{}
	pattern := "(\\d+\\.\\d+) IP (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+) > (\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3})\\.(\\d+)"
	compiled_pattern,err := regexp.Compile(pattern)
	if err != nil{
		return err
	}
	for{
		platform = <- receiving
		matches := compiled_pattern.FindSubmatch(platform)
		fmt.Println(string(matches[2]), string(matches[1]))
		carl := ipdata{
			Epoch: string(matches[1]),
			SourceIP: string(matches[2]),
			SourcePort: string(matches[3]),
			ReceiveIP: string(matches[4]),
			ReceivePort: string(matches[5]),
		}
		carlson,err = json.Marshal(carl)
		if err != nil{
			return err
		}
		fmt.Println("Sending json: ", string(carlson))
		sending<-carlson
	}

	return nil
}
func startIPProcess(messages chan<- []byte, procId string, watcherName string,
	watcherArgs ...string) {
	arguments := []string{"netns", "exec", procId, watcherName}
	arguments = append(arguments, watcherArgs...)
	err := runCommandAndChannelOutput("ip", arguments, messages)
	if err != nil {
		panic(err)
	}
}

func setSymlink(procId string, destination string) error {
	nameSpaceDir := "/var/run/netns/"
	target := nameSpaceDir+destination
	source := "/proc/"+procId+"/ns/net"

	err := os.MkdirAll(nameSpaceDir, 0777)
	if err != nil {
		return err
	}
	_, err = os.Stat(target)
	if err == nil {
		err = os.Remove(target)
		if err != nil {
			return err
		}
	}

	_, err = os.Stat(source)
	if os.IsNotExist(err) {
		return errors.New("Could not set symlink: Target directory " +
			target + " does not exist")
	}

	err = os.Symlink(source, target)
	return err
}

/* This function is going to need some comments describing why we chose this
 * approach, because it will be hairy.
 * FIXME: Break out functionality into multiple functions so we can test this
 * easily!
 */
func (m FSMonitor) getDockerFSDirectory(dockerComposeName string) []string {
	ids := getDockerContainerIds(dockerComposeName)
	/*
		dockerInfoCommand := exec.Command("docker", "info")
		infoOutPipe, err := dockerInfoCommand.StdoutPipe()

		if err != nil {
			panic(err)
		}

		dockerInfoCommand.Start()
		defer dockerInfoCommand.Wait()

		infoBuf, err := ioutil.ReadAll(infoOutPipe)

		if err != nil {
			panic(err)
		}

		info := string(infoBuf)

		re := regexp.MustCompile("Docker Root Dir: (.*)")

		submatch := re.FindStringSubmatch(info)

		if len(submatch) < 2 {
			fmt.Println(submatch, info)
			panic("Couldn't find the docker root directory")
		}
		dockerRootPath := submatch[1]
	*/
	dockerRootPath := "/var/lib/docker/aufs"

	for i := 0; i < len(ids); i++ {
		ids[i] = dockerRootPath + "/mnt/" + ids[i]
	}
	return ids
}

/* Start the process running on the honeypot host to monitor the Docker
 * container. The Docker container's filesystem is mounted on the host. Find
 * the location of this filesystem with the getDockerFSDirectory function and
 * store it in the struct. Then create and start the process and forward
 * the output of the process on to the messages channel.
 */
func (m FSMonitor) Start(messages chan<- []byte, dockerComposeName string) {
	m.DockerDirs = m.getDockerFSDirectory(dockerComposeName)
	// FIXME Make arguments configurable
	//arguments := append([]string{"-wte"}, m.DockerDirs...)
	arguments := m.DockerDirs
	fmt.Println(m.MonitorName, arguments)
	m.fsWatcherProc = exec.Command(m.MonitorName, arguments...)
	defer m.fsWatcherProc.Wait()

	outpipe, err := m.fsWatcherProc.StdoutPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stdout pipe")
		panic(err)
	}
	stderr, err := m.fsWatcherProc.StderrPipe()
	if err != nil {
		log.Println("Could not open the ", m.MonitorName, "stderr pipe")
		panic(err)
	}
	go io.Copy(os.Stderr, stderr)

	m.fsWatcherProc.Start()

	stdoutReader := bufio.NewReader(outpipe)

	for {
		fetch := true
		line := []byte{}
		for fetch {
			partial_line, f, err := stdoutReader.ReadLine()
			if err != nil {
				fmt.Println("File monitor stopped")
				panic(err)
			}
			fetch = f
			line = append(line, partial_line...)
		}
		messages <- line

	}
}

/* Stop the filesystem monitor. Kill the process monitoring the Docker
 * container's filesysem.
 */
func (m FSMonitor) Stop() {
	err := m.fsWatcherProc.Process.Kill()
	if err != nil {
		log.Fatal("Could not kill ", m.MonitorName, err)
	}
}
<|endoftext|>