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

// Package listener's nodes.go contains a subset of listener functionality
// where it connects to kubernetes using kubernetes client-go and watches
// node creation/deletion events and then adds/deletes the nodes
// to/from romana cluster appropriately.
package listener

import (
	"errors"
	"fmt"
	"net"
	"time"

	log "github.com/romana/rlog"

	romanaApi "github.com/romana/core/common/api"

	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
)

func (l *KubeListener) kubeClientInit() error {
	var err error

	// Try generating config for kubernetes client-go from
	// in-cluster variables like KUBERNETES_SERVICE_HOST, etc
	// so that we can connect to kubernetes using them.
	kConfig, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("Error: %s", err)
	}

	// Get a set of REST clients which connect to kubernetes services
	// from the config generated above.
	l.kubeClientSet, err = kubernetes.NewForConfig(kConfig)
	if err != nil {
		return fmt.Errorf("Error while connecting to kubernetes: %s", err)
	}

	return nil
}

// ProcessNodeEvents processes kubernetes node events, there by
// adding/deleting nodes to/from romana cluster automatically
// when they are added/removed to/from kubernetes cluster.
func (l *KubeListener) ProcessNodeEvents(done <-chan struct{}) {
	log.Debug("In ProcessNodeEvents()")

	// nodeWatcher is a new ListWatch object created from the specified
	// kubeClientSet which k8s.io/client-go exports for watching node events.
	nodeWatcher := cache.NewListWatchFromClient(
		l.kubeClientSet.CoreV1Client.RESTClient(),
		"nodes",
		api.NamespaceAll,
		fields.Everything())

	// Setup a notifications for specific events using NewInformer.
	_, nodeInformer := cache.NewInformer(
		nodeWatcher,
		&v1.Node{},
		time.Minute,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    l.kubernetesAddNodeEventHandler,
			UpdateFunc: l.kubernetesUpdateNodeEventHandler,
			DeleteFunc: l.kubernetesDeleteNodeEventHandler,
		},
	)

	log.Infof("Starting receving node events.")
	go nodeInformer.Run(done)
}

// kubernetesAddNodeEventHandler is called when Kubernetes reports an
// add-node event It connects to the Romana REST API and adds the node
// to the Romana cluster.
func (l *KubeListener) kubernetesAddNodeEventHandler(n interface{}) {
	log.Debug("In kubernetesAddNodeEventHandler()")
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Add Event received for node(%s) ", n)
		return
	}

	log.Debugf("Add Event received for node(%s, %s, %d) ",
		node.Name, node.Status.Addresses, len(node.Status.Addresses))

	if err := romanaHostAdd(l, node); err != nil {
		log.Errorf("Error processing Add Event received for node(%s): %s",
			node.Name, err)
		return
	}

	log.Infof("Node (%s) successful added to romana cluster.", node.Name)
}

// kubernetesUpdateNodeEventHandler currently doesn't do anything yet.
// TODO: If node shows up with new IP or romana CIDR,
//       then accommodate it if possible.
func (l *KubeListener) kubernetesUpdateNodeEventHandler(o, n interface{}) {
	log.Debug("In kubernetesUpdateNodeEventHandler()")

	// node, ok := n.(*v1.Node)
	_, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Update Event received for node(%s) ", n)
		return
	}

	// Disable this for now, update events are sent every
	// 10 seconds per node, thus this could fill up the log
	// easily in very small amount of time.
	// log.Info("Update Event received for node: ",node.Name)
}

// kubernetesDeleteNodeEventHandler is called when Kubernetes reports a
// delete-node event It connects to the Romana REST API and deletes the
// node from the Romana cluster.
func (l *KubeListener) kubernetesDeleteNodeEventHandler(n interface{}) {
	log.Debug("In kubernetesDeleteNodeEventHandler()")
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Delete Event received for node(%s) ", n)
		return
	}

	log.Debugf("Delete Event received for node(%s, %s) ",
		node.Name, node.Status.Addresses)

	if err := romanaHostRemove(l, node.Name); err != nil {
		log.Errorf("Error processing Delete Event received for node(%s) ",
			node.Name)
		return
	}

	log.Infof("Node (%s) successful removed from romana cluster.", node.Name)
}

// romanaHostAdd connects to romana API and adds a node to
// the romana cluster.
func romanaHostAdd(l *KubeListener, node *v1.Node) error {
	if node.Name == "" || len(node.Status.Addresses) < 1 {
		log.Errorf("Error: received invalid host name or IP Address: (%s)", node)
		return errors.New("Error: received invalid host name or IP Address.")
	}
	hostname := node.Name
	hostIP := net.ParseIP(node.Status.Addresses[0].Address)
	host := romanaApi.Host{IP: hostIP,
		Name: hostname,
		Tags: node.GetAnnotations(),
	}
	return l.client.IPAM.AddHost(host)
}

// romanaHostRemove connects to romana API and removes a node from
// the romana cluster.
func romanaHostRemove(l *KubeListener, node string) error {
	if node == "" {
		log.Errorf("Error: received invalid node name (%s)", node)
		return errors.New("Error: received invalid node name.")
	}
	return l.client.IPAM.RemoveHost(romanaApi.Host{Name: node})
}
<commit_msg>Fix listener to use node.Labels instead of node.Annotation<commit_after>// Copyright (c) 2016 Pani Networks Inc
// All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations
// under the License.

// Package listener's nodes.go contains a subset of listener functionality
// where it connects to kubernetes using kubernetes client-go and watches
// node creation/deletion events and then adds/deletes the nodes
// to/from romana cluster appropriately.
package listener

import (
	"errors"
	"fmt"
	"net"
	"time"

	log "github.com/romana/rlog"

	romanaApi "github.com/romana/core/common/api"

	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
)

func (l *KubeListener) kubeClientInit() error {
	var err error

	// Try generating config for kubernetes client-go from
	// in-cluster variables like KUBERNETES_SERVICE_HOST, etc
	// so that we can connect to kubernetes using them.
	kConfig, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("Error: %s", err)
	}

	// Get a set of REST clients which connect to kubernetes services
	// from the config generated above.
	l.kubeClientSet, err = kubernetes.NewForConfig(kConfig)
	if err != nil {
		return fmt.Errorf("Error while connecting to kubernetes: %s", err)
	}

	return nil
}

// ProcessNodeEvents processes kubernetes node events, there by
// adding/deleting nodes to/from romana cluster automatically
// when they are added/removed to/from kubernetes cluster.
func (l *KubeListener) ProcessNodeEvents(done <-chan struct{}) {
	log.Debug("In ProcessNodeEvents()")

	// nodeWatcher is a new ListWatch object created from the specified
	// kubeClientSet which k8s.io/client-go exports for watching node events.
	nodeWatcher := cache.NewListWatchFromClient(
		l.kubeClientSet.CoreV1Client.RESTClient(),
		"nodes",
		api.NamespaceAll,
		fields.Everything())

	// Setup a notifications for specific events using NewInformer.
	_, nodeInformer := cache.NewInformer(
		nodeWatcher,
		&v1.Node{},
		time.Minute,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    l.kubernetesAddNodeEventHandler,
			UpdateFunc: l.kubernetesUpdateNodeEventHandler,
			DeleteFunc: l.kubernetesDeleteNodeEventHandler,
		},
	)

	log.Infof("Starting receving node events.")
	go nodeInformer.Run(done)
}

// kubernetesAddNodeEventHandler is called when Kubernetes reports an
// add-node event It connects to the Romana REST API and adds the node
// to the Romana cluster.
func (l *KubeListener) kubernetesAddNodeEventHandler(n interface{}) {
	log.Debug("In kubernetesAddNodeEventHandler()")
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Add Event received for node(%s) ", n)
		return
	}

	log.Debugf("Add Event received for node(%s, %s, %d) ",
		node.Name, node.Status.Addresses, len(node.Status.Addresses))

	if err := romanaHostAdd(l, node); err != nil {
		log.Errorf("Error processing Add Event received for node(%s): %s",
			node.Name, err)
		return
	}

	log.Infof("Node (%s) successful added to romana cluster.", node.Name)
}

// kubernetesUpdateNodeEventHandler currently doesn't do anything yet.
// TODO: If node shows up with new IP or romana CIDR,
//       then accommodate it if possible.
func (l *KubeListener) kubernetesUpdateNodeEventHandler(o, n interface{}) {
	log.Debug("In kubernetesUpdateNodeEventHandler()")

	// node, ok := n.(*v1.Node)
	_, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Update Event received for node(%s) ", n)
		return
	}

	// Disable this for now, update events are sent every
	// 10 seconds per node, thus this could fill up the log
	// easily in very small amount of time.
	// log.Info("Update Event received for node: ",node.Name)
}

// kubernetesDeleteNodeEventHandler is called when Kubernetes reports a
// delete-node event It connects to the Romana REST API and deletes the
// node from the Romana cluster.
func (l *KubeListener) kubernetesDeleteNodeEventHandler(n interface{}) {
	log.Debug("In kubernetesDeleteNodeEventHandler()")
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Error processing Delete Event received for node(%s) ", n)
		return
	}

	log.Debugf("Delete Event received for node(%s, %s) ",
		node.Name, node.Status.Addresses)

	if err := romanaHostRemove(l, node.Name); err != nil {
		log.Errorf("Error processing Delete Event received for node(%s) ",
			node.Name)
		return
	}

	log.Infof("Node (%s) successful removed from romana cluster.", node.Name)
}

// romanaHostAdd connects to romana API and adds a node to
// the romana cluster.
func romanaHostAdd(l *KubeListener, node *v1.Node) error {
	if node.Name == "" || len(node.Status.Addresses) < 1 {
		log.Errorf("Error: received invalid host name or IP Address: (%s)", node)
		return errors.New("Error: received invalid host name or IP Address.")
	}
	hostname := node.Name
	hostIP := net.ParseIP(node.Status.Addresses[0].Address)
	host := romanaApi.Host{IP: hostIP,
		Name: hostname,
		Tags: node.GetLabels(),
	}
	return l.client.IPAM.AddHost(host)
}

// romanaHostRemove connects to romana API and removes a node from
// the romana cluster.
func romanaHostRemove(l *KubeListener, node string) error {
	if node == "" {
		log.Errorf("Error: received invalid node name (%s)", node)
		return errors.New("Error: received invalid node name.")
	}
	return l.client.IPAM.RemoveHost(romanaApi.Host{Name: node})
}
<|endoftext|>


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

package gg

// LayerLines is like LayerPaths, but connects data points in order by
// the "x" property.
type LayerLines LayerPaths

func (l LayerLines) Apply(p *Plot) {
	defer p.Save().Restore()
	// TODO: This doesn't fill in the default correctly if X is "".
	p = p.SortBy(l.X)
	LayerPaths(l).Apply(p)
}

//go:generate stringer -type StepMode

// StepMode controls how LayerSteps connects subsequent points.
type StepMode int

const (
	// StepHV makes LayerSteps connect subsequent points with a
	// horizontal segment and then a vertical segment.
	StepHV StepMode = iota

	// StepVH makes LayerSteps connect subsequent points with a
	// vertical segment and then a horizontal segment.
	StepVH

	// StepHMid makes LayerSteps connect subsequent points A and B
	// with three segments: a horizontal segment from A to the
	// midpoint between A and B, followed by vertical segment,
	// followed by a horizontal segment from the midpoint to B.
	StepHMid

	// StepVMid makes LayerSteps connect subsequent points A and B
	// with three segments: a vertical segment from A to the
	// midpoint between A and B, followed by horizontal segment,
	// followed by a vertical segment from the midpoint to B.
	StepVMid
)

// LayerSteps is like LayerPaths, but connects data points with a path
// consisting only of horizontal and vertical segments.
type LayerSteps struct {
	LayerPaths

	Step StepMode
}

func (l LayerSteps) Apply(p *Plot) {
	// TODO: Should this also support only showing horizontal or
	// vertical segments?
	//
	// TODO: This could be a data transform instead of a layer.
	// Then it could be used in conjunction with, for example,
	// ribbons.

	l.resolveDefaults()
	p.marks = append(p.marks, plotMark{&markSteps{
		l.Step,
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPaths connects successive data points in each group with a
// path and/or a filled polygon.
type LayerPaths struct {
	// X and Y name columns that define the location of each point
	// on the path. If these are empty, they default to the first
	// and second columns, respectively.
	X, Y string

	// Color names a column that defines the stroke color of each
	// path. If Color is "", it defaults to constant black.
	// Otherwise, the data is grouped by Color.
	Color string

	// Fill names a column that defines the fill color of each
	// path. If Fill is "", it defaults to none. Otherwise, the
	// data is grouped by Fill.
	Fill string

	// XXX Perhaps the theme should provide default values for
	// things like "color". That would suggest we need to resolve
	// defaults like that at render time. Possibly a special scale
	// that gets values from the theme could be used to resolve
	// them.
	//
	// XXX strokeOpacity, fillOpacity, strokeWidth, what other
	// properties do SVG strokes have?
	//
	// XXX Should the set of known styling bindings be fixed, and
	// all possible rendering targets have to know what to do with
	// them, or should the rendering target be able to have
	// different styling bindings they understand (presumably with
	// some reasonable base set)? If the renderer can determine
	// the known bindings, we would probably just capture the
	// environment here (and make it so a captured environment
	// does not change) and hand that to the renderer later.
}

func (l *LayerPaths) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerPaths) Apply(p *Plot) {
	l.resolveDefaults()
	defer p.Save().Restore()
	if l.Color != "" {
		p.GroupBy(l.Color)
	}
	if l.Fill != "" {
		p.GroupBy(l.Fill)
	}
	p.marks = append(p.marks, plotMark{&markPath{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPoints layers a point mark at each data point.
type LayerPoints struct {
	// X and Y name columns that define the location of each
	// point. If these are empty, they default to the first and
	// second columns, respectively.
	X, Y string

	// XXX Color, fill, size, shape
}

func (l *LayerPoints) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerPoints) Apply(p *Plot) {
	l.resolveDefaults()
	p.marks = append(p.marks, plotMark{&markPoint{
		p.use("x", l.X),
		p.use("y", l.Y),
	}, p.Data().Tables()})
}

// LayerTiles layers a rectangle at each data point. The rectangle is
// specified by its center, width, and height.
type LayerTiles struct {
	// X and Y name columns that define the center of each
	// rectangle. If they are "", they default to the first and
	// second columns, respectively.
	X, Y string

	// Width and Height name columns that define the width and
	// height of each rectangle. If they are "", the width and/or
	// height are automatically determined from the smallest
	// spacing between distinct X and Y points.
	Width, Height string

	// Fill names a column that defines the fill color of each
	// rectangle. If it is "", the default fill is black.
	Fill string

	// XXX Stroke color/width, opacity, center adjustment.
}

func (l *LayerTiles) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerTiles) Apply(p *Plot) {
	l.resolveDefaults()
	if l.Width != "" || l.Height != "" {
		// TODO: What scale are these in? (x+width) is in the
		// X scale, but width itself is not. It doesn't make
		// sense to train the X scale on width, and if there's
		// a scale transform, (x+width) has to happen before
		// the transform. OTOH, if x is discrete, I can't do
		// (x+width); maybe in that case you just can't
		// specify a width. OTOOH, if width is specified and
		// the value is unscaled, I could still do something
		// reasonable with that if x is discrete.
		panic("not implemented: non-default width/height")
	}
	p.marks = append(p.marks, plotMark{&markTiles{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}
<commit_msg>gg: make LayerPaths leaving color/fill grouping in place<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gg

// LayerLines is like LayerPaths, but connects data points in order by
// the "x" property.
type LayerLines LayerPaths

func (l LayerLines) Apply(p *Plot) {
	LayerPaths(l).apply(p, true)
}

//go:generate stringer -type StepMode

// StepMode controls how LayerSteps connects subsequent points.
type StepMode int

const (
	// StepHV makes LayerSteps connect subsequent points with a
	// horizontal segment and then a vertical segment.
	StepHV StepMode = iota

	// StepVH makes LayerSteps connect subsequent points with a
	// vertical segment and then a horizontal segment.
	StepVH

	// StepHMid makes LayerSteps connect subsequent points A and B
	// with three segments: a horizontal segment from A to the
	// midpoint between A and B, followed by vertical segment,
	// followed by a horizontal segment from the midpoint to B.
	StepHMid

	// StepVMid makes LayerSteps connect subsequent points A and B
	// with three segments: a vertical segment from A to the
	// midpoint between A and B, followed by horizontal segment,
	// followed by a vertical segment from the midpoint to B.
	StepVMid
)

// LayerSteps is like LayerPaths, but connects data points with a path
// consisting only of horizontal and vertical segments.
type LayerSteps struct {
	LayerPaths

	Step StepMode
}

func (l LayerSteps) Apply(p *Plot) {
	// TODO: Should this also support only showing horizontal or
	// vertical segments?
	//
	// TODO: This could be a data transform instead of a layer.
	// Then it could be used in conjunction with, for example,
	// ribbons.

	l.resolveDefaults()
	p.marks = append(p.marks, plotMark{&markSteps{
		l.Step,
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPaths groups by Color and Fill, and then connects successive
// data points in each group with a path and/or a filled polygon.
type LayerPaths struct {
	// X and Y name columns that define the location of each point
	// on the path. If these are empty, they default to the first
	// and second columns, respectively.
	X, Y string

	// Color names a column that defines the stroke color of each
	// path. If Color is "", it defaults to constant black.
	// Otherwise, the data is grouped by Color.
	Color string

	// Fill names a column that defines the fill color of each
	// path. If Fill is "", it defaults to none. Otherwise, the
	// data is grouped by Fill.
	Fill string

	// XXX Perhaps the theme should provide default values for
	// things like "color". That would suggest we need to resolve
	// defaults like that at render time. Possibly a special scale
	// that gets values from the theme could be used to resolve
	// them.
	//
	// XXX strokeOpacity, fillOpacity, strokeWidth, what other
	// properties do SVG strokes have?
	//
	// XXX Should the set of known styling bindings be fixed, and
	// all possible rendering targets have to know what to do with
	// them, or should the rendering target be able to have
	// different styling bindings they understand (presumably with
	// some reasonable base set)? If the renderer can determine
	// the known bindings, we would probably just capture the
	// environment here (and make it so a captured environment
	// does not change) and hand that to the renderer later.
}

func (l *LayerPaths) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerPaths) Apply(p *Plot) {
	l.apply(p, false)
}

func (l LayerPaths) apply(p *Plot, sort bool) {
	l.resolveDefaults()
	if l.Color != "" {
		p.GroupBy(l.Color)
	}
	if l.Fill != "" {
		p.GroupBy(l.Fill)
	}
	if sort {
		defer p.Save().Restore()
		// TODO: This doesn't understand @0, so this fails if
		// X is defaulted.
		p = p.SortBy(l.X)
	}

	p.marks = append(p.marks, plotMark{&markPath{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPoints layers a point mark at each data point.
type LayerPoints struct {
	// X and Y name columns that define the location of each
	// point. If these are empty, they default to the first and
	// second columns, respectively.
	X, Y string

	// XXX Color, fill, size, shape
}

func (l *LayerPoints) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerPoints) Apply(p *Plot) {
	l.resolveDefaults()
	p.marks = append(p.marks, plotMark{&markPoint{
		p.use("x", l.X),
		p.use("y", l.Y),
	}, p.Data().Tables()})
}

// LayerTiles layers a rectangle at each data point. The rectangle is
// specified by its center, width, and height.
type LayerTiles struct {
	// X and Y name columns that define the center of each
	// rectangle. If they are "", they default to the first and
	// second columns, respectively.
	X, Y string

	// Width and Height name columns that define the width and
	// height of each rectangle. If they are "", the width and/or
	// height are automatically determined from the smallest
	// spacing between distinct X and Y points.
	Width, Height string

	// Fill names a column that defines the fill color of each
	// rectangle. If it is "", the default fill is black.
	Fill string

	// XXX Stroke color/width, opacity, center adjustment.
}

func (l *LayerTiles) resolveDefaults() {
	if l.X == "" {
		l.X = "@0"
	}
	if l.Y == "" {
		l.Y = "@1"
	}
}

func (l LayerTiles) Apply(p *Plot) {
	l.resolveDefaults()
	if l.Width != "" || l.Height != "" {
		// TODO: What scale are these in? (x+width) is in the
		// X scale, but width itself is not. It doesn't make
		// sense to train the X scale on width, and if there's
		// a scale transform, (x+width) has to happen before
		// the transform. OTOH, if x is discrete, I can't do
		// (x+width); maybe in that case you just can't
		// specify a width. OTOOH, if width is specified and
		// the value is unscaled, I could still do something
		// reasonable with that if x is discrete.
		panic("not implemented: non-default width/height")
	}
	p.marks = append(p.marks, plotMark{&markTiles{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}
<|endoftext|>


<commit_before>// Package run implements the "plz run" command.
package run

import (
	"context"
	"fmt"
	"math"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"golang.org/x/sync/errgroup"
	"gopkg.in/op/go-logging.v1"

	"github.com/thought-machine/please/src/core"
	"github.com/thought-machine/please/src/output"
	"github.com/thought-machine/please/src/process"
)

var log = logging.MustGetLogger("run")

// Run implements the running part of 'plz run'.
func Run(state *core.BuildState, label core.BuildLabel, args []string, remote, env bool) {
	run(context.Background(), state, label, args, false, false, remote, env, false)
}

// Parallel runs a series of targets in parallel.
// Returns a relevant exit code (i.e. if at least one subprocess exited unsuccessfully, it will be
// that code, otherwise 0 if all were successful).
// The given context can be used to control the lifetime of the subprocesses.
func Parallel(ctx context.Context, state *core.BuildState, labels []core.BuildLabel, args []string, numTasks int, quiet, remote, env, detach bool) int {
	limiter := make(chan struct{}, numTasks)
	var g errgroup.Group
	for _, label := range labels {
		label := label // capture locally
		g.Go(func() error {
			limiter <- struct{}{}
			defer func() { <-limiter }()
			return run(ctx, state, label, args, true, quiet, remote, env, detach)
		})
	}
	if err := g.Wait(); err != nil {
		if ctx.Err() != context.Canceled { // Don't error if the context killed the process.
			log.Error("Command failed: %s", err)
		}
		return err.(*exitError).code
	}
	return 0
}

// Sequential runs a series of targets sequentially.
// Returns a relevant exit code (i.e. if at least one subprocess exited unsuccessfully, it will be
// that code, otherwise 0 if all were successful).
func Sequential(state *core.BuildState, labels []core.BuildLabel, args []string, quiet, remote, env bool) int {
	for _, label := range labels {
		log.Notice("Running %s", label)
		if err := run(context.Background(), state, label, args, true, quiet, remote, env, false); err != nil {
			log.Error("%s", err)
			return err.(*exitError).code
		}
	}
	return 0
}

// run implements the internal logic about running a target.
// If fork is true then we fork to run the target and return any error from the subprocesses.
// If it's false this function never returns (because we either win or die; it's like
// Game of Thrones except rather less glamorous).
func run(ctx context.Context, state *core.BuildState, label core.BuildLabel, args []string, fork, quiet, remote, setenv, detach bool) error {
	target := state.Graph.TargetOrDie(label)
	if !target.IsBinary {
		log.Fatalf("Target %s cannot be run; it's not marked as binary", label)
	}
	if len(target.Outputs()) != 1 && target.IsTest {
		log.Fatalf("Targets %s cannot be run as it has %d outputs; Only tests with 1 output can be run.", label, len(target.Outputs()))
	}
	if remote {
		// Send this off to be done remotely.
		// This deliberately misses the out_exe bit below, but also doesn't pick up whatever's going on with java -jar;
		// however that will be obsolete post #920 anyway.
		if state.RemoteClient == nil {
			log.Fatalf("You must configure remote execution to use plz run --remote")
		}
		return state.RemoteClient.Run(target)
	}
	// ReplaceSequences always quotes stuff in case it contains spaces or special characters,
	// that works fine if we interpret it as a shell but not to pass it as an argument here.
	command, _ := core.ReplaceSequences(state, target, fmt.Sprintf("$(out_exe %s)", target.Label))
	arg0 := strings.Trim(command, "\"")
	// Handle targets where $(exe ...) returns something nontrivial
	splitCmd := strings.Split(arg0, " ")
	if !strings.Contains(splitCmd[0], "/") {
		// Probably it's a java -jar, we need an absolute path to it.
		cmd, err := exec.LookPath(splitCmd[0])
		if err != nil {
			log.Fatalf("Can't find binary %s", splitCmd[0])
		}
		splitCmd[0] = cmd
	}
	args = append(splitCmd, args...)
	log.Info("Running target %s...", strings.Join(args, " "))
	output.SetWindowTitle("plz run: " + strings.Join(args, " "))
	env := environ(state.Config, setenv)
	if !fork {
		// Plain 'plz run'. One way or another we never return from the following line.
		must(syscall.Exec(splitCmd[0], args, env), args)
	} else if detach {
		// Bypass the whole process management system since we explicitly aim not to manage this subprocess.
		cmd := exec.Command(splitCmd[0], args[1:]...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return toExitError(cmd.Start(), args, nil)
	}
	// Run as a normal subcommand.
	// Note that we don't connect stdin. It doesn't make sense for multiple processes.
	// The process executor doesn't actually support not having a timeout, but the max is ~290 years so nobody
	// should know the difference.
	_, output, err := process.New("").ExecWithTimeout(target, "", env, time.Duration(math.MaxInt64), false, false, !quiet, args)
	return toExitError(err, args, output)
}

// environ returns an appropriate environment for a command.
func environ(config *core.Configuration, setenv bool) []string {
	env := os.Environ()
	for _, e := range adRunEnviron {
		env = addEnv(env, e)
	}
	if setenv {
		for _, e := range core.GeneralBuildEnvironment(config) {
			env = addEnv(env, e)
		}
	}
	return env
}

// adRunEnviron returns values that are appended to the environment for a command.
var adRunEnviron = []string{
	"PEX_NOCACHE=true",
}

// addEnv adds an env var to an existing set, with replacement.
func addEnv(env []string, e string) []string {
	name := e[:strings.IndexRune(e, '=')+1]
	for i, existing := range env {
		if strings.HasPrefix(existing, name) {
			env[i] = e
			return env
		}
	}
	return append(env, e)
}

// must dies if the given error is non-nil.
func must(err error, cmd []string) {
	if err != nil {
		log.Fatalf("Error running command %s: %s", strings.Join(cmd, " "), err)
	}
}

// toExitError attempts to extract the exit code from an error.
func toExitError(err error, cmd []string, out []byte) error {
	exitCode := 1
	if err == nil {
		return nil
	} else if exitError, ok := err.(*exec.ExitError); ok {
		// This is a little hairy; there isn't a good way of getting the exit code,
		// but this should be reasonably portable (at least to the platforms we care about).
		if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		}
	}
	return &exitError{
		msg:  fmt.Sprintf("Error running command %s: %s\n%s", strings.Join(cmd, " "), err, string(out)),
		code: exitCode,
	}
}

type exitError struct {
	msg  string
	code int
}

func (e *exitError) Error() string {
	return e.msg
}
<commit_msg>Don't pass target into ExecWithTimeout from `plz run`<commit_after>// Package run implements the "plz run" command.
package run

import (
	"context"
	"fmt"
	"math"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"golang.org/x/sync/errgroup"
	"gopkg.in/op/go-logging.v1"

	"github.com/thought-machine/please/src/core"
	"github.com/thought-machine/please/src/output"
	"github.com/thought-machine/please/src/process"
)

var log = logging.MustGetLogger("run")

// Run implements the running part of 'plz run'.
func Run(state *core.BuildState, label core.BuildLabel, args []string, remote, env bool) {
	run(context.Background(), state, label, args, false, false, remote, env, false)
}

// Parallel runs a series of targets in parallel.
// Returns a relevant exit code (i.e. if at least one subprocess exited unsuccessfully, it will be
// that code, otherwise 0 if all were successful).
// The given context can be used to control the lifetime of the subprocesses.
func Parallel(ctx context.Context, state *core.BuildState, labels []core.BuildLabel, args []string, numTasks int, quiet, remote, env, detach bool) int {
	limiter := make(chan struct{}, numTasks)
	var g errgroup.Group
	for _, label := range labels {
		label := label // capture locally
		g.Go(func() error {
			limiter <- struct{}{}
			defer func() { <-limiter }()
			return run(ctx, state, label, args, true, quiet, remote, env, detach)
		})
	}
	if err := g.Wait(); err != nil {
		if ctx.Err() != context.Canceled { // Don't error if the context killed the process.
			log.Error("Command failed: %s", err)
		}
		return err.(*exitError).code
	}
	return 0
}

// Sequential runs a series of targets sequentially.
// Returns a relevant exit code (i.e. if at least one subprocess exited unsuccessfully, it will be
// that code, otherwise 0 if all were successful).
func Sequential(state *core.BuildState, labels []core.BuildLabel, args []string, quiet, remote, env bool) int {
	for _, label := range labels {
		log.Notice("Running %s", label)
		if err := run(context.Background(), state, label, args, true, quiet, remote, env, false); err != nil {
			log.Error("%s", err)
			return err.(*exitError).code
		}
	}
	return 0
}

// run implements the internal logic about running a target.
// If fork is true then we fork to run the target and return any error from the subprocesses.
// If it's false this function never returns (because we either win or die; it's like
// Game of Thrones except rather less glamorous).
func run(ctx context.Context, state *core.BuildState, label core.BuildLabel, args []string, fork, quiet, remote, setenv, detach bool) error {
	target := state.Graph.TargetOrDie(label)
	if !target.IsBinary {
		log.Fatalf("Target %s cannot be run; it's not marked as binary", label)
	}
	if len(target.Outputs()) != 1 && target.IsTest {
		log.Fatalf("Targets %s cannot be run as it has %d outputs; Only tests with 1 output can be run.", label, len(target.Outputs()))
	}
	if remote {
		// Send this off to be done remotely.
		// This deliberately misses the out_exe bit below, but also doesn't pick up whatever's going on with java -jar;
		// however that will be obsolete post #920 anyway.
		if state.RemoteClient == nil {
			log.Fatalf("You must configure remote execution to use plz run --remote")
		}
		return state.RemoteClient.Run(target)
	}
	// ReplaceSequences always quotes stuff in case it contains spaces or special characters,
	// that works fine if we interpret it as a shell but not to pass it as an argument here.
	command, _ := core.ReplaceSequences(state, target, fmt.Sprintf("$(out_exe %s)", target.Label))
	arg0 := strings.Trim(command, "\"")
	// Handle targets where $(exe ...) returns something nontrivial
	splitCmd := strings.Split(arg0, " ")
	if !strings.Contains(splitCmd[0], "/") {
		// Probably it's a java -jar, we need an absolute path to it.
		cmd, err := exec.LookPath(splitCmd[0])
		if err != nil {
			log.Fatalf("Can't find binary %s", splitCmd[0])
		}
		splitCmd[0] = cmd
	}
	args = append(splitCmd, args...)
	log.Info("Running target %s...", strings.Join(args, " "))
	output.SetWindowTitle("plz run: " + strings.Join(args, " "))
	env := environ(state.Config, setenv)
	if !fork {
		// Plain 'plz run'. One way or another we never return from the following line.
		must(syscall.Exec(splitCmd[0], args, env), args)
	} else if detach {
		// Bypass the whole process management system since we explicitly aim not to manage this subprocess.
		cmd := exec.Command(splitCmd[0], args[1:]...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		return toExitError(cmd.Start(), args, nil)
	}
	// Run as a normal subcommand.
	// Note that we don't connect stdin. It doesn't make sense for multiple processes.
	// The process executor doesn't actually support not having a timeout, but the max is ~290 years so nobody
	// should know the difference.
	_, output, err := process.New("").ExecWithTimeout(nil, "", env, time.Duration(math.MaxInt64), false, false, !quiet, args)
	return toExitError(err, args, output)
}

// environ returns an appropriate environment for a command.
func environ(config *core.Configuration, setenv bool) []string {
	env := os.Environ()
	for _, e := range adRunEnviron {
		env = addEnv(env, e)
	}
	if setenv {
		for _, e := range core.GeneralBuildEnvironment(config) {
			env = addEnv(env, e)
		}
	}
	return env
}

// adRunEnviron returns values that are appended to the environment for a command.
var adRunEnviron = []string{
	"PEX_NOCACHE=true",
}

// addEnv adds an env var to an existing set, with replacement.
func addEnv(env []string, e string) []string {
	name := e[:strings.IndexRune(e, '=')+1]
	for i, existing := range env {
		if strings.HasPrefix(existing, name) {
			env[i] = e
			return env
		}
	}
	return append(env, e)
}

// must dies if the given error is non-nil.
func must(err error, cmd []string) {
	if err != nil {
		log.Fatalf("Error running command %s: %s", strings.Join(cmd, " "), err)
	}
}

// toExitError attempts to extract the exit code from an error.
func toExitError(err error, cmd []string, out []byte) error {
	exitCode := 1
	if err == nil {
		return nil
	} else if exitError, ok := err.(*exec.ExitError); ok {
		// This is a little hairy; there isn't a good way of getting the exit code,
		// but this should be reasonably portable (at least to the platforms we care about).
		if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		}
	}
	return &exitError{
		msg:  fmt.Sprintf("Error running command %s: %s\n%s", strings.Join(cmd, " "), err, string(out)),
		code: exitCode,
	}
}

type exitError struct {
	msg  string
	code int
}

func (e *exitError) Error() string {
	return e.msg
}
<|endoftext|>


<commit_before>/*
 *    Copyright (C) 2014-2017 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// Package bees is Beehive's central module system.
package bees

import (
	"fmt"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	uuid "github.com/nu7hatch/gouuid"
)

// BeeInterface is an interface all bees implement.
type BeeInterface interface {
	// Name of the bee
	Name() string
	// Namespace of the bee
	Namespace() string

	// Description of the bee
	Description() string
	// SetDescription sets a description
	SetDescription(s string)

	// Config returns this bees config
	Config() BeeConfig
	// Options of the bee
	Options() BeeOptions
	// SetOptions to configure the bee
	SetOptions(options BeeOptions)

	// ReloadOptions gets called after a bee's options get updated
	ReloadOptions(options BeeOptions)

	// Activates the bee
	Run(eventChannel chan Event)
	// Running returns the current state of the bee
	IsRunning() bool
	// Start the bee
	Start()
	// Stop the bee
	Stop()

	LastEvent() time.Time
	LogEvent()
	LastAction() time.Time
	LogAction()

	Logln(args ...interface{})
	Logf(format string, args ...interface{})
	LogErrorf(format string, args ...interface{})
	LogFatal(args ...interface{})

	SetSigChan(c chan bool)
	WaitGroup() *sync.WaitGroup

	// Handles an action
	Action(action Action) []Placeholder
}

// Bee is the base-struct to be embedded by bee implementations.
type Bee struct {
	config BeeConfig

	lastEvent  time.Time
	lastAction time.Time

	Running   bool
	SigChan   chan bool
	waitGroup *sync.WaitGroup
}

var (
	bees      = make(map[string]*BeeInterface)
	factories = make(map[string]*BeeFactoryInterface)
)

// RegisterBee gets called by Bees to register themselves.
func RegisterBee(bee BeeInterface) {
	log.Println("Worker bee ready:", bee.Name(), "-", bee.Description())

	bees[bee.Name()] = &bee
}

// GetBee returns a bee with a specific name.
func GetBee(identifier string) *BeeInterface {
	bee, ok := bees[identifier]
	if ok {
		return bee
	}

	return nil
}

// GetBees returns all known bees.
func GetBees() []*BeeInterface {
	r := []*BeeInterface{}
	for _, bee := range bees {
		r = append(r, bee)
	}

	return r
}

// startBee starts a bee and recovers from panics.
func startBee(bee *BeeInterface, fatals int) {
	if fatals >= 3 {
		log.Println("Terminating evil bee", (*bee).Name(), "after", fatals, "failed tries!")
		(*bee).Stop()
		return
	}

	(*bee).WaitGroup().Add(1)
	defer (*bee).WaitGroup().Done()

	defer func(bee *BeeInterface) {
		if e := recover(); e != nil {
			log.Println("Fatal bee event:", e, fatals)
			go startBee(bee, fatals+1)
		}
	}(bee)

	(*bee).Run(eventsIn)
}

// NewBeeInstance sets up a new Bee with supplied config.
func NewBeeInstance(bee BeeConfig) *BeeInterface {
	factory := GetFactory(bee.Class)
	if factory == nil {
		panic("Unknown bee-class in config file: " + bee.Class)
	}
	mod := (*factory).New(bee.Name, bee.Description, bee.Options)
	RegisterBee(mod)

	return &mod
}

// DeleteBee removes a Bee instance.
func DeleteBee(bee *BeeInterface) {
	(*bee).Stop()

	delete(bees, (*bee).Name())
}

// StartBee starts a bee.
func StartBee(bee BeeConfig) *BeeInterface {
	b := NewBeeInstance(bee)

	(*b).Start()
	go func(mod *BeeInterface) {
		startBee(mod, 0)
	}(b)

	return b
}

// StartBees starts all registered bees.
func StartBees(beeList []BeeConfig) {
	eventsIn = make(chan Event)
	go handleEvents()

	for _, bee := range beeList {
		StartBee(bee)
	}
}

// StopBees stops all bees gracefully.
func StopBees() {
	for _, bee := range bees {
		log.Println("Stopping bee:", (*bee).Name())
		(*bee).Stop()
	}

	close(eventsIn)
	bees = make(map[string]*BeeInterface)
}

// RestartBee restarts a Bee.
func RestartBee(bee *BeeInterface) {
	(*bee).Stop()

	(*bee).SetSigChan(make(chan bool))
	(*bee).Start()
	go func(mod *BeeInterface) {
		startBee(mod, 0)
	}(bee)
}

// RestartBees stops all running bees and restarts a new set of bees.
func RestartBees(bees []BeeConfig) {
	StopBees()
	StartBees(bees)
}

// NewBee returns a new bee and sets up sig-channel & waitGroup.
func NewBee(name, factoryName, description string, options []BeeOption) Bee {
	c := BeeConfig{
		Name:        name,
		Class:       factoryName,
		Description: description,
		Options:     options,
	}
	b := Bee{
		config:    c,
		SigChan:   make(chan bool),
		waitGroup: &sync.WaitGroup{},
	}

	return b
}

// Name returns the configured name for a bee.
func (bee *Bee) Name() string {
	return bee.config.Name
}

// Namespace returns the namespace for a bee.
func (bee *Bee) Namespace() string {
	return bee.config.Class
}

// Description returns the description for a bee.
func (bee *Bee) Description() string {
	return bee.config.Description
}

// SetDescription sets the description for a bee.
func (bee *Bee) SetDescription(s string) {
	bee.config.Description = s
}

// Config returns the config for a bee.
func (bee *Bee) Config() BeeConfig {
	return bee.config
}

// Options returns the options for a bee.
func (bee *Bee) Options() BeeOptions {
	return bee.config.Options
}

// SetOptions sets the options for a bee.
func (bee *Bee) SetOptions(options BeeOptions) {
	bee.config.Options = options
}

// SetSigChan sets the signaling channel for a bee.
func (bee *Bee) SetSigChan(c chan bool) {
	bee.SigChan = c
}

// WaitGroup returns the WaitGroup for a bee.
func (bee *Bee) WaitGroup() *sync.WaitGroup {
	return bee.waitGroup
}

// Run is the default, empty implementation of a Bee's Run method.
func (bee *Bee) Run(chan Event) {
	select {
	case <-bee.SigChan:
		return
	}
}

// Action is the default, empty implementation of a Bee's Action method.
func (bee *Bee) Action(action Action) []Placeholder {
	return []Placeholder{}
}

// IsRunning returns whether a Bee is currently running.
func (bee *Bee) IsRunning() bool {
	return bee.Running
}

// Start gets called when a Bee gets started.
func (bee *Bee) Start() {
	bee.Running = true
}

// Stop gracefully stops a Bee.
func (bee *Bee) Stop() {
	if !bee.IsRunning() {
		return
	}
	log.Println(bee.Name(), "stopping gracefully!")

	close(bee.SigChan)
	bee.waitGroup.Wait()
	bee.Running = false
	log.Println(bee.Name(), "stopped gracefully!")
}

// LastEvent returns the timestamp of the last triggered event.
func (bee *Bee) LastEvent() time.Time {
	return bee.lastEvent
}

// LastAction returns the timestamp of the last triggered action.
func (bee *Bee) LastAction() time.Time {
	return bee.lastAction
}

// LogEvent logs the last triggered event.
func (bee *Bee) LogEvent() {
	bee.lastEvent = time.Now()
}

// LogAction logs the last triggered action.
func (bee *Bee) LogAction() {
	bee.lastAction = time.Now()
}

// Logln logs args
func (bee *Bee) Logln(args ...interface{}) {
	a := []interface{}{"[" + bee.Name() + "]:"}
	for _, v := range args {
		a = append(a, v)
	}

	log.Println(a...)
}

// Logf logs a formatted string
func (bee *Bee) Logf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	log.Printf("[%s]: %s", bee.Name(), s)
}

// LogErrorf logs a formatted error string
func (bee *Bee) LogErrorf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	log.Errorf("[%s]: %s", bee.Name(), s)
}

// LogFatal logs a fatal error
func (bee *Bee) LogFatal(args ...interface{}) {
	a := []interface{}{"[" + bee.Name() + "]:"}
	for _, v := range args {
		a = append(a, v)
	}
	log.Fatal(a...)
}

// UUID generates a new unique ID.
func UUID() string {
	u, _ := uuid.NewV4()
	return u.String()
}
<commit_msg>Don't call log.Fatal() in bees, use log.Panicln() instead<commit_after>/*
 *    Copyright (C) 2014-2017 Christian Muehlhaeuser
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors:
 *      Christian Muehlhaeuser <muesli@gmail.com>
 */

// Package bees is Beehive's central module system.
package bees

import (
	"fmt"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	uuid "github.com/nu7hatch/gouuid"
)

// BeeInterface is an interface all bees implement.
type BeeInterface interface {
	// Name of the bee
	Name() string
	// Namespace of the bee
	Namespace() string

	// Description of the bee
	Description() string
	// SetDescription sets a description
	SetDescription(s string)

	// Config returns this bees config
	Config() BeeConfig
	// Options of the bee
	Options() BeeOptions
	// SetOptions to configure the bee
	SetOptions(options BeeOptions)

	// ReloadOptions gets called after a bee's options get updated
	ReloadOptions(options BeeOptions)

	// Activates the bee
	Run(eventChannel chan Event)
	// Running returns the current state of the bee
	IsRunning() bool
	// Start the bee
	Start()
	// Stop the bee
	Stop()

	LastEvent() time.Time
	LogEvent()
	LastAction() time.Time
	LogAction()

	Logln(args ...interface{})
	Logf(format string, args ...interface{})
	LogErrorf(format string, args ...interface{})
	LogFatal(args ...interface{})

	SetSigChan(c chan bool)
	WaitGroup() *sync.WaitGroup

	// Handles an action
	Action(action Action) []Placeholder
}

// Bee is the base-struct to be embedded by bee implementations.
type Bee struct {
	config BeeConfig

	lastEvent  time.Time
	lastAction time.Time

	Running   bool
	SigChan   chan bool
	waitGroup *sync.WaitGroup
}

var (
	bees      = make(map[string]*BeeInterface)
	factories = make(map[string]*BeeFactoryInterface)
)

// RegisterBee gets called by Bees to register themselves.
func RegisterBee(bee BeeInterface) {
	log.Println("Worker bee ready:", bee.Name(), "-", bee.Description())

	bees[bee.Name()] = &bee
}

// GetBee returns a bee with a specific name.
func GetBee(identifier string) *BeeInterface {
	bee, ok := bees[identifier]
	if ok {
		return bee
	}

	return nil
}

// GetBees returns all known bees.
func GetBees() []*BeeInterface {
	r := []*BeeInterface{}
	for _, bee := range bees {
		r = append(r, bee)
	}

	return r
}

// startBee starts a bee and recovers from panics.
func startBee(bee *BeeInterface, fatals int) {
	if fatals >= 3 {
		log.Println("Terminating evil bee", (*bee).Name(), "after", fatals, "failed tries!")
		(*bee).Stop()
		return
	}

	(*bee).WaitGroup().Add(1)
	defer (*bee).WaitGroup().Done()

	defer func(bee *BeeInterface) {
		if e := recover(); e != nil {
			log.Println("Fatal bee event:", e, fatals)
			go startBee(bee, fatals+1)
		}
	}(bee)

	(*bee).Run(eventsIn)
}

// NewBeeInstance sets up a new Bee with supplied config.
func NewBeeInstance(bee BeeConfig) *BeeInterface {
	factory := GetFactory(bee.Class)
	if factory == nil {
		panic("Unknown bee-class in config file: " + bee.Class)
	}
	mod := (*factory).New(bee.Name, bee.Description, bee.Options)
	RegisterBee(mod)

	return &mod
}

// DeleteBee removes a Bee instance.
func DeleteBee(bee *BeeInterface) {
	(*bee).Stop()

	delete(bees, (*bee).Name())
}

// StartBee starts a bee.
func StartBee(bee BeeConfig) *BeeInterface {
	b := NewBeeInstance(bee)

	(*b).Start()
	go func(mod *BeeInterface) {
		startBee(mod, 0)
	}(b)

	return b
}

// StartBees starts all registered bees.
func StartBees(beeList []BeeConfig) {
	eventsIn = make(chan Event)
	go handleEvents()

	for _, bee := range beeList {
		StartBee(bee)
	}
}

// StopBees stops all bees gracefully.
func StopBees() {
	for _, bee := range bees {
		log.Println("Stopping bee:", (*bee).Name())
		(*bee).Stop()
	}

	close(eventsIn)
	bees = make(map[string]*BeeInterface)
}

// RestartBee restarts a Bee.
func RestartBee(bee *BeeInterface) {
	(*bee).Stop()

	(*bee).SetSigChan(make(chan bool))
	(*bee).Start()
	go func(mod *BeeInterface) {
		startBee(mod, 0)
	}(bee)
}

// RestartBees stops all running bees and restarts a new set of bees.
func RestartBees(bees []BeeConfig) {
	StopBees()
	StartBees(bees)
}

// NewBee returns a new bee and sets up sig-channel & waitGroup.
func NewBee(name, factoryName, description string, options []BeeOption) Bee {
	c := BeeConfig{
		Name:        name,
		Class:       factoryName,
		Description: description,
		Options:     options,
	}
	b := Bee{
		config:    c,
		SigChan:   make(chan bool),
		waitGroup: &sync.WaitGroup{},
	}

	return b
}

// Name returns the configured name for a bee.
func (bee *Bee) Name() string {
	return bee.config.Name
}

// Namespace returns the namespace for a bee.
func (bee *Bee) Namespace() string {
	return bee.config.Class
}

// Description returns the description for a bee.
func (bee *Bee) Description() string {
	return bee.config.Description
}

// SetDescription sets the description for a bee.
func (bee *Bee) SetDescription(s string) {
	bee.config.Description = s
}

// Config returns the config for a bee.
func (bee *Bee) Config() BeeConfig {
	return bee.config
}

// Options returns the options for a bee.
func (bee *Bee) Options() BeeOptions {
	return bee.config.Options
}

// SetOptions sets the options for a bee.
func (bee *Bee) SetOptions(options BeeOptions) {
	bee.config.Options = options
}

// SetSigChan sets the signaling channel for a bee.
func (bee *Bee) SetSigChan(c chan bool) {
	bee.SigChan = c
}

// WaitGroup returns the WaitGroup for a bee.
func (bee *Bee) WaitGroup() *sync.WaitGroup {
	return bee.waitGroup
}

// Run is the default, empty implementation of a Bee's Run method.
func (bee *Bee) Run(chan Event) {
	select {
	case <-bee.SigChan:
		return
	}
}

// Action is the default, empty implementation of a Bee's Action method.
func (bee *Bee) Action(action Action) []Placeholder {
	return []Placeholder{}
}

// IsRunning returns whether a Bee is currently running.
func (bee *Bee) IsRunning() bool {
	return bee.Running
}

// Start gets called when a Bee gets started.
func (bee *Bee) Start() {
	bee.Running = true
}

// Stop gracefully stops a Bee.
func (bee *Bee) Stop() {
	if !bee.IsRunning() {
		return
	}
	log.Println(bee.Name(), "stopping gracefully!")

	close(bee.SigChan)
	bee.waitGroup.Wait()
	bee.Running = false
	log.Println(bee.Name(), "stopped gracefully!")
}

// LastEvent returns the timestamp of the last triggered event.
func (bee *Bee) LastEvent() time.Time {
	return bee.lastEvent
}

// LastAction returns the timestamp of the last triggered action.
func (bee *Bee) LastAction() time.Time {
	return bee.lastAction
}

// LogEvent logs the last triggered event.
func (bee *Bee) LogEvent() {
	bee.lastEvent = time.Now()
}

// LogAction logs the last triggered action.
func (bee *Bee) LogAction() {
	bee.lastAction = time.Now()
}

// Logln logs args
func (bee *Bee) Logln(args ...interface{}) {
	a := []interface{}{"[" + bee.Name() + "]:"}
	for _, v := range args {
		a = append(a, v)
	}

	log.Println(a...)
}

// Logf logs a formatted string
func (bee *Bee) Logf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	log.Printf("[%s]: %s", bee.Name(), s)
}

// LogErrorf logs a formatted error string
func (bee *Bee) LogErrorf(format string, args ...interface{}) {
	s := fmt.Sprintf(format, args...)
	log.Errorf("[%s]: %s", bee.Name(), s)
}

// LogFatal logs a fatal error
func (bee *Bee) LogFatal(args ...interface{}) {
	a := []interface{}{"[" + bee.Name() + "]:"}
	for _, v := range args {
		a = append(a, v)
	}
	log.Panicln(a...)
}

// UUID generates a new unique ID.
func UUID() string {
	u, _ := uuid.NewV4()
	return u.String()
}
<|endoftext|>


<commit_before>package blocks

import (
	"fmt"
	"github.com/nytlabs/streamtools/st/loghub"
	"log"
	"net/url"
	"reflect"
	"strconv"
	"time"
)

type Msg struct {
	Msg   interface{}
	Route string
}

type MsgChan chan interface{}

func (c MsgChan) MarshalJSON() ([]byte, error) {
	outBytes := []byte("{\"channel\":" + strconv.Itoa(len(c)) + "}")
	return outBytes, nil
}

type AddChanMsg struct {
	Route   string
	Channel chan *Msg
}

type QueryMsg struct {
	Route   string
	MsgChan MsgChan
}

type QueryParamMsg struct {
	Route    string
	RespChan chan interface{}
	Params   url.Values
}

type Query struct {
	RespChan chan interface{}
	Params   url.Values
}

type BlockChans struct {
	InChan         chan *Msg
	QueryChan      chan *QueryMsg
	QueryParamChan chan *QueryParamMsg
	AddChan        chan *AddChanMsg
	DelChan        chan *Msg
	ErrChan        chan error
	QuitChan       chan bool
}

type LogStreams struct {
	log MsgChan
	ui  MsgChan
}

type Block struct {
	Id               string // the name of the block specifed by the user (like MyBlock)
	Kind             string // the kind of block this is (like count, toFile, fromSQS)
	Desc             string // the description of block ('counts the number of messages it has seen')
	inRoutes         map[string]MsgChan
	queryRoutes      map[string]chan MsgChan
	queryParamRoutes map[string]chan Query
	broadcast        MsgChan
	quit             MsgChan
	doesBroadcast    bool
	BlockChans
	LogStreams
}

type BlockDef struct {
	Type             string
	Desc             string
	InRoutes         []string
	QueryRoutes      []string
	QueryParamRoutes []string
	OutRoutes        []string
}

type BlockInterface interface {
	Setup()
	Run()
	CleanUp()
	Build(BlockChans)
	Quit() MsgChan
	Broadcast() MsgChan
	InRoute(string) MsgChan
	QueryRoute(string) chan MsgChan
	QueryParamRoute(string) chan Query
	GetBlock() *Block
	GetDef() *BlockDef
	Log(interface{})
	Error(interface{})
	SetId(string)
}

func (b *Block) Build(c BlockChans) {
	// block channels
	b.InChan = c.InChan
	b.QueryChan = c.QueryChan
	b.QueryParamChan = c.QueryParamChan
	b.AddChan = c.AddChan
	b.DelChan = c.DelChan
	b.ErrChan = c.ErrChan
	b.QuitChan = c.QuitChan

	// route maps
	b.inRoutes = make(map[string]MsgChan) // necessary to stop locking...
	b.queryRoutes = make(map[string]chan MsgChan)
	b.queryParamRoutes = make(map[string]chan Query)

	// broadcast channel
	b.broadcast = make(MsgChan, 10) // necessary to stop locking...

	// quit chan
	b.quit = make(MsgChan)

	b.ui = make(MsgChan)
	b.log = make(MsgChan)
}

func (b *Block) SetId(Id string) {
	b.Id = Id
}

func (b *Block) InRoute(routeName string) MsgChan {
	route := make(MsgChan, 1000)
	b.inRoutes[routeName] = route
	return route
}

func (b *Block) QueryRoute(routeName string) chan MsgChan {
	route := make(chan MsgChan, 1000)
	b.queryRoutes[routeName] = route
	return route
}

func (b *Block) QueryParamRoute(routeName string) chan Query {
	route := make(chan Query, 1000)
	b.queryParamRoutes[routeName] = route
	return route
}

func (b *Block) Broadcast() MsgChan {
	b.doesBroadcast = true
	return b.broadcast
}

func (b *Block) Quit() MsgChan {
	return b.quit
}

func (b *Block) GetBlock() *Block {
	return b
}

func (b *Block) GetDef() *BlockDef {
	inRoutes := []string{}
	queryRoutes := []string{}
	queryParamRoutes := []string{}
	outRoutes := []string{}

	for k, _ := range b.inRoutes {
		inRoutes = append(inRoutes, k)
	}

	for k, _ := range b.queryRoutes {
		queryRoutes = append(queryRoutes, k)
	}

	for k, _ := range b.queryParamRoutes {
		queryParamRoutes = append(queryParamRoutes, k)
	}

	if b.doesBroadcast {
		outRoutes = []string{"out"}
	}

	return &BlockDef{
		Type:             b.Kind,
		Desc:             b.Desc,
		InRoutes:         inRoutes,
		QueryRoutes:      queryRoutes,
		QueryParamRoutes: queryParamRoutes,
		OutRoutes:        outRoutes,
	}
}

func (b *Block) CleanUp() {
	for route := range b.inRoutes {
		defer close(b.inRoutes[route])
	}
	for route := range b.queryRoutes {
		defer close(b.queryRoutes[route])
	}
	defer close(b.InChan)
	defer close(b.QueryChan)
	defer close(b.QueryParamChan)
	defer close(b.AddChan)
	defer close(b.DelChan)
	defer close(b.ErrChan)
	defer close(b.QuitChan)
	defer close(b.broadcast)

	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: fmt.Sprintf("Block %s Quitting...", b.Id),
			Id:   id,
		}
	}(b.Id)
	log.Println("about to close channels on", b.Kind)
}

func (b *Block) Error(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.ERROR,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func (b *Block) Log(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func BlockRoutine(bi BlockInterface) {
	var dropped int64
	dropTicker := time.NewTicker(time.Duration(1 * time.Second))
	dropTicker.Stop()

	outChans := make(map[string]chan *Msg)
	b := bi.GetBlock()
	bi.Setup()
	go bi.Run()

	for {
		select {
		case <-dropTicker.C:
			go func(id string, count int64) {
				loghub.Log <- &loghub.LogMsg{
					Type: loghub.ERROR,
					Data: fmt.Sprintf("Dropped messages: %d (cannot keep up with stream)", count),
					Id:   id,
				}
			}(b.Id, dropped)

			if dropped == 0 {
				dropTicker.Stop()
			}

			dropped = 0
		case msg := <-b.InChan:
			_, ok := b.inRoutes[msg.Route]
			if !ok {
				break
			}

			// every in channel is buffered a 1000 messages.
			// if we cannot immediately send to that in channel we place the msg
			// in a go routine and notify the user that the block routine's
			// buffer has overflowed. this still allows for unrecoverable
			// overflows (for example: a stuck run() function), but at least it
			// offloads to memory/cpu instead of blocking.
			select {
			case b.inRoutes[msg.Route] <- msg.Msg:
			default:
				if dropped == 0 {
					dropTicker.Stop()
					dropTicker = time.NewTicker(1 * time.Second)
				}

				dropped++
			}

			if msg.Route == "rule" {
				go func(id string) {
					loghub.UI <- &loghub.LogMsg{
						Type: loghub.RULE_UPDATED,
						Data: map[string]interface{}{},
						Id:   id,
					}
				}(b.Id)
			}

		case msg := <-b.QueryChan:

			if msg.Route == "ping" {
				msg.MsgChan <- "OK"
				continue
			}

			_, ok := b.queryRoutes[msg.Route]
			if !ok {
				break
			}

			log.Println(reflect.TypeOf(msg.MsgChan))

			select {
			case b.queryRoutes[msg.Route] <- msg.MsgChan:
			default:
				go func() {
					b.queryRoutes[msg.Route] <- msg.MsgChan
				}()
			}
		case msg := <-b.QueryParamChan:

			if msg.Route == "ping" {
				msg.RespChan <- "OK"
				continue
			}

			_, ok := b.queryParamRoutes[msg.Route]
			if !ok {
				break
			}
			q := Query{
				RespChan: msg.RespChan,
				Params:   msg.Params,
			}

			select {
			case b.queryParamRoutes[msg.Route] <- q:
			default:
				go func() {
					b.queryParamRoutes[msg.Route] <- q
				}()
			}

		case msg := <-b.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-b.DelChan:
			delete(outChans, msg.Route)
		case msg := <-b.broadcast:
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg,
					Route: "",
				}
			}
		case <-b.QuitChan:
			b.quit <- true
			b.CleanUp()
			return
		}
	}
}

type Connection struct {
	Id      string
	ToRoute string
	BlockChans
	LogStreams
}

func (c *Connection) SetId(Id string) {
	c.Id = Id
}

func (c *Connection) Build(chans BlockChans) {
	c.InChan = chans.InChan
	c.QueryChan = chans.QueryChan
	c.QueryParamChan = chans.QueryParamChan
	c.AddChan = chans.AddChan
	c.DelChan = chans.DelChan
	c.QuitChan = chans.QuitChan
}

func (c *Connection) CleanUp() {
	defer close(c.InChan)
	defer close(c.QueryChan)
	defer close(c.QueryParamChan)
	defer close(c.AddChan)
	defer close(c.DelChan)
	defer close(c.QuitChan)

	loghub.Log <- &loghub.LogMsg{
		Type: loghub.INFO,
		Data: fmt.Sprintf("Connection %s Quitting...", c.Id),
		Id:   c.Id,
	}
}

func ConnectionRoutine(c *Connection) {
	var last interface{}
	var rate float64

	outChans := make(map[string]chan *Msg)
	times := make([]int64, 100, 100)
	timesIdx := len(times)
	rateReport := time.NewTicker(200 * time.Millisecond)

	for {
		select {
		case <-rateReport.C:
			if timesIdx == len(times) {
				rate = 0
			} else {
				rate = 1000000000.0 * float64(len(times)-timesIdx) / float64(time.Now().UnixNano()-times[timesIdx])
			}

			go func(id string, r float64) {
				loghub.UI <- &loghub.LogMsg{
					Type: loghub.UPDATE_RATE,
					Data: map[string]interface{}{
						"Rate": r,
					},
					Id: id,
				}
			}(c.Id, rate)

		case msg := <-c.InChan:
			last = msg.Msg
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg.Msg,
					Route: c.ToRoute,
				}
			}

			times = times[1:]
			times = append(times, time.Now().UnixNano())

			if timesIdx > 0 {
				timesIdx--
			}

		case msg := <-c.QueryChan:
			switch msg.Route {
			case "rate":
				msg.MsgChan <- map[string]interface{}{
					"Rate": rate,
				}
			case "last":
				msg.MsgChan <- map[string]interface{}{
					"Last": last,
				}
			}
		case msg := <-c.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-c.DelChan:
			delete(outChans, msg.Route)
		case <-c.QuitChan:
			c.CleanUp()
			return
		}
	}
}
<commit_msg>tidied up queryparam routes in light of recent updates<commit_after>package blocks

import (
	"fmt"
	"github.com/nytlabs/streamtools/st/loghub"
	"log"
	"net/url"
	"reflect"
	"strconv"
	"time"
)

type Msg struct {
	Msg   interface{}
	Route string
}

type MsgChan chan interface{}

func (c MsgChan) MarshalJSON() ([]byte, error) {
	outBytes := []byte("{\"channel\":" + strconv.Itoa(len(c)) + "}")
	return outBytes, nil
}

type AddChanMsg struct {
	Route   string
	Channel chan *Msg
}

type QueryMsg struct {
	Route   string
	MsgChan MsgChan
}

type QueryParamMsg struct {
	Route    string
	RespChan chan interface{}
	Params   url.Values
}

type Query struct {
	RespChan chan interface{}
	Params   url.Values
}

type BlockChans struct {
	InChan         chan *Msg
	QueryChan      chan *QueryMsg
	QueryParamChan chan *QueryParamMsg
	AddChan        chan *AddChanMsg
	DelChan        chan *Msg
	ErrChan        chan error
	QuitChan       chan bool
}

type LogStreams struct {
	log MsgChan
	ui  MsgChan
}

type Block struct {
	Id               string // the name of the block specifed by the user (like MyBlock)
	Kind             string // the kind of block this is (like count, toFile, fromSQS)
	Desc             string // the description of block ('counts the number of messages it has seen')
	inRoutes         map[string]MsgChan
	queryRoutes      map[string]chan MsgChan
	queryParamRoutes map[string]chan Query
	broadcast        MsgChan
	quit             MsgChan
	doesBroadcast    bool
	BlockChans
	LogStreams
}

type BlockDef struct {
	Type             string
	Desc             string
	InRoutes         []string
	QueryRoutes      []string
	QueryParamRoutes []string
	OutRoutes        []string
}

type BlockInterface interface {
	Setup()
	Run()
	CleanUp()
	Build(BlockChans)
	Quit() MsgChan
	Broadcast() MsgChan
	InRoute(string) MsgChan
	QueryRoute(string) chan MsgChan
	QueryParamRoute(string) chan Query
	GetBlock() *Block
	GetDef() *BlockDef
	Log(interface{})
	Error(interface{})
	SetId(string)
}

func (b *Block) Build(c BlockChans) {
	// block channels
	b.InChan = c.InChan
	b.QueryChan = c.QueryChan
	b.QueryParamChan = c.QueryParamChan
	b.AddChan = c.AddChan
	b.DelChan = c.DelChan
	b.ErrChan = c.ErrChan
	b.QuitChan = c.QuitChan

	// route maps
	b.inRoutes = make(map[string]MsgChan) // necessary to stop locking...
	b.queryRoutes = make(map[string]chan MsgChan)
	b.queryParamRoutes = make(map[string]chan Query)

	// broadcast channel
	b.broadcast = make(MsgChan, 10) // necessary to stop locking...

	// quit chan
	b.quit = make(MsgChan)

	b.ui = make(MsgChan)
	b.log = make(MsgChan)
}

func (b *Block) SetId(Id string) {
	b.Id = Id
}

func (b *Block) InRoute(routeName string) MsgChan {
	route := make(MsgChan, 1000)
	b.inRoutes[routeName] = route
	return route
}

func (b *Block) QueryRoute(routeName string) chan MsgChan {
	route := make(chan MsgChan, 1000)
	b.queryRoutes[routeName] = route
	return route
}

func (b *Block) QueryParamRoute(routeName string) chan Query {
	route := make(chan Query, 1000)
	b.queryParamRoutes[routeName] = route
	return route
}

func (b *Block) Broadcast() MsgChan {
	b.doesBroadcast = true
	return b.broadcast
}

func (b *Block) Quit() MsgChan {
	return b.quit
}

func (b *Block) GetBlock() *Block {
	return b
}

func (b *Block) GetDef() *BlockDef {
	inRoutes := []string{}
	queryRoutes := []string{}
	queryParamRoutes := []string{}
	outRoutes := []string{}

	for k, _ := range b.inRoutes {
		inRoutes = append(inRoutes, k)
	}

	for k, _ := range b.queryRoutes {
		queryRoutes = append(queryRoutes, k)
	}

	for k, _ := range b.queryParamRoutes {
		queryParamRoutes = append(queryParamRoutes, k)
	}

	if b.doesBroadcast {
		outRoutes = []string{"out"}
	}

	return &BlockDef{
		Type:             b.Kind,
		Desc:             b.Desc,
		InRoutes:         inRoutes,
		QueryRoutes:      queryRoutes,
		QueryParamRoutes: queryParamRoutes,
		OutRoutes:        outRoutes,
	}
}

func (b *Block) CleanUp() {
	for route := range b.inRoutes {
		defer close(b.inRoutes[route])
	}
	for route := range b.queryRoutes {
		defer close(b.queryRoutes[route])
	}
	defer close(b.InChan)
	defer close(b.QueryChan)
	defer close(b.QueryParamChan)
	defer close(b.AddChan)
	defer close(b.DelChan)
	defer close(b.ErrChan)
	defer close(b.QuitChan)
	defer close(b.broadcast)

	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: fmt.Sprintf("Block %s Quitting...", b.Id),
			Id:   id,
		}
	}(b.Id)
}

func (b *Block) Error(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.ERROR,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func (b *Block) Log(msg interface{}) {
	go func(id string) {
		loghub.Log <- &loghub.LogMsg{
			Type: loghub.INFO,
			Data: msg,
			Id:   id,
		}
	}(b.Id)
}

func BlockRoutine(bi BlockInterface) {
	var dropped int64
	dropTicker := time.NewTicker(time.Duration(1 * time.Second))
	dropTicker.Stop()

	outChans := make(map[string]chan *Msg)
	b := bi.GetBlock()
	bi.Setup()
	go bi.Run()

	for {
		select {
		case <-dropTicker.C:
			go func(id string, count int64) {
				loghub.Log <- &loghub.LogMsg{
					Type: loghub.ERROR,
					Data: fmt.Sprintf("Dropped messages: %d (cannot keep up with stream)", count),
					Id:   id,
				}
			}(b.Id, dropped)

			if dropped == 0 {
				dropTicker.Stop()
			}

			dropped = 0
		case msg := <-b.InChan:
			_, ok := b.inRoutes[msg.Route]
			if !ok {
				break
			}

			// every in channel is buffered a 1000 messages.
			// if we cannot immediately send to that in channel we place the msg
			// in a go routine and notify the user that the block routine's
			// buffer has overflowed. this still allows for unrecoverable
			// overflows (for example: a stuck run() function), but at least it
			// offloads to memory/cpu instead of blocking.
			select {
			case b.inRoutes[msg.Route] <- msg.Msg:
			default:
				if dropped == 0 {
					dropTicker.Stop()
					dropTicker = time.NewTicker(1 * time.Second)
				}

				dropped++
			}

			if msg.Route == "rule" {
				go func(id string) {
					loghub.UI <- &loghub.LogMsg{
						Type: loghub.RULE_UPDATED,
						Data: map[string]interface{}{},
						Id:   id,
					}
				}(b.Id)
			}

		case msg := <-b.QueryChan:

			if msg.Route == "ping" {
				msg.MsgChan <- "OK"
				continue
			}

			_, ok := b.queryRoutes[msg.Route]
			if !ok {
				break
			}

			select {
			case b.queryRoutes[msg.Route] <- msg.MsgChan:
			default:
				go func() {
					b.queryRoutes[msg.Route] <- msg.MsgChan
				}()
			}
		case msg := <-b.QueryParamChan:

			if msg.Route == "ping" {
				msg.RespChan <- "OK"
				continue
			}

			_, ok := b.queryParamRoutes[msg.Route]
			if !ok {
				break
			}
			q := Query{
				RespChan: msg.RespChan,
				Params:   msg.Params,
			}

			select {
			case b.queryParamRoutes[msg.Route] <- q:
			default:
				go func() {
					b.queryParamRoutes[msg.Route] <- q
				}()
			}

		case msg := <-b.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-b.DelChan:
			delete(outChans, msg.Route)
		case msg := <-b.broadcast:
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg,
					Route: "",
				}
			}
		case <-b.QuitChan:
			b.quit <- true
			b.CleanUp()
			return
		}
	}
}

type Connection struct {
	Id      string
	ToRoute string
	BlockChans
	LogStreams
}

func (c *Connection) SetId(Id string) {
	c.Id = Id
}

func (c *Connection) Build(chans BlockChans) {
	c.InChan = chans.InChan
	c.QueryChan = chans.QueryChan
	c.QueryParamChan = chans.QueryParamChan
	c.AddChan = chans.AddChan
	c.DelChan = chans.DelChan
	c.QuitChan = chans.QuitChan
}

func (c *Connection) CleanUp() {
	defer close(c.InChan)
	defer close(c.QueryChan)
	defer close(c.QueryParamChan)
	defer close(c.AddChan)
	defer close(c.DelChan)
	defer close(c.QuitChan)

	loghub.Log <- &loghub.LogMsg{
		Type: loghub.INFO,
		Data: fmt.Sprintf("Connection %s Quitting...", c.Id),
		Id:   c.Id,
	}
}

func ConnectionRoutine(c *Connection) {
	var last interface{}
	var rate float64

	outChans := make(map[string]chan *Msg)
	times := make([]int64, 100, 100)
	timesIdx := len(times)
	rateReport := time.NewTicker(200 * time.Millisecond)

	for {
		select {
		case <-rateReport.C:
			if timesIdx == len(times) {
				rate = 0
			} else {
				rate = 1000000000.0 * float64(len(times)-timesIdx) / float64(time.Now().UnixNano()-times[timesIdx])
			}

			go func(id string, r float64) {
				loghub.UI <- &loghub.LogMsg{
					Type: loghub.UPDATE_RATE,
					Data: map[string]interface{}{
						"Rate": r,
					},
					Id: id,
				}
			}(c.Id, rate)

		case msg := <-c.InChan:
			last = msg.Msg
			for _, v := range outChans {
				v <- &Msg{
					Msg:   msg.Msg,
					Route: c.ToRoute,
				}
			}

			times = times[1:]
			times = append(times, time.Now().UnixNano())

			if timesIdx > 0 {
				timesIdx--
			}

		case msg := <-c.QueryChan:
			switch msg.Route {
			case "rate":
				msg.MsgChan <- map[string]interface{}{
					"Rate": rate,
				}
			case "last":
				msg.MsgChan <- map[string]interface{}{
					"Last": last,
				}
			}
		case msg := <-c.AddChan:
			outChans[msg.Route] = msg.Channel
		case msg := <-c.DelChan:
			delete(outChans, msg.Route)
		case <-c.QuitChan:
			c.CleanUp()
			return
		}
	}
}
<|endoftext|>


<commit_before>package scanner

import (
	"crypto/sha512"
	"errors"
	"github.com/Symantec/Dominator/sub/fsrateio"
	"io"
	"os"
	"path"
	"sort"
	"syscall"
)

func (fileSystem *FileSystem) getInode(stat *syscall.Stat_t) (*Inode, bool) {
	inode := fileSystem.InodeTable[stat.Ino]
	new := false
	if inode == nil {
		var _inode Inode
		inode = &_inode
		_inode.stat = *stat
		fileSystem.InodeTable[stat.Ino] = inode
		new = true
	}
	return inode, new
}

func (directory *Directory) scan(fileSystem *FileSystem,
	parentName string) error {
	myPathName := path.Join(parentName, directory.name)
	file, err := os.Open(myPathName)
	if err != nil {
		return err
	}
	names, err := file.Readdirnames(-1)
	if err != nil {
		return err
	}
	file.Close()
	sort.Strings(names)
	for _, name := range names {
		filename := path.Join(myPathName, name)
		var stat syscall.Stat_t
		err := syscall.Lstat(filename, &stat)
		if err != nil {
			if err == syscall.ENOENT {
				continue
			}
			return err
		}
		inode, isNewInode := fileSystem.getInode(&stat)
		if stat.Dev == directory.inode.stat.Dev {
			if stat.Mode&syscall.S_IFMT == syscall.S_IFDIR {
				if !isNewInode {
					return errors.New("Hardlinked directory: " + filename)
				}
				var dir Directory
				dir.name = name
				dir.inode = inode
				err := dir.scan(fileSystem, myPathName)
				if err != nil {
					return err
				}
				directory.DirectoryList = append(directory.DirectoryList, &dir)
			} else {
				var file File
				file.name = name
				file.inode = inode
				if isNewInode {
					err := file.scan(fileSystem, myPathName)
					if err != nil {
						if err == syscall.ENOENT {
							continue
						}
						return err
					}
				}
				directory.FileList = append(directory.FileList, &file)
			}
		}
	}
	return nil
}

func (file *File) scan(fileSystem *FileSystem, parentName string) error {
	myPathName := path.Join(parentName, file.name)
	if file.inode.stat.Mode&syscall.S_IFMT == syscall.S_IFREG {
		if len(file.inode.hash) > 0 {
			return nil
		}
		f, err := os.Open(myPathName)
		if err != nil {
			return err
		}
		reader := fsrateio.NewReader(f, fileSystem.ctx)
		hash := sha512.New()
		io.Copy(hash, reader)
		f.Close()
		file.inode.hash = hash.Sum(nil)
	} else if file.inode.stat.Mode&syscall.S_IFMT == syscall.S_IFLNK {
		symlink, err := os.Readlink(myPathName)
		if err != nil {
			return err
		}
		file.inode.symlink = symlink
	}
	return nil
}
<commit_msg>Remove obsolete check for existing hash scanner:file.scan().<commit_after>package scanner

import (
	"crypto/sha512"
	"errors"
	"github.com/Symantec/Dominator/sub/fsrateio"
	"io"
	"os"
	"path"
	"sort"
	"syscall"
)

func (fileSystem *FileSystem) getInode(stat *syscall.Stat_t) (*Inode, bool) {
	inode := fileSystem.InodeTable[stat.Ino]
	new := false
	if inode == nil {
		var _inode Inode
		inode = &_inode
		_inode.stat = *stat
		fileSystem.InodeTable[stat.Ino] = inode
		new = true
	}
	return inode, new
}

func (directory *Directory) scan(fileSystem *FileSystem,
	parentName string) error {
	myPathName := path.Join(parentName, directory.name)
	file, err := os.Open(myPathName)
	if err != nil {
		return err
	}
	names, err := file.Readdirnames(-1)
	if err != nil {
		return err
	}
	file.Close()
	sort.Strings(names)
	for _, name := range names {
		filename := path.Join(myPathName, name)
		var stat syscall.Stat_t
		err := syscall.Lstat(filename, &stat)
		if err != nil {
			if err == syscall.ENOENT {
				continue
			}
			return err
		}
		inode, isNewInode := fileSystem.getInode(&stat)
		if stat.Dev == directory.inode.stat.Dev {
			if stat.Mode&syscall.S_IFMT == syscall.S_IFDIR {
				if !isNewInode {
					return errors.New("Hardlinked directory: " + filename)
				}
				var dir Directory
				dir.name = name
				dir.inode = inode
				err := dir.scan(fileSystem, myPathName)
				if err != nil {
					return err
				}
				directory.DirectoryList = append(directory.DirectoryList, &dir)
			} else {
				var file File
				file.name = name
				file.inode = inode
				if isNewInode {
					err := file.scan(fileSystem, myPathName)
					if err != nil {
						if err == syscall.ENOENT {
							continue
						}
						return err
					}
				}
				directory.FileList = append(directory.FileList, &file)
			}
		}
	}
	return nil
}

func (file *File) scan(fileSystem *FileSystem, parentName string) error {
	myPathName := path.Join(parentName, file.name)
	if file.inode.stat.Mode&syscall.S_IFMT == syscall.S_IFREG {
		f, err := os.Open(myPathName)
		if err != nil {
			return err
		}
		reader := fsrateio.NewReader(f, fileSystem.ctx)
		hash := sha512.New()
		io.Copy(hash, reader)
		f.Close()
		file.inode.hash = hash.Sum(nil)
	} else if file.inode.stat.Mode&syscall.S_IFMT == syscall.S_IFLNK {
		symlink, err := os.Readlink(myPathName)
		if err != nil {
			return err
		}
		file.inode.symlink = symlink
	}
	return nil
}
<|endoftext|>


<commit_before>package water

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"sync"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows/registry"
)

// To use it with windows, you need a tap driver installed on windows.
// https://github.com/OpenVPN/tap-windows6
// or just install OpenVPN
// https://github.com/OpenVPN/openvpn

const (
	// tapDriverKey is a location of the TAP driver key.
	tapDriverKey = `SYSTEM\CurrentControlSet\Control\Class\{4D36E972-E325-11CE-BFC1-08002BE10318}`
)

var (
	errIfceNameNotFound = errors.New("Failed to find the name of interface")
	// Device Control Codes
	tap_win_ioctl_get_mac             = tap_control_code(1, 0)
	tap_win_ioctl_get_version         = tap_control_code(2, 0)
	tap_win_ioctl_get_mtu             = tap_control_code(3, 0)
	tap_win_ioctl_get_info            = tap_control_code(4, 0)
	tap_ioctl_config_point_to_point   = tap_control_code(5, 0)
	tap_ioctl_set_media_status        = tap_control_code(6, 0)
	tap_win_ioctl_config_dhcp_masq    = tap_control_code(7, 0)
	tap_win_ioctl_get_log_line        = tap_control_code(8, 0)
	tap_win_ioctl_config_dhcp_set_opt = tap_control_code(9, 0)
	tap_ioctl_config_tun              = tap_control_code(10, 0)
	// w32 api
	file_device_unknown = uint32(0x00000022)
	nCreateEvent,
	nResetEvent,
	nGetOverlappedResult uintptr
)

func init() {
	k32, err := syscall.LoadLibrary("kernel32.dll")
	if err != nil {
		panic("LoadLibrary " + err.Error())
	}
	defer syscall.FreeLibrary(k32)

	nCreateEvent = getProcAddr(k32, "CreateEventW")
	nResetEvent = getProcAddr(k32, "ResetEvent")
	nGetOverlappedResult = getProcAddr(k32, "GetOverlappedResult")
}

func getProcAddr(lib syscall.Handle, name string) uintptr {
	addr, err := syscall.GetProcAddress(lib, name)
	if err != nil {
		panic(name + " " + err.Error())
	}
	return addr
}

func resetEvent(h syscall.Handle) error {
	r, _, err := syscall.Syscall(nResetEvent, 1, uintptr(h), 0, 0)
	if r == 0 {
		return err
	}
	return nil
}

func getOverlappedResult(h syscall.Handle, overlapped *syscall.Overlapped) (int, error) {
	var n int
	r, _, err := syscall.Syscall6(nGetOverlappedResult, 4,
		uintptr(h),
		uintptr(unsafe.Pointer(overlapped)),
		uintptr(unsafe.Pointer(&n)), 1, 0, 0)
	if r == 0 {
		return n, err
	}

	return n, nil
}

func newOverlapped() (*syscall.Overlapped, error) {
	var overlapped syscall.Overlapped
	r, _, err := syscall.Syscall6(nCreateEvent, 4, 0, 1, 0, 0, 0, 0)
	if r == 0 {
		return nil, err
	}
	overlapped.HEvent = syscall.Handle(r)
	return &overlapped, nil
}

type wfile struct {
	fd syscall.Handle
	rl sync.Mutex
	wl sync.Mutex
	ro *syscall.Overlapped
	wo *syscall.Overlapped
}

func (f *wfile) Close() error {
	return syscall.Close(f.fd)
}

func (f *wfile) Write(b []byte) (int, error) {
	f.wl.Lock()
	defer f.wl.Unlock()

	if err := resetEvent(f.wo.HEvent); err != nil {
		return 0, err
	}
	var n uint32
	err := syscall.WriteFile(f.fd, b, &n, f.wo)
	if err != nil && err != syscall.ERROR_IO_PENDING {
		return int(n), err
	}
	return getOverlappedResult(f.fd, f.wo)
}

func (f *wfile) Read(b []byte) (int, error) {
	f.rl.Lock()
	defer f.rl.Unlock()

	if err := resetEvent(f.ro.HEvent); err != nil {
		return 0, err
	}
	var done uint32
	err := syscall.ReadFile(f.fd, b, &done, f.ro)
	if err != nil && err != syscall.ERROR_IO_PENDING {
		return int(done), err
	}
	return getOverlappedResult(f.fd, f.ro)
}

func ctl_code(device_type, function, method, access uint32) uint32 {
	return (device_type << 16) | (access << 14) | (function << 2) | method
}

func tap_control_code(request, method uint32) uint32 {
	return ctl_code(file_device_unknown, request, method, 0)
}

// getdeviceid finds out a TAP device from registry, it *may* requires privileged right to prevent some weird issue.
func getdeviceid(componentID string, interfaceName string) (deviceid string, err error) {
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, tapDriverKey, registry.READ)
	if err != nil {
		return "", fmt.Errorf("Failed to open the adapter registry, TAP driver may be not installed, %v", err)
	}
	defer k.Close()
	// read all subkeys, it should not return an err here
	keys, err := k.ReadSubKeyNames(-1)
	if err != nil {
		return "", err
	}
	// find the one matched ComponentId
	for _, v := range keys {
		key, err := registry.OpenKey(registry.LOCAL_MACHINE, tapDriverKey+"\\"+v, registry.READ)
		if err != nil {
			continue
		}
		val, _, err := key.GetStringValue("ComponentId")
		if err != nil {
			key.Close()
			continue
		}
		if val == componentID {
			val, _, err = key.GetStringValue("NetCfgInstanceId")
			if err != nil {
				key.Close()
				continue
			}
			if len(interfaceName) > 0 {
				key2 := fmt.Sprintf("%s\\%s\\Connection", tapDriverKey, val)
				k2, err := registry.OpenKey(registry.LOCAL_MACHINE, key2, registry.READ)
				if err != nil {
					continue
				}
				defer k2.Close()
				val, _, err := k2.GetStringValue("Name")
				if err != nil || val != interfaceName {
					continue
				}
			}
			key.Close()
			return val, nil
		}
		key.Close()
	}
	if len(interfaceName) > 0 {
		return "", fmt.Errorf("Failed to find the tap device in registry with specified ComponentId '%s' and InterfaceName '%s', TAP driver may be not installed or you may have specified an interface name that doesn't exist", componentID, interfaceName)
	}

	return "", fmt.Errorf("Failed to find the tap device in registry with specified ComponentId '%s', TAP driver may be not installed", componentID)
}

// setStatus is used to bring up or bring down the interface
func setStatus(fd syscall.Handle, status bool) error {
	var bytesReturned uint32
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)
	code := []byte{0x00, 0x00, 0x00, 0x00}
	if status {
		code[0] = 0x01
	}
	return syscall.DeviceIoControl(fd, tap_ioctl_set_media_status, &code[0], uint32(4), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
}

// setTUN is used to configure the IP address in the underlying driver when using TUN
func setTUN(fd syscall.Handle, network string) error {
	var bytesReturned uint32
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)

	localIP, remoteNet, err := net.ParseCIDR(network)
	if err != nil {
		return fmt.Errorf("Failed to parse network CIDR in config, %v", err)
	}
	if localIP.To4() == nil {
		return fmt.Errorf("Provided network(%s) is not a valid IPv4 address", network)
	}
	code2 := make([]byte, 0, 12)
	code2 = append(code2, localIP.To4()[:4]...)
	code2 = append(code2, remoteNet.IP.To4()[:4]...)
	code2 = append(code2, remoteNet.Mask[:4]...)
	if len(code2) != 12 {
		return fmt.Errorf("Provided network(%s) is not valid", network)
	}
	if err := syscall.DeviceIoControl(fd, tap_ioctl_config_tun, &code2[0], uint32(12), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil); err != nil {
		return err
	}
	return nil
}

// openDev find and open an interface.
func openDev(config Config) (ifce *Interface, err error) {
	// find the device in registry.
	deviceid, err := getdeviceid(config.PlatformSpecificParams.ComponentID, config.PlatformSpecificParams.InterfaceName)
	if err != nil {
		return nil, err
	}
	path := "\\\\.\\Global\\" + deviceid + ".tap"
	pathp, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return nil, err
	}
	// type Handle uintptr
	file, err := syscall.CreateFile(pathp, syscall.GENERIC_READ|syscall.GENERIC_WRITE, uint32(syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE), nil, syscall.OPEN_EXISTING, syscall.FILE_ATTRIBUTE_SYSTEM|syscall.FILE_FLAG_OVERLAPPED, 0)
	// if err hanppens, close the interface.
	defer func() {
		if err != nil {
			syscall.Close(file)
		}
		if err := recover(); err != nil {
			syscall.Close(file)
		}
	}()
	if err != nil {
		return nil, err
	}
	var bytesReturned uint32

	// find the mac address of tap device, use this to find the name of interface
	mac := make([]byte, 6)
	err = syscall.DeviceIoControl(file, tap_win_ioctl_get_mac, &mac[0], uint32(len(mac)), &mac[0], uint32(len(mac)), &bytesReturned, nil)
	if err != nil {
		return nil, err
	}

	// fd := os.NewFile(uintptr(file), path)
	ro, err := newOverlapped()
	if err != nil {
		return
	}
	wo, err := newOverlapped()
	if err != nil {
		return
	}
	fd := &wfile{fd: file, ro: ro, wo: wo}
	ifce = &Interface{isTAP: (config.DeviceType == TAP), ReadWriteCloser: fd}

	// bring up device.
	if err := setStatus(file, true); err != nil {
		return nil, err
	}

	//TUN
	if config.DeviceType == TUN {
		if err := setTUN(file, config.PlatformSpecificParams.Network); err != nil {
			return nil, err
		}
	}

	// find the name of tap interface(u need it to set the ip or other command)
	ifces, err := net.Interfaces()
	if err != nil {
		return
	}

	for _, v := range ifces {
		if bytes.Equal(v.HardwareAddr[:6], mac[:6]) {
			ifce.name = v.Name
			return
		}
	}

	return nil, errIfceNameNotFound
}
<commit_msg>Fix Windows interface selection by name as changed in #59<commit_after>package water

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"sync"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows/registry"
)

// To use it with windows, you need a tap driver installed on windows.
// https://github.com/OpenVPN/tap-windows6
// or just install OpenVPN
// https://github.com/OpenVPN/openvpn

const (
	// tapDriverKey is the location of the TAP driver key.
	tapDriverKey = `SYSTEM\CurrentControlSet\Control\Class\{4D36E972-E325-11CE-BFC1-08002BE10318}`
	// netConfigKey is the location of the TAP adapter's network config.
	netConfigKey = `SYSTEM\CurrentControlSet\Control\Network\{4D36E972-E325-11CE-BFC1-08002BE10318}`
)

var (
	errIfceNameNotFound = errors.New("Failed to find the name of interface")
	// Device Control Codes
	tap_win_ioctl_get_mac             = tap_control_code(1, 0)
	tap_win_ioctl_get_version         = tap_control_code(2, 0)
	tap_win_ioctl_get_mtu             = tap_control_code(3, 0)
	tap_win_ioctl_get_info            = tap_control_code(4, 0)
	tap_ioctl_config_point_to_point   = tap_control_code(5, 0)
	tap_ioctl_set_media_status        = tap_control_code(6, 0)
	tap_win_ioctl_config_dhcp_masq    = tap_control_code(7, 0)
	tap_win_ioctl_get_log_line        = tap_control_code(8, 0)
	tap_win_ioctl_config_dhcp_set_opt = tap_control_code(9, 0)
	tap_ioctl_config_tun              = tap_control_code(10, 0)
	// w32 api
	file_device_unknown = uint32(0x00000022)
	nCreateEvent,
	nResetEvent,
	nGetOverlappedResult uintptr
)

func init() {
	k32, err := syscall.LoadLibrary("kernel32.dll")
	if err != nil {
		panic("LoadLibrary " + err.Error())
	}
	defer syscall.FreeLibrary(k32)

	nCreateEvent = getProcAddr(k32, "CreateEventW")
	nResetEvent = getProcAddr(k32, "ResetEvent")
	nGetOverlappedResult = getProcAddr(k32, "GetOverlappedResult")
}

func getProcAddr(lib syscall.Handle, name string) uintptr {
	addr, err := syscall.GetProcAddress(lib, name)
	if err != nil {
		panic(name + " " + err.Error())
	}
	return addr
}

func resetEvent(h syscall.Handle) error {
	r, _, err := syscall.Syscall(nResetEvent, 1, uintptr(h), 0, 0)
	if r == 0 {
		return err
	}
	return nil
}

func getOverlappedResult(h syscall.Handle, overlapped *syscall.Overlapped) (int, error) {
	var n int
	r, _, err := syscall.Syscall6(nGetOverlappedResult, 4,
		uintptr(h),
		uintptr(unsafe.Pointer(overlapped)),
		uintptr(unsafe.Pointer(&n)), 1, 0, 0)
	if r == 0 {
		return n, err
	}

	return n, nil
}

func newOverlapped() (*syscall.Overlapped, error) {
	var overlapped syscall.Overlapped
	r, _, err := syscall.Syscall6(nCreateEvent, 4, 0, 1, 0, 0, 0, 0)
	if r == 0 {
		return nil, err
	}
	overlapped.HEvent = syscall.Handle(r)
	return &overlapped, nil
}

type wfile struct {
	fd syscall.Handle
	rl sync.Mutex
	wl sync.Mutex
	ro *syscall.Overlapped
	wo *syscall.Overlapped
}

func (f *wfile) Close() error {
	return syscall.Close(f.fd)
}

func (f *wfile) Write(b []byte) (int, error) {
	f.wl.Lock()
	defer f.wl.Unlock()

	if err := resetEvent(f.wo.HEvent); err != nil {
		return 0, err
	}
	var n uint32
	err := syscall.WriteFile(f.fd, b, &n, f.wo)
	if err != nil && err != syscall.ERROR_IO_PENDING {
		return int(n), err
	}
	return getOverlappedResult(f.fd, f.wo)
}

func (f *wfile) Read(b []byte) (int, error) {
	f.rl.Lock()
	defer f.rl.Unlock()

	if err := resetEvent(f.ro.HEvent); err != nil {
		return 0, err
	}
	var done uint32
	err := syscall.ReadFile(f.fd, b, &done, f.ro)
	if err != nil && err != syscall.ERROR_IO_PENDING {
		return int(done), err
	}
	return getOverlappedResult(f.fd, f.ro)
}

func ctl_code(device_type, function, method, access uint32) uint32 {
	return (device_type << 16) | (access << 14) | (function << 2) | method
}

func tap_control_code(request, method uint32) uint32 {
	return ctl_code(file_device_unknown, request, method, 0)
}

// getdeviceid finds out a TAP device from registry, it *may* requires privileged right to prevent some weird issue.
func getdeviceid(componentID string, interfaceName string) (deviceid string, err error) {
	k, err := registry.OpenKey(registry.LOCAL_MACHINE, tapDriverKey, registry.READ)
	if err != nil {
		return "", fmt.Errorf("Failed to open the adapter registry, TAP driver may be not installed, %v", err)
	}
	defer k.Close()
	// read all subkeys, it should not return an err here
	keys, err := k.ReadSubKeyNames(-1)
	if err != nil {
		return "", err
	}
	// find the one matched ComponentId
	for _, v := range keys {
		key, err := registry.OpenKey(registry.LOCAL_MACHINE, tapDriverKey+"\\"+v, registry.READ)
		if err != nil {
			continue
		}
		val, _, err := key.GetStringValue("ComponentId")
		if err != nil {
			key.Close()
			continue
		}
		if val == componentID {
			val, _, err = key.GetStringValue("NetCfgInstanceId")
			if err != nil {
				key.Close()
				continue
			}
			if len(interfaceName) > 0 {
				key2 := fmt.Sprintf("%s\\%s\\Connection", netConfigKey, val)
				k2, err := registry.OpenKey(registry.LOCAL_MACHINE, key2, registry.READ)
				if err != nil {
					continue
				}
				defer k2.Close()
				val, _, err := k2.GetStringValue("Name")
				if err != nil || val != interfaceName {
					continue
				}
			}
			key.Close()
			return val, nil
		}
		key.Close()
	}
	if len(interfaceName) > 0 {
		return "", fmt.Errorf("Failed to find the tap device in registry with specified ComponentId '%s' and InterfaceName '%s', TAP driver may be not installed or you may have specified an interface name that doesn't exist", componentID, interfaceName)
	}

	return "", fmt.Errorf("Failed to find the tap device in registry with specified ComponentId '%s', TAP driver may be not installed", componentID)
}

// setStatus is used to bring up or bring down the interface
func setStatus(fd syscall.Handle, status bool) error {
	var bytesReturned uint32
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)
	code := []byte{0x00, 0x00, 0x00, 0x00}
	if status {
		code[0] = 0x01
	}
	return syscall.DeviceIoControl(fd, tap_ioctl_set_media_status, &code[0], uint32(4), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil)
}

// setTUN is used to configure the IP address in the underlying driver when using TUN
func setTUN(fd syscall.Handle, network string) error {
	var bytesReturned uint32
	rdbbuf := make([]byte, syscall.MAXIMUM_REPARSE_DATA_BUFFER_SIZE)

	localIP, remoteNet, err := net.ParseCIDR(network)
	if err != nil {
		return fmt.Errorf("Failed to parse network CIDR in config, %v", err)
	}
	if localIP.To4() == nil {
		return fmt.Errorf("Provided network(%s) is not a valid IPv4 address", network)
	}
	code2 := make([]byte, 0, 12)
	code2 = append(code2, localIP.To4()[:4]...)
	code2 = append(code2, remoteNet.IP.To4()[:4]...)
	code2 = append(code2, remoteNet.Mask[:4]...)
	if len(code2) != 12 {
		return fmt.Errorf("Provided network(%s) is not valid", network)
	}
	if err := syscall.DeviceIoControl(fd, tap_ioctl_config_tun, &code2[0], uint32(12), &rdbbuf[0], uint32(len(rdbbuf)), &bytesReturned, nil); err != nil {
		return err
	}
	return nil
}

// openDev find and open an interface.
func openDev(config Config) (ifce *Interface, err error) {
	// find the device in registry.
	deviceid, err := getdeviceid(config.PlatformSpecificParams.ComponentID, config.PlatformSpecificParams.InterfaceName)
	if err != nil {
		return nil, err
	}
	path := "\\\\.\\Global\\" + deviceid + ".tap"
	pathp, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return nil, err
	}
	// type Handle uintptr
	file, err := syscall.CreateFile(pathp, syscall.GENERIC_READ|syscall.GENERIC_WRITE, uint32(syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE), nil, syscall.OPEN_EXISTING, syscall.FILE_ATTRIBUTE_SYSTEM|syscall.FILE_FLAG_OVERLAPPED, 0)
	// if err hanppens, close the interface.
	defer func() {
		if err != nil {
			syscall.Close(file)
		}
		if err := recover(); err != nil {
			syscall.Close(file)
		}
	}()
	if err != nil {
		return nil, err
	}
	var bytesReturned uint32

	// find the mac address of tap device, use this to find the name of interface
	mac := make([]byte, 6)
	err = syscall.DeviceIoControl(file, tap_win_ioctl_get_mac, &mac[0], uint32(len(mac)), &mac[0], uint32(len(mac)), &bytesReturned, nil)
	if err != nil {
		return nil, err
	}

	// fd := os.NewFile(uintptr(file), path)
	ro, err := newOverlapped()
	if err != nil {
		return
	}
	wo, err := newOverlapped()
	if err != nil {
		return
	}
	fd := &wfile{fd: file, ro: ro, wo: wo}
	ifce = &Interface{isTAP: (config.DeviceType == TAP), ReadWriteCloser: fd}

	// bring up device.
	if err := setStatus(file, true); err != nil {
		return nil, err
	}

	//TUN
	if config.DeviceType == TUN {
		if err := setTUN(file, config.PlatformSpecificParams.Network); err != nil {
			return nil, err
		}
	}

	// find the name of tap interface(u need it to set the ip or other command)
	ifces, err := net.Interfaces()
	if err != nil {
		return
	}

	for _, v := range ifces {
		if bytes.Equal(v.HardwareAddr[:6], mac[:6]) {
			ifce.name = v.Name
			return
		}
	}

	return nil, errIfceNameNotFound
}
<|endoftext|>