<commit_before>// Copyright 2016-2017 Authors of Cilium
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/bpf/policymap"
	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/apierror"
	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/policy"

	"github.com/go-openapi/runtime/middleware"
	"github.com/op/go-logging"
)

func validPath(path string) bool {
	return strings.HasPrefix(path, common.GlobalLabelPrefix)
}

// findNode returns node and its parent or an error
func (d *Daemon) findNode(path string) (*policy.Node, *policy.Node) {
	var parent *policy.Node

	newPath := strings.Replace(path, common.GlobalLabelPrefix, "", 1)
	if newPath == "" {
		return d.policy.Root, nil
	}

	current := d.policy.Root
	parent = nil

	for _, nodeName := range strings.Split(newPath, ".") {
		if nodeName == "" {
			continue
		}
		if child, ok := current.Children[nodeName]; ok {
			parent = current
			current = child
		} else {
			return nil, nil
		}
	}

	return current, parent
}

func (d *Daemon) GetCachedLabelList(ID policy.NumericIdentity) ([]labels.Label, error) {
	// Check if we have the source security context in our local
	// consumable cache
	if c := d.consumableCache.Lookup(ID); c != nil {
		return c.LabelList, nil
	}

	// No cache entry or labels not available, do full lookup of labels
	// via KV store
	lbls, err := d.LookupIdentity(ID)
	if err != nil {
		return nil, err
	}

	// ID is not associated with anything, skip...
	if lbls == nil {
		return nil, nil
	}

	l := lbls.Labels.ToSlice()

	return l, nil
}

func (d *Daemon) invalidateCache() {
	d.consumableCache.Iteration++
	if d.consumableCache.Iteration == 0 {
		d.consumableCache.Iteration = 1
	}
}

func (d *Daemon) triggerPolicyUpdates(added []policy.NumericIdentity) {
	d.endpointsMU.Lock()
	defer d.endpointsMU.Unlock()

	if len(added) == 0 {
		log.Debugf("Full policy recalculation triggered")
		d.invalidateCache()
	} else {
		log.Debugf("Partial policy recalculation triggered: %d\n", added)
		// FIXME: Invalidate only cache that is affected
		d.invalidateCache()
	}

	log.Debugf("Iterating over endpoints...")

	for _, ep := range d.endpoints {
		log.Debugf("Triggering policy update for ep %+v", ep)
		err := ep.TriggerPolicyUpdates(d)
		if err != nil {
			log.Warningf("Error while handling policy updates for endpoint %s: %s\n",
				ep.StringID(), err)
			ep.LogStatus(endpoint.Failure, err.Error())
		} else {
			ep.LogStatusOK("Policy regenerated")
		}
	}

	log.Debugf("End")
}

// PolicyCanConsume calculates if the ctx allows the consumer to be consumed. This public
// function returns a SearchContextReply with the consumable decision and the tracing log
// if ctx.Trace was set.
func (d *Daemon) PolicyCanConsume(ctx *policy.SearchContext) (*policy.SearchContextReply, error) {
	buffer := new(bytes.Buffer)
	if ctx.Trace != policy.TRACE_DISABLED {
		ctx.Logging = logging.NewLogBackend(buffer, "", 0)
	}
	scr := policy.SearchContextReply{}
	d.policy.Mutex.RLock()
	scr.Decision = d.policy.Allows(ctx)
	d.policy.Mutex.RUnlock()

	if ctx.Trace != policy.TRACE_DISABLED {
		scr.Logging = buffer.Bytes()
	}
	return &scr, nil
}

type getPolicyResolve struct {
	daemon *Daemon
}

func NewGetPolicyResolveHandler(d *Daemon) GetPolicyResolveHandler {
	return &getPolicyResolve{daemon: d}
}

func (h *getPolicyResolve) Handle(params GetPolicyResolveParams) middleware.Responder {
	d := h.daemon
	buffer := new(bytes.Buffer)
	ctx := params.IdentityContext
	search := policy.SearchContext{
		Trace:   policy.TRACE_ENABLED,
		Logging: logging.NewLogBackend(buffer, "", 0),
		From:    labels.NewLabelsFromModel(ctx.From).ToSlice(),
		To:      labels.NewLabelsFromModel(ctx.To).ToSlice(),
	}

	d.policy.Mutex.RLock()
	verdict := d.policy.Allows(&search)
	d.policy.Mutex.RUnlock()

	result := models.PolicyTraceResult{
		Verdict: verdict.String(),
		Log:     buffer.String(),
	}

	return NewGetPolicyResolveOK().WithPayload(&result)
}

func (d *Daemon) policyAddNode(path string, node *policy.Node) (bool, error) {
	var (
		currNode, parentNode *policy.Node
		policyModified       bool
		err                  error
	)

	if node.Name == "" {
		path, node.Name = policy.SplitNodePath(path)
	} else if strings.Contains(node.Name, ".") && node.Name != common.GlobalLabelPrefix {
		path, node.Name = policy.SplitNodePath(path + "." + node.Name)
	}

	currNode, parentNode = d.findNode(path)
	log.Debugf("Policy currNode %+v, parentNode %+v", currNode, parentNode)

	// eg. path = io.cilium.lizards.foo.db and io.cilium.lizards doesn't exist
	if (currNode == nil && parentNode == nil) ||
		// eg. path = io.cilium.lizards.foo and io.cilium.lizards.foo doesn't exist
		(currNode == nil && parentNode != nil) {

		pn := policy.NewNode("", nil)
		policyModified, err = d.policyAddNode(path, pn)
		if err != nil {
			return false, err
		}
		currNode, parentNode = d.findNode(path)
		log.Debugf("Policy currNode %+v, parentNode %+v", currNode, parentNode)
	}
	// eg. path = io.cilium
	if currNode != nil && parentNode == nil {
		if currNode.Name == node.Name {
			node.Path()
			policyModified, err = currNode.Merge(node)
			if err != nil {
				return false, err
			}
		} else {
			policyModified, err = currNode.AddChild(node.Name, node)
			if err != nil {
				return false, err
			}
		}
	} else if currNode != nil && parentNode != nil {
		// eg. path = io.cilium.lizards.db exists
		policyModified, err = currNode.AddChild(node.Name, node)
		if err != nil {
			return false, err
		}
	}

	return policyModified, nil
}

func (d *Daemon) policyAdd(path string, node *policy.Node) (bool, error) {
	d.policy.Mutex.Lock()
	defer d.policy.Mutex.Unlock()

	if modified, err := d.policyAddNode(path, node); err != nil {
		return false, err
	} else if modified {
		return modified, node.ResolveTree()
	}

	return false, nil
}

func (d *Daemon) enablePolicyEnforcement() {
	d.conf.Opts.Set(endpoint.OptionPolicy, true)

	d.endpointsMU.Lock()
	defer d.endpointsMU.Unlock()

	enablePolicy := map[string]string{endpoint.OptionPolicy: "enabled"}

	for _, ep := range d.endpoints {
		if ep.ApplyOpts(enablePolicy) {
			ep.RegenerateIfReady(d)
		}
	}
}

func (d *Daemon) PolicyAdd(path string, node *policy.Node) *apierror.ApiError {
	log.Debugf("Policy Add Request: %s %+v", path, node)

	if !strings.HasPrefix(path, common.GlobalLabelPrefix) {
		return apierror.New(PutPolicyPathInvalidPathCode,
			"Invalid path %s: must start with %s", path, common.GlobalLabelPrefix)
	}

	// Enable policy if not already enabled
	if !d.conf.Opts.IsEnabled(endpoint.OptionPolicy) {
		d.enablePolicyEnforcement()
	}

	if policyModified, err := d.policyAdd(path, node); err != nil {
		return apierror.Error(PutPolicyPathFailureCode, err)
	} else if policyModified {
		log.Info("New policy imported, regenerating...")
		d.triggerPolicyUpdates([]policy.NumericIdentity{})
	}

	return nil
}

func (d *Daemon) deleteNode(node *policy.Node, parent *policy.Node) {
	if node == d.policy.Root {
		d.policy.Root = policy.NewNode(common.GlobalLabelPrefix, nil)
		d.policy.Root.Path()
	} else {
		delete(parent.Children, node.Name)
	}
}

// PolicyDelete deletes the policy set in the given path from the policy tree. If
// cover256Sum is set it finds the rule with the respective coverage that rule from the
// node. If the path's node becomes ruleless it is removed from the tree.
func (d *Daemon) PolicyDelete(path, cover256Sum string) *apierror.ApiError {
	log.Debugf("Policy Delete Request: %s, cover256Sum %s", path, cover256Sum)

	d.policy.Mutex.Lock()
	node, parent := d.findNode(path)
	if node == nil {
		d.policy.Mutex.Unlock()
		return apierror.New(DeletePolicyPathNotFoundCode, "Policy node not found")
	}

	// Deletion request of a specific rule of a node
	if cover256Sum != "" {
		if len(cover256Sum) != policy.CoverageSHASize {
			d.policy.Mutex.Unlock()
			return apierror.New(DeletePolicyPathInvalidCode,
				"Invalid length of hash, must be %d", policy.CoverageSHASize)
		}

		for i, pr := range node.Rules {
			if prCover256Sum, err := pr.CoverageSHA256Sum(); err == nil &&
				prCover256Sum == cover256Sum {
				node.Rules = append(node.Rules[:i], node.Rules[i+1:]...)

				// If the rule was the last remaining, delete the node
				if !node.HasRules() {
					d.deleteNode(node, parent)
				}

				d.policy.Mutex.Unlock()
				d.triggerPolicyUpdates([]policy.NumericIdentity{})
				return nil
			}
		}

		d.policy.Mutex.Unlock()
		return apierror.New(DeletePolicyPathNotFoundCode, "policy not found")
	}

	// Deletion request for entire node
	d.deleteNode(node, parent)
	d.policy.Mutex.Unlock()

	d.triggerPolicyUpdates([]policy.NumericIdentity{})
	return nil
}

type deletePolicyPath struct {
	daemon *Daemon
}

func NewDeletePolicyPathHandler(d *Daemon) DeletePolicyPathHandler {
	return &deletePolicyPath{daemon: d}
}

func (h *deletePolicyPath) Handle(params DeletePolicyPathParams) middleware.Responder {
	d := h.daemon
	if err := d.PolicyDelete(params.Path, ""); err != nil {
		return apierror.Error(DeletePolicyPathFailureCode, err)
	}

	return NewDeletePolicyPathNoContent()
}

type putPolicyPath struct {
	daemon *Daemon
}

func NewPutPolicyPathHandler(d *Daemon) PutPolicyPathHandler {
	return &putPolicyPath{daemon: d}
}

func (h *putPolicyPath) Handle(params PutPolicyPathParams) middleware.Responder {
	d := h.daemon
	if !validPath(params.Path) {
		return apierror.New(PutPolicyPathInvalidPathCode,
			"path must have prefix %s", common.GlobalLabelPrefix)
	}

	var node policy.Node
	if err := json.Unmarshal([]byte(*params.Policy), &node); err != nil {
		return NewPutPolicyPathInvalidPolicy()
	}

	if err := d.PolicyAdd(params.Path, &node); err != nil {
		return apierror.Error(PutPolicyPathFailureCode, err)
	}

	return NewPutPolicyPathOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
}

type getPolicy struct {
	daemon *Daemon
}

func NewGetPolicyHandler(d *Daemon) GetPolicyHandler {
	return &getPolicy{daemon: d}
}

// Returns the entire policy tree
func (h *getPolicy) Handle(params GetPolicyParams) middleware.Responder {
	d := h.daemon
	d.policy.Mutex.RLock()
	defer d.policy.Mutex.RUnlock()
	node := d.policy.Root
	return NewGetPolicyOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
}

type getPolicyPath struct {
	daemon *Daemon
}

func NewGetPolicyPathHandler(d *Daemon) GetPolicyPathHandler {
	return &getPolicyPath{daemon: d}
}

func (h *getPolicyPath) Handle(params GetPolicyPathParams) middleware.Responder {
	d := h.daemon
	d.policy.Mutex.RLock()
	defer d.policy.Mutex.RUnlock()

	if !validPath(params.Path) {
		return apierror.New(GetPolicyPathInvalidCode,
			"path must have prefix %s", common.GlobalLabelPrefix)
	}

	if node, _ := d.findNode(params.Path); node == nil {
		return NewGetPolicyPathNotFound()
	} else {
		return NewGetPolicyPathOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
	}
}

func (d *Daemon) PolicyInit() error {
	for k, v := range policy.ReservedIdentities {
		key := policy.NumericIdentity(v).String()
		lbl := labels.NewLabel(
			key, "", common.ReservedLabelSource,
		)
		secLbl := policy.NewIdentity()
		secLbl.ID = v
		secLbl.AssociateEndpoint(lbl.String())
		secLbl.Labels[k] = lbl

		policyMapPath := bpf.MapPath(fmt.Sprintf("%sreserved_%d", policymap.MapName, int(v)))

		policyMap, _, err := policymap.OpenMap(policyMapPath)
		if err != nil {
			return fmt.Errorf("Could not create policy BPF map '%s': %s", policyMapPath, err)
		}

		if c := d.consumableCache.GetOrCreate(v, secLbl); c == nil {
			return fmt.Errorf("Unable to initialize consumable for %v", secLbl)
		} else {
			d.consumableCache.AddReserved(c)
			c.AddMap(policyMap)
		}
	}

	return nil
}
<commit_msg>daemon: changed debug message to be less verbose<commit_after>// Copyright 2016-2017 Authors of Cilium
//
// 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 main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/cilium/cilium/api/v1/models"
	. "github.com/cilium/cilium/api/v1/server/restapi/policy"
	"github.com/cilium/cilium/bpf/policymap"
	"github.com/cilium/cilium/common"
	"github.com/cilium/cilium/pkg/apierror"
	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/endpoint"
	"github.com/cilium/cilium/pkg/labels"
	"github.com/cilium/cilium/pkg/policy"

	"github.com/go-openapi/runtime/middleware"
	"github.com/op/go-logging"
)

func validPath(path string) bool {
	return strings.HasPrefix(path, common.GlobalLabelPrefix)
}

// findNode returns node and its parent or an error
func (d *Daemon) findNode(path string) (*policy.Node, *policy.Node) {
	var parent *policy.Node

	newPath := strings.Replace(path, common.GlobalLabelPrefix, "", 1)
	if newPath == "" {
		return d.policy.Root, nil
	}

	current := d.policy.Root
	parent = nil

	for _, nodeName := range strings.Split(newPath, ".") {
		if nodeName == "" {
			continue
		}
		if child, ok := current.Children[nodeName]; ok {
			parent = current
			current = child
		} else {
			return nil, nil
		}
	}

	return current, parent
}

func (d *Daemon) GetCachedLabelList(ID policy.NumericIdentity) ([]labels.Label, error) {
	// Check if we have the source security context in our local
	// consumable cache
	if c := d.consumableCache.Lookup(ID); c != nil {
		return c.LabelList, nil
	}

	// No cache entry or labels not available, do full lookup of labels
	// via KV store
	lbls, err := d.LookupIdentity(ID)
	if err != nil {
		return nil, err
	}

	// ID is not associated with anything, skip...
	if lbls == nil {
		return nil, nil
	}

	l := lbls.Labels.ToSlice()

	return l, nil
}

func (d *Daemon) invalidateCache() {
	d.consumableCache.Iteration++
	if d.consumableCache.Iteration == 0 {
		d.consumableCache.Iteration = 1
	}
}

func (d *Daemon) triggerPolicyUpdates(added []policy.NumericIdentity) {
	d.endpointsMU.Lock()
	defer d.endpointsMU.Unlock()

	if len(added) == 0 {
		log.Debugf("Full policy recalculation triggered")
		d.invalidateCache()
	} else {
		log.Debugf("Partial policy recalculation triggered: %d\n", added)
		// FIXME: Invalidate only cache that is affected
		d.invalidateCache()
	}

	log.Debugf("Iterating over endpoints...")

	for _, ep := range d.endpoints {
		log.Debugf("Triggering policy update for ep %s", ep.StringID())
		err := ep.TriggerPolicyUpdates(d)
		if err != nil {
			log.Warningf("Error while handling policy updates for endpoint %s: %s\n",
				ep.StringID(), err)
			ep.LogStatus(endpoint.Failure, err.Error())
		} else {
			ep.LogStatusOK("Policy regenerated")
		}
	}

	log.Debugf("End")
}

// PolicyCanConsume calculates if the ctx allows the consumer to be consumed. This public
// function returns a SearchContextReply with the consumable decision and the tracing log
// if ctx.Trace was set.
func (d *Daemon) PolicyCanConsume(ctx *policy.SearchContext) (*policy.SearchContextReply, error) {
	buffer := new(bytes.Buffer)
	if ctx.Trace != policy.TRACE_DISABLED {
		ctx.Logging = logging.NewLogBackend(buffer, "", 0)
	}
	scr := policy.SearchContextReply{}
	d.policy.Mutex.RLock()
	scr.Decision = d.policy.Allows(ctx)
	d.policy.Mutex.RUnlock()

	if ctx.Trace != policy.TRACE_DISABLED {
		scr.Logging = buffer.Bytes()
	}
	return &scr, nil
}

type getPolicyResolve struct {
	daemon *Daemon
}

func NewGetPolicyResolveHandler(d *Daemon) GetPolicyResolveHandler {
	return &getPolicyResolve{daemon: d}
}

func (h *getPolicyResolve) Handle(params GetPolicyResolveParams) middleware.Responder {
	d := h.daemon
	buffer := new(bytes.Buffer)
	ctx := params.IdentityContext
	search := policy.SearchContext{
		Trace:   policy.TRACE_ENABLED,
		Logging: logging.NewLogBackend(buffer, "", 0),
		From:    labels.NewLabelsFromModel(ctx.From).ToSlice(),
		To:      labels.NewLabelsFromModel(ctx.To).ToSlice(),
	}

	d.policy.Mutex.RLock()
	verdict := d.policy.Allows(&search)
	d.policy.Mutex.RUnlock()

	result := models.PolicyTraceResult{
		Verdict: verdict.String(),
		Log:     buffer.String(),
	}

	return NewGetPolicyResolveOK().WithPayload(&result)
}

func (d *Daemon) policyAddNode(path string, node *policy.Node) (bool, error) {
	var (
		currNode, parentNode *policy.Node
		policyModified       bool
		err                  error
	)

	if node.Name == "" {
		path, node.Name = policy.SplitNodePath(path)
	} else if strings.Contains(node.Name, ".") && node.Name != common.GlobalLabelPrefix {
		path, node.Name = policy.SplitNodePath(path + "." + node.Name)
	}

	currNode, parentNode = d.findNode(path)
	log.Debugf("Policy currNode %+v, parentNode %+v", currNode, parentNode)

	// eg. path = io.cilium.lizards.foo.db and io.cilium.lizards doesn't exist
	if (currNode == nil && parentNode == nil) ||
		// eg. path = io.cilium.lizards.foo and io.cilium.lizards.foo doesn't exist
		(currNode == nil && parentNode != nil) {

		pn := policy.NewNode("", nil)
		policyModified, err = d.policyAddNode(path, pn)
		if err != nil {
			return false, err
		}
		currNode, parentNode = d.findNode(path)
		log.Debugf("Policy currNode %+v, parentNode %+v", currNode, parentNode)
	}
	// eg. path = io.cilium
	if currNode != nil && parentNode == nil {
		if currNode.Name == node.Name {
			node.Path()
			policyModified, err = currNode.Merge(node)
			if err != nil {
				return false, err
			}
		} else {
			policyModified, err = currNode.AddChild(node.Name, node)
			if err != nil {
				return false, err
			}
		}
	} else if currNode != nil && parentNode != nil {
		// eg. path = io.cilium.lizards.db exists
		policyModified, err = currNode.AddChild(node.Name, node)
		if err != nil {
			return false, err
		}
	}

	return policyModified, nil
}

func (d *Daemon) policyAdd(path string, node *policy.Node) (bool, error) {
	d.policy.Mutex.Lock()
	defer d.policy.Mutex.Unlock()

	if modified, err := d.policyAddNode(path, node); err != nil {
		return false, err
	} else if modified {
		return modified, node.ResolveTree()
	}

	return false, nil
}

func (d *Daemon) enablePolicyEnforcement() {
	d.conf.Opts.Set(endpoint.OptionPolicy, true)

	d.endpointsMU.Lock()
	defer d.endpointsMU.Unlock()

	enablePolicy := map[string]string{endpoint.OptionPolicy: "enabled"}

	for _, ep := range d.endpoints {
		if ep.ApplyOpts(enablePolicy) {
			ep.RegenerateIfReady(d)
		}
	}
}

func (d *Daemon) PolicyAdd(path string, node *policy.Node) *apierror.ApiError {
	log.Debugf("Policy Add Request: %s %+v", path, node)

	if !strings.HasPrefix(path, common.GlobalLabelPrefix) {
		return apierror.New(PutPolicyPathInvalidPathCode,
			"Invalid path %s: must start with %s", path, common.GlobalLabelPrefix)
	}

	// Enable policy if not already enabled
	if !d.conf.Opts.IsEnabled(endpoint.OptionPolicy) {
		d.enablePolicyEnforcement()
	}

	if policyModified, err := d.policyAdd(path, node); err != nil {
		return apierror.Error(PutPolicyPathFailureCode, err)
	} else if policyModified {
		log.Info("New policy imported, regenerating...")
		d.triggerPolicyUpdates([]policy.NumericIdentity{})
	}

	return nil
}

func (d *Daemon) deleteNode(node *policy.Node, parent *policy.Node) {
	if node == d.policy.Root {
		d.policy.Root = policy.NewNode(common.GlobalLabelPrefix, nil)
		d.policy.Root.Path()
	} else {
		delete(parent.Children, node.Name)
	}
}

// PolicyDelete deletes the policy set in the given path from the policy tree. If
// cover256Sum is set it finds the rule with the respective coverage that rule from the
// node. If the path's node becomes ruleless it is removed from the tree.
func (d *Daemon) PolicyDelete(path, cover256Sum string) *apierror.ApiError {
	log.Debugf("Policy Delete Request: %s, cover256Sum %s", path, cover256Sum)

	d.policy.Mutex.Lock()
	node, parent := d.findNode(path)
	if node == nil {
		d.policy.Mutex.Unlock()
		return apierror.New(DeletePolicyPathNotFoundCode, "Policy node not found")
	}

	// Deletion request of a specific rule of a node
	if cover256Sum != "" {
		if len(cover256Sum) != policy.CoverageSHASize {
			d.policy.Mutex.Unlock()
			return apierror.New(DeletePolicyPathInvalidCode,
				"Invalid length of hash, must be %d", policy.CoverageSHASize)
		}

		for i, pr := range node.Rules {
			if prCover256Sum, err := pr.CoverageSHA256Sum(); err == nil &&
				prCover256Sum == cover256Sum {
				node.Rules = append(node.Rules[:i], node.Rules[i+1:]...)

				// If the rule was the last remaining, delete the node
				if !node.HasRules() {
					d.deleteNode(node, parent)
				}

				d.policy.Mutex.Unlock()
				d.triggerPolicyUpdates([]policy.NumericIdentity{})
				return nil
			}
		}

		d.policy.Mutex.Unlock()
		return apierror.New(DeletePolicyPathNotFoundCode, "policy not found")
	}

	// Deletion request for entire node
	d.deleteNode(node, parent)
	d.policy.Mutex.Unlock()

	d.triggerPolicyUpdates([]policy.NumericIdentity{})
	return nil
}

type deletePolicyPath struct {
	daemon *Daemon
}

func NewDeletePolicyPathHandler(d *Daemon) DeletePolicyPathHandler {
	return &deletePolicyPath{daemon: d}
}

func (h *deletePolicyPath) Handle(params DeletePolicyPathParams) middleware.Responder {
	d := h.daemon
	if err := d.PolicyDelete(params.Path, ""); err != nil {
		return apierror.Error(DeletePolicyPathFailureCode, err)
	}

	return NewDeletePolicyPathNoContent()
}

type putPolicyPath struct {
	daemon *Daemon
}

func NewPutPolicyPathHandler(d *Daemon) PutPolicyPathHandler {
	return &putPolicyPath{daemon: d}
}

func (h *putPolicyPath) Handle(params PutPolicyPathParams) middleware.Responder {
	d := h.daemon
	if !validPath(params.Path) {
		return apierror.New(PutPolicyPathInvalidPathCode,
			"path must have prefix %s", common.GlobalLabelPrefix)
	}

	var node policy.Node
	if err := json.Unmarshal([]byte(*params.Policy), &node); err != nil {
		return NewPutPolicyPathInvalidPolicy()
	}

	if err := d.PolicyAdd(params.Path, &node); err != nil {
		return apierror.Error(PutPolicyPathFailureCode, err)
	}

	return NewPutPolicyPathOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
}

type getPolicy struct {
	daemon *Daemon
}

func NewGetPolicyHandler(d *Daemon) GetPolicyHandler {
	return &getPolicy{daemon: d}
}

// Returns the entire policy tree
func (h *getPolicy) Handle(params GetPolicyParams) middleware.Responder {
	d := h.daemon
	d.policy.Mutex.RLock()
	defer d.policy.Mutex.RUnlock()
	node := d.policy.Root
	return NewGetPolicyOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
}

type getPolicyPath struct {
	daemon *Daemon
}

func NewGetPolicyPathHandler(d *Daemon) GetPolicyPathHandler {
	return &getPolicyPath{daemon: d}
}

func (h *getPolicyPath) Handle(params GetPolicyPathParams) middleware.Responder {
	d := h.daemon
	d.policy.Mutex.RLock()
	defer d.policy.Mutex.RUnlock()

	if !validPath(params.Path) {
		return apierror.New(GetPolicyPathInvalidCode,
			"path must have prefix %s", common.GlobalLabelPrefix)
	}

	if node, _ := d.findNode(params.Path); node == nil {
		return NewGetPolicyPathNotFound()
	} else {
		return NewGetPolicyPathOK().WithPayload(models.PolicyTree(node.JSONMarshal()))
	}
}

func (d *Daemon) PolicyInit() error {
	for k, v := range policy.ReservedIdentities {
		key := policy.NumericIdentity(v).String()
		lbl := labels.NewLabel(
			key, "", common.ReservedLabelSource,
		)
		secLbl := policy.NewIdentity()
		secLbl.ID = v
		secLbl.AssociateEndpoint(lbl.String())
		secLbl.Labels[k] = lbl

		policyMapPath := bpf.MapPath(fmt.Sprintf("%sreserved_%d", policymap.MapName, int(v)))

		policyMap, _, err := policymap.OpenMap(policyMapPath)
		if err != nil {
			return fmt.Errorf("Could not create policy BPF map '%s': %s", policyMapPath, err)
		}

		if c := d.consumableCache.GetOrCreate(v, secLbl); c == nil {
			return fmt.Errorf("Unable to initialize consumable for %v", secLbl)
		} else {
			d.consumableCache.AddReserved(c)
			c.AddMap(policyMap)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package proxy

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/hellofresh/janus/pkg/router"
	log "github.com/sirupsen/logrus"
)

const (
	methodAll = "ALL"
)

// Register handles the register of proxies into the chosen router.
// It also handles the conversion from a proxy to an http.HandlerFunc
type Register struct {
	Router router.Router
	params Params
}

// NewRegister creates a new instance of Register
func NewRegister(router router.Router, params Params) *Register {
	return &Register{router, params}
}

// AddMany registers many proxies at once
func (p *Register) AddMany(routes []*Route) error {
	for _, r := range routes {
		err := p.Add(r)
		if nil != err {
			return err
		}
	}

	return nil
}

// Add register a new route
func (p *Register) Add(route *Route) error {
	definition := route.Proxy

	var balancer Balancer
	if definition.IsBalancerDefined() {
		log.WithField("balancing_alg", definition.Upstreams.Balancing).Debug("Using a load balancing algorithm")

		var err error
		balancer, err = NewBalancer(definition.Upstreams.Balancing)
		if err != nil {
			log.WithError(err).Error("Could not create a balancer")
			return err
		}
	}

	p.params.Outbound = route.Outbound
	p.params.InsecureSkipVerify = definition.InsecureSkipVerify
	handler := &httputil.ReverseProxy{
		Director:  p.createDirector(definition, balancer),
		Transport: NewTransportWithParams(p.params),
	}

	matcher := router.NewListenPathMatcher()
	if matcher.Match(definition.ListenPath) {
		p.doRegister(matcher.Extract(definition.ListenPath), handler.ServeHTTP, definition.Methods, route.Inbound)
	}

	p.doRegister(definition.ListenPath, handler.ServeHTTP, definition.Methods, route.Inbound)
	return nil
}

func (p *Register) createDirector(proxyDefinition *Definition, balancer Balancer) func(req *http.Request) {
	return func(req *http.Request) {
		var upstreamURL string

		// TODO: find better solution
		// maybe create "proxyDefinition.Upstreams.Targets every time",
		// but currently we have several points of definition creation
		if proxyDefinition.IsBalancerDefined() {
			upstream, err := balancer.Elect(proxyDefinition.Upstreams.Targets)
			if err != nil {
				log.WithError(err).Error("Could not elect one upstream")
				return
			}
			log.WithField("target", upstream.Target).Debug("Target upstream elected ")
			upstreamURL = upstream.Target
		} else {
			log.Warn("The upstream URL is deprecated. Use Upstreams instead")
			upstreamURL = proxyDefinition.UpstreamURL
		}

		target, err := url.Parse(upstreamURL)
		if err != nil {
			log.WithError(err).Error("Could not parse the target URL")
			return
		}

		targetQuery := target.RawQuery
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		path := target.Path

		if proxyDefinition.AppendPath {
			log.Debug("Appending listen path to the target url")
			path = singleJoiningSlash(target.Path, req.URL.Path)
		}

		if proxyDefinition.StripPath {
			path = singleJoiningSlash(target.Path, req.URL.Path)
			matcher := router.NewListenPathMatcher()
			listenPath := matcher.Extract(proxyDefinition.ListenPath)

			log.WithField("listen_path", listenPath).Debug("Stripping listen path")
			path = strings.Replace(path, listenPath, "", 1)
			if !strings.HasSuffix(target.Path, "/") && strings.HasSuffix(path, "/") {
				path = path[:len(path)-1]
			}
		}

		log.WithField("path", path).Debug("Upstream Path")
		req.URL.Path = path

		// This is very important to avoid problems with ssl verification for the HOST header
		if !proxyDefinition.PreserveHost {
			log.Debug("Preserving the host header")
			req.Host = target.Host
		}

		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
}

func (p *Register) doRegister(listenPath string, handler http.HandlerFunc, methods []string, handlers InChain) {
	log.WithFields(log.Fields{
		"listen_path": listenPath,
	}).Debug("Registering a route")

	if strings.Index(listenPath, "/") != 0 {
		log.WithField("listen_path", listenPath).
			Error("Route listen path must begin with '/'.Skipping invalid route.")
	} else {
		for _, method := range methods {
			if strings.ToUpper(method) == methodAll {
				p.Router.Any(listenPath, handler, handlers...)
			} else {
				p.Router.Handle(strings.ToUpper(method), listenPath, handler, handlers...)
			}
		}
	}
}

func cleanSlashes(a string) string {
	endSlash := strings.HasSuffix(a, "//")
	startSlash := strings.HasPrefix(a, "//")

	if startSlash {
		a = "/" + strings.TrimPrefix(a, "//")
	}

	if endSlash {
		a = strings.TrimSuffix(a, "//") + "/"
	}

	return a
}

func singleJoiningSlash(a, b string) string {
	a = cleanSlashes(a)
	b = cleanSlashes(b)

	aSlash := strings.HasSuffix(a, "/")
	bSlash := strings.HasPrefix(b, "/")

	switch {
	case aSlash && bSlash:
		return a + b[1:]
	case !aSlash && !bSlash:
		if len(b) > 0 {
			return a + "/" + b
		}
		return a
	}
	return a + b
}
<commit_msg>Apply CR changes<commit_after>package proxy

import (
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"github.com/hellofresh/janus/pkg/router"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

const (
	methodAll = "ALL"
)

// Register handles the register of proxies into the chosen router.
// It also handles the conversion from a proxy to an http.HandlerFunc
type Register struct {
	Router router.Router
	params Params
}

// NewRegister creates a new instance of Register
func NewRegister(router router.Router, params Params) *Register {
	return &Register{router, params}
}

// AddMany registers many proxies at once
func (p *Register) AddMany(routes []*Route) error {
	for _, r := range routes {
		err := p.Add(r)
		if nil != err {
			return err
		}
	}

	return nil
}

// Add register a new route
func (p *Register) Add(route *Route) error {
	definition := route.Proxy

	var balancer Balancer
	if definition.IsBalancerDefined() {
		log.WithField("balancing_alg", definition.Upstreams.Balancing).Debug("Using a load balancing algorithm")

		var err error
		balancer, err = NewBalancer(definition.Upstreams.Balancing)
		if err != nil {
			msg := "Could not create a balancer"
			log.WithError(err).Error(msg)
			return errors.Wrap(err, msg)
		}
	}

	p.params.Outbound = route.Outbound
	p.params.InsecureSkipVerify = definition.InsecureSkipVerify
	handler := &httputil.ReverseProxy{
		Director:  p.createDirector(definition, balancer),
		Transport: NewTransportWithParams(p.params),
	}

	matcher := router.NewListenPathMatcher()
	if matcher.Match(definition.ListenPath) {
		p.doRegister(matcher.Extract(definition.ListenPath), handler.ServeHTTP, definition.Methods, route.Inbound)
	}

	p.doRegister(definition.ListenPath, handler.ServeHTTP, definition.Methods, route.Inbound)
	return nil
}

func (p *Register) createDirector(proxyDefinition *Definition, balancer Balancer) func(req *http.Request) {
	return func(req *http.Request) {
		var upstreamURL string
		if proxyDefinition.IsBalancerDefined() && balancer != nil {
			upstream, err := balancer.Elect(proxyDefinition.Upstreams.Targets)
			if err != nil {
				log.WithError(err).Error("Could not elect one upstream")
				return
			}
			log.WithField("target", upstream.Target).Debug("Target upstream elected ")
			upstreamURL = upstream.Target
		} else {
			log.Warn("The upstream URL is deprecated. Use Upstreams instead")
			upstreamURL = proxyDefinition.UpstreamURL
		}

		target, err := url.Parse(upstreamURL)
		if err != nil {
			log.WithError(err).Error("Could not parse the target URL")
			return
		}

		targetQuery := target.RawQuery
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host
		path := target.Path

		if proxyDefinition.AppendPath {
			log.Debug("Appending listen path to the target url")
			path = singleJoiningSlash(target.Path, req.URL.Path)
		}

		if proxyDefinition.StripPath {
			path = singleJoiningSlash(target.Path, req.URL.Path)
			matcher := router.NewListenPathMatcher()
			listenPath := matcher.Extract(proxyDefinition.ListenPath)

			log.WithField("listen_path", listenPath).Debug("Stripping listen path")
			path = strings.Replace(path, listenPath, "", 1)
			if !strings.HasSuffix(target.Path, "/") && strings.HasSuffix(path, "/") {
				path = path[:len(path)-1]
			}
		}

		log.WithField("path", path).Debug("Upstream Path")
		req.URL.Path = path

		// This is very important to avoid problems with ssl verification for the HOST header
		if !proxyDefinition.PreserveHost {
			log.Debug("Preserving the host header")
			req.Host = target.Host
		}

		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}
}

func (p *Register) doRegister(listenPath string, handler http.HandlerFunc, methods []string, handlers InChain) {
	log.WithFields(log.Fields{
		"listen_path": listenPath,
	}).Debug("Registering a route")

	if strings.Index(listenPath, "/") != 0 {
		log.WithField("listen_path", listenPath).
			Error("Route listen path must begin with '/'.Skipping invalid route.")
	} else {
		for _, method := range methods {
			if strings.ToUpper(method) == methodAll {
				p.Router.Any(listenPath, handler, handlers...)
			} else {
				p.Router.Handle(strings.ToUpper(method), listenPath, handler, handlers...)
			}
		}
	}
}

func cleanSlashes(a string) string {
	endSlash := strings.HasSuffix(a, "//")
	startSlash := strings.HasPrefix(a, "//")

	if startSlash {
		a = "/" + strings.TrimPrefix(a, "//")
	}

	if endSlash {
		a = strings.TrimSuffix(a, "//") + "/"
	}

	return a
}

func singleJoiningSlash(a, b string) string {
	a = cleanSlashes(a)
	b = cleanSlashes(b)

	aSlash := strings.HasSuffix(a, "/")
	bSlash := strings.HasPrefix(b, "/")

	switch {
	case aSlash && bSlash:
		return a + b[1:]
	case !aSlash && !bSlash:
		if len(b) > 0 {
			return a + "/" + b
		}
		return a
	}
	return a + b
}
<|endoftext|>


<commit_before>package scheduler

import (
	"errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/hortonworks/gohadoop/hadoop_common/security"
	"github.com/hortonworks/gohadoop/hadoop_yarn"
	"github.com/hortonworks/gohadoop/hadoop_yarn/conf"
	"github.com/hortonworks/gohadoop/hadoop_yarn/yarn_client"
	"log"
	"os"
	"time"
)

func YARNInit() (*yarn_client.YarnClient, *yarn_client.AMRMClient) {
	var err error

	//Hack! This should be external, but doing this here for demo purposes
	hadoopHome := "/home/vagrant/hadoop/install/hadoop-2.6.0-SNAPSHOT"
	os.Setenv("HADOOP_HOME", hadoopHome)
	os.Setenv("HADOOP_COMMON_HOME", hadoopHome)
	os.Setenv("HADOOP_CONF_DIR", hadoopHome+"/etc/hadoop")
	os.Setenv("HADOOP_HDFS_HOME", hadoopHome)
	os.Setenv("HADOOP_MAPRED_HOME", hadoopHome)

	// Create YarnConfiguration
	conf, _ := conf.NewYarnConfiguration()

	// Create YarnClient
	yarnClient, _ := yarn_client.CreateYarnClient(conf)

	// Create new application to get ApplicationSubmissionContext
	_, asc, _ := yarnClient.CreateNewApplication()

	// Some useful information
	queue := "default"
	appName := "kubernetes"
	appType := "PAAS"
	unmanaged := true
	clc := hadoop_yarn.ContainerLaunchContextProto{}

	// Setup ApplicationSubmissionContext for the application
	asc.AmContainerSpec = &clc
	asc.ApplicationName = &appName
	asc.Queue = &queue
	asc.ApplicationType = &appType
	asc.UnmanagedAm = &unmanaged

	// Submit!
	err = yarnClient.SubmitApplication(asc)
	if err != nil {
		log.Fatal("yarnClient.SubmitApplication ", err)
	}
	log.Println("Successfully submitted unmanaged application: ", asc.ApplicationId)
	time.Sleep(1 * time.Second)

	appReport, err := yarnClient.GetApplicationReport(asc.ApplicationId)
	if err != nil {
		log.Fatal("yarnClient.GetApplicationReport ", err)
	}
	appState := appReport.GetYarnApplicationState()
	for appState != hadoop_yarn.YarnApplicationStateProto_ACCEPTED {
		log.Println("Application in state ", appState)
		time.Sleep(1 * time.Second)
		appReport, err = yarnClient.GetApplicationReport(asc.ApplicationId)
		appState = appReport.GetYarnApplicationState()
		if appState == hadoop_yarn.YarnApplicationStateProto_FAILED || appState == hadoop_yarn.YarnApplicationStateProto_KILLED {
			log.Fatal("Application in state ", appState)
		}
	}

	amRmToken := appReport.GetAmRmToken()

	if amRmToken != nil {
		savedAmRmToken := *amRmToken
		service, _ := conf.GetRMSchedulerAddress()
		savedAmRmToken.Service = &service
		security.GetCurrentUser().AddUserToken(&savedAmRmToken)
	}

	log.Println("Application in state ", appState)

	// Create AMRMClient
	var attemptId int32
	attemptId = 1
	applicationAttemptId := hadoop_yarn.ApplicationAttemptIdProto{ApplicationId: asc.ApplicationId, AttemptId: &attemptId}

	rmClient, _ := yarn_client.CreateAMRMClient(conf, &applicationAttemptId)
	log.Println("Created RM client: ", rmClient)

	// Wait for ApplicationAttempt to be in Launched state
	appAttemptReport, err := yarnClient.GetApplicationAttemptReport(&applicationAttemptId)
	appAttemptState := appAttemptReport.GetYarnApplicationAttemptState()
	for appAttemptState != hadoop_yarn.YarnApplicationAttemptStateProto_APP_ATTEMPT_LAUNCHED {
		log.Println("ApplicationAttempt in state ", appAttemptState)
		time.Sleep(1 * time.Second)
		appAttemptReport, err = yarnClient.GetApplicationAttemptReport(&applicationAttemptId)
		appAttemptState = appAttemptReport.GetYarnApplicationAttemptState()
	}
	log.Println("ApplicationAttempt in state ", appAttemptState)

	// Register with ResourceManager
	log.Println("About to register application master.")
	err = rmClient.RegisterApplicationMaster("", -1, "")
	if err != nil {
		log.Fatal("rmClient.RegisterApplicationMaster ", err)
	}
	log.Println("Successfully registered application master.")

	return yarnClient, rmClient
}

type YARNScheduler struct {
	yarnClient          *yarn_client.YarnClient
	rmClient            *yarn_client.AMRMClient
	podsToContainersMap map[string]*hadoop_yarn.ContainerIdProto
}

func NewYARNScheduler() Scheduler {
	yarnC, rmC := YARNInit()
	podsToContainers := make(map[string]*hadoop_yarn.ContainerIdProto)

	return &YARNScheduler{
		yarnClient:          yarnC,
		rmClient:            rmC,
		podsToContainersMap: podsToContainers}
}

func (yarnScheduler *YARNScheduler) Delete(id string) error {
	log.Println("yarn delete hook")

	return nil
}

func (yarnScheduler *YARNScheduler) Schedule(pod api.Pod, minionLister MinionLister) (string, error) {

	rmClient := yarnScheduler.rmClient

	// Add resource requests
	const numContainers = int32(1)
	memory := int32(128)
	resource := hadoop_yarn.ResourceProto{Memory: &memory}
	rmClient.AddRequest(1, "*", &resource, numContainers)

	// Now call ResourceManager.allocate
	allocateResponse, err := rmClient.Allocate()
	if err == nil {
		log.Println("allocateResponse: ", *allocateResponse)
	}
	log.Println("#containers allocated: ", len(allocateResponse.AllocatedContainers))

	numAllocatedContainers := int32(0)
	allocatedContainers := make([]*hadoop_yarn.ContainerProto, numContainers, numContainers)
	for numAllocatedContainers < numContainers {
		// Sleep for a while
		log.Println("Sleeping...")
		time.Sleep(3 * time.Second)
		log.Println("Sleeping... done!")

		// Try to get containers now...
		allocateResponse, err = rmClient.Allocate()
		if err == nil {
			log.Println("allocateResponse: ", *allocateResponse)
		}

		for _, container := range allocateResponse.AllocatedContainers {
			allocatedContainers[numAllocatedContainers] = container
			numAllocatedContainers++
			log.Println("#containers allocated so far: ", numAllocatedContainers)

			//We have the hostname available. return from here.
			yarnScheduler.podsToContainersMap[pod.ID] = container.GetId()
			return *container.NodeId.Host, nil
		}

		log.Println("#containers allocated: ", len(allocateResponse.AllocatedContainers))
		log.Println("Total #containers allocated so far: ", numAllocatedContainers)
	}
	log.Println("Final #containers allocated: ", numAllocatedContainers)

	return "<invalid_host>", errors.New("invalid_host")
}
<commit_msg>cleaning up yarn.go a little. adding support for YARN node/kubernetes minion mapping<commit_after>package scheduler

import (
	"errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/hortonworks/gohadoop/hadoop_common/security"
	"github.com/hortonworks/gohadoop/hadoop_yarn"
	"github.com/hortonworks/gohadoop/hadoop_yarn/conf"
	"github.com/hortonworks/gohadoop/hadoop_yarn/yarn_client"

	"log"
	"net"
	"os"
	"time"
)

func YARNInit() (*yarn_client.YarnClient, *yarn_client.AMRMClient) {
	var err error

	//Hack! This should be external, but doing this here for demo purposes
	hadoopHome := "/home/vagrant/hadoop/install/hadoop-2.6.0-SNAPSHOT"
	os.Setenv("HADOOP_HOME", hadoopHome)
	os.Setenv("HADOOP_COMMON_HOME", hadoopHome)
	os.Setenv("HADOOP_CONF_DIR", hadoopHome+"/etc/hadoop")
	os.Setenv("HADOOP_HDFS_HOME", hadoopHome)
	os.Setenv("HADOOP_MAPRED_HOME", hadoopHome)

	// Create YarnConfiguration
	conf, _ := conf.NewYarnConfiguration()

	// Create YarnClient
	yarnClient, _ := yarn_client.CreateYarnClient(conf)

	// Create new application to get ApplicationSubmissionContext
	_, asc, _ := yarnClient.CreateNewApplication()

	// Some useful information
	queue := "default"
	appName := "kubernetes"
	appType := "PAAS"
	unmanaged := true
	clc := hadoop_yarn.ContainerLaunchContextProto{}

	// Setup ApplicationSubmissionContext for the application
	asc.AmContainerSpec = &clc
	asc.ApplicationName = &appName
	asc.Queue = &queue
	asc.ApplicationType = &appType
	asc.UnmanagedAm = &unmanaged

	// Submit!
	err = yarnClient.SubmitApplication(asc)
	if err != nil {
		log.Fatal("yarnClient.SubmitApplication ", err)
	}
	log.Println("Successfully submitted unmanaged application: ", asc.ApplicationId)
	time.Sleep(1 * time.Second)

	appReport, err := yarnClient.GetApplicationReport(asc.ApplicationId)
	if err != nil {
		log.Fatal("yarnClient.GetApplicationReport ", err)
	}
	appState := appReport.GetYarnApplicationState()
	for appState != hadoop_yarn.YarnApplicationStateProto_ACCEPTED {
		log.Println("Application in state ", appState)
		time.Sleep(1 * time.Second)
		appReport, err = yarnClient.GetApplicationReport(asc.ApplicationId)
		appState = appReport.GetYarnApplicationState()
		if appState == hadoop_yarn.YarnApplicationStateProto_FAILED || appState == hadoop_yarn.YarnApplicationStateProto_KILLED {
			log.Fatal("Application in state ", appState)
		}
	}

	amRmToken := appReport.GetAmRmToken()

	if amRmToken != nil {
		savedAmRmToken := *amRmToken
		service, _ := conf.GetRMSchedulerAddress()
		savedAmRmToken.Service = &service
		security.GetCurrentUser().AddUserToken(&savedAmRmToken)
	}

	log.Println("Application in state ", appState)

	// Create AMRMClient
	var attemptId int32
	attemptId = 1
	applicationAttemptId := hadoop_yarn.ApplicationAttemptIdProto{ApplicationId: asc.ApplicationId, AttemptId: &attemptId}

	rmClient, _ := yarn_client.CreateAMRMClient(conf, &applicationAttemptId)
	log.Println("Created RM client: ", rmClient)

	// Wait for ApplicationAttempt to be in Launched state
	appAttemptReport, err := yarnClient.GetApplicationAttemptReport(&applicationAttemptId)
	appAttemptState := appAttemptReport.GetYarnApplicationAttemptState()
	for appAttemptState != hadoop_yarn.YarnApplicationAttemptStateProto_APP_ATTEMPT_LAUNCHED {
		log.Println("ApplicationAttempt in state ", appAttemptState)
		time.Sleep(1 * time.Second)
		appAttemptReport, err = yarnClient.GetApplicationAttemptReport(&applicationAttemptId)
		appAttemptState = appAttemptReport.GetYarnApplicationAttemptState()
	}
	log.Println("ApplicationAttempt in state ", appAttemptState)

	// Register with ResourceManager
	log.Println("About to register application master.")
	err = rmClient.RegisterApplicationMaster("", -1, "")
	if err != nil {
		log.Fatal("rmClient.RegisterApplicationMaster ", err)
	}
	log.Println("Successfully registered application master.")

	return yarnClient, rmClient
}

type YARNScheduler struct {
	yarnClient          *yarn_client.YarnClient
	rmClient            *yarn_client.AMRMClient
	podsToContainersMap map[string]*hadoop_yarn.ContainerIdProto
}

func NewYARNScheduler() Scheduler {
	yarnC, rmC := YARNInit()
	podsToContainers := make(map[string]*hadoop_yarn.ContainerIdProto)

	return &YARNScheduler{
		yarnClient:          yarnC,
		rmClient:            rmC,
		podsToContainersMap: podsToContainers}
}

func (yarnScheduler *YARNScheduler) Delete(id string) error {
	log.Println("yarn delete hook")

	return nil
}

func (yarnScheduler *YARNScheduler) Schedule(pod api.Pod, minionLister MinionLister) (string, error) {
	rmClient := yarnScheduler.rmClient

	// Add resource requests
	const numContainers = int32(1)
	memory := int32(128)
	resource := hadoop_yarn.ResourceProto{Memory: &memory}
	numAllocatedContainers := int32(0)
	const maxAllocationAttempts = int(5)
	allocationAttempts := 0
	allocatedContainers := make([]*hadoop_yarn.ContainerProto, numContainers, numContainers)

	rmClient.AddRequest(1, "*", &resource, numContainers)

	for numAllocatedContainers < numContainers && allocationAttempts < maxAllocationAttempts {
		// Sleep for a while
		log.Println("Sleeping...")
		time.Sleep(3 * time.Second)
		log.Println("Sleeping... done!")

		// Try to get containers now...
		allocateResponse, err := rmClient.Allocate()
		if err == nil {
			log.Println("allocateResponse: ", *allocateResponse)
		}

		for _, container := range allocateResponse.AllocatedContainers {
			allocatedContainers[numAllocatedContainers] = container
			numAllocatedContainers++
			log.Println("#containers allocated so far: ", numAllocatedContainers)

			//We have the hostname available. return from here.
			yarnScheduler.podsToContainersMap[pod.ID] = container.GetId()
			host := *container.NodeId.Host

			log.Println("allocated container on: ", host)

			return findMinionForHost(host, minionLister)
		}

		allocationAttempts++
		log.Println("#containers allocated: ", len(allocateResponse.AllocatedContainers))
		log.Println("Total #containers allocated so far: ", numAllocatedContainers)
	}

	log.Println("Final #containers allocated: ", numAllocatedContainers)

	return "<invalid_host>", errors.New("unable to schedule pod! YARN didn't allocate a container")
}

/* YARN returns hostnames, but minions maybe using IPs.
TODO: This is an expensive mechanism to find the right minion corresponding to the YARN node.
      Find a better mechanism if possible
*/
func findMinionForHost(host string, minionLister MinionLister) (string, error) {
	hostIPs, err := net.LookupIP(host)

	if err != nil {
		return "<invalid_host>", errors.New("unable to lookup IPs for YARN host: " + host)
	}

	for _, hostIP := range hostIPs {
		minions, err := minionLister.List()
		if err != nil {
			return "<invalid_host>", errors.New("update to list minions")
		}

		for _, minion := range minions {
			minionIPs, err := net.LookupIP(minion)

			if err != nil {
				return "<invalid_host>", errors.New("unable to lookup IPs for minion: " + minion)
			}

			for _, minionIP := range minionIPs {
				if hostIP.Equal(minionIP) {
					log.Printf("YARN host %s maps to minion: %s", host, minion)
					return minion, nil
				}
			}
		}
	}

	return "<invalid_host>", errors.New("unable to find minion for YARN host: " + host)
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Etix Labs - All Rights Reserved.
// All information contained herein is, and remains the property of Etix Labs and its suppliers,
// if any. The intellectual and technical concepts contained herein are proprietary to Etix Labs
// Dissemination of this information or reproduction of this material is strictly forbidden unless
// prior written permission is obtained from Etix Labs.

package jsh

import (
	"reflect"
	"strings"
	"unicode"
)

const (
	tagNameJSON    = "json"
	tagIgnore      = "-"
	tagNameJSH     = "jsh"
	tagSep         = ","
	tagToOne       = "one"
	tagToMany      = "many"
	tagCreate      = "create"
	tagUpdate      = "update"
	optionSep      = "/"
	optionRequired = "required"
	fieldSep       = "/"
)

// tagOptions represents the options that can be passed to JSH tags.
type tagOptions struct {
	required bool
}

// tags represents the tag options by tag name of a struct field
type tags map[string]*tagOptions

// decodeJSONTag returns the first JSON tag of the given struct field. If there is none, the field name is returned.
func decodeJSONTag(f reflect.StructField) string {
	rawTags := f.Tag.Get(tagNameJSON)
	tags := strings.SplitN(rawTags, tagSep, -1)
	if len(tags) == 0 {
		return f.Name
	}
	return tags[0]
}

// decodeFieldTags decodes all JSH tags from the struct field to a tag struct.
func decodeFieldTags(rawTags string) tags {
	var result = make(tags)
	options := strings.SplitN(rawTags, tagSep, -1)
	for _, option := range options {
		jshTag := strings.SplitN(option, optionSep, -1)
		if !isValidTag(jshTag[0]) {
			continue
		}
		options := &tagOptions{}
		if len(jshTag) == 2 {
			options.required = jshTag[1] == optionRequired
		}
		result[jshTag[0]] = options
	}
	return result
}

// decodeFieldTag decodes the JSH tag from the struct field to a tag struct.
// It returns nil if the tag was not found.
func decodeFieldTag(tags, tagName string) *tagOptions {
	result := decodeFieldTags(tags)
	return result[tagName]
}

// isValidTag returns false if the tag is empty or contains invalid characters.
func isValidTag(s string) bool {
	if s == "" {
		return false
	}
	for _, c := range s {
		switch {
		case strings.ContainsRune("!#$%&()*+-./:<=>?@[]^_{|}~ ", c):
			// Backslash and quote chars are reserved, but
			// otherwise any punctuation chars are allowed
			// in a tag name.
		default:
			if !unicode.IsLetter(c) && !unicode.IsDigit(c) {
				return false
			}
		}
	}
	return true
}
<commit_msg>#10 Remove EtixLabs copyright<commit_after>package jsh

import (
	"reflect"
	"strings"
	"unicode"
)

const (
	tagNameJSON    = "json"
	tagIgnore      = "-"
	tagNameJSH     = "jsh"
	tagSep         = ","
	tagToOne       = "one"
	tagToMany      = "many"
	tagCreate      = "create"
	tagUpdate      = "update"
	optionSep      = "/"
	optionRequired = "required"
	fieldSep       = "/"
)

// tagOptions represents the options that can be passed to JSH tags.
type tagOptions struct {
	required bool
}

// tags represents the tag options by tag name of a struct field
type tags map[string]*tagOptions

// decodeJSONTag returns the first JSON tag of the given struct field. If there is none, the field name is returned.
func decodeJSONTag(f reflect.StructField) string {
	rawTags := f.Tag.Get(tagNameJSON)
	tags := strings.SplitN(rawTags, tagSep, -1)
	if len(tags) == 0 {
		return f.Name
	}
	return tags[0]
}

// decodeFieldTags decodes all JSH tags from the struct field to a tag struct.
func decodeFieldTags(rawTags string) tags {
	var result = make(tags)
	options := strings.SplitN(rawTags, tagSep, -1)
	for _, option := range options {
		jshTag := strings.SplitN(option, optionSep, -1)
		if !isValidTag(jshTag[0]) {
			continue
		}
		options := &tagOptions{}
		if len(jshTag) == 2 {
			options.required = jshTag[1] == optionRequired
		}
		result[jshTag[0]] = options
	}
	return result
}

// decodeFieldTag decodes the JSH tag from the struct field to a tag struct.
// It returns nil if the tag was not found.
func decodeFieldTag(tags, tagName string) *tagOptions {
	result := decodeFieldTags(tags)
	return result[tagName]
}

// isValidTag returns false if the tag is empty or contains invalid characters.
func isValidTag(s string) bool {
	if s == "" {
		return false
	}
	for _, c := range s {
		switch {
		case strings.ContainsRune("!#$%&()*+-./:<=>?@[]^_{|}~ ", c):
			// Backslash and quote chars are reserved, but
			// otherwise any punctuation chars are allowed
			// in a tag name.
		default:
			if !unicode.IsLetter(c) && !unicode.IsDigit(c) {
				return false
			}
		}
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
)

var conn dbox.IConnection
var count int
var ratioTableName string

var (
	sourcetablename = "salespls-summary"
	calctablename   = "salespls-summary"
	desttablename   = "salespls-summary"
	t0              time.Time
	masters         = toolkit.M{}
	sgaalloc        = map[string]float64{
		"EXP": 0.08,
		"I4":  0.08,
		"I6":  0.105,
	}
)

type plalloc struct {
	ID               string `bson:"_id" json:"_id"`
	Key              string
	Key1, Key2, Key3 string
	Ref1             float64
	Current          float64
	Expect           float64
	Absorbed         float64
}

type allocmap map[string]*plalloc

var (
	plallocs = allocmap{}
	totals1  = allocmap{}
	totals2  = allocmap{}
)

func main() {
	setinitialconnection()
	prepmastercalc()
	buildratio()
	processTable()
}

func buildratio() {
	cursor, _ := conn.NewQuery().From(calctablename).Select().Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	mstone := 0
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		e := cursor.Fetch(&mr, 1, false)
		if e != nil {
			break
		}
		i++
		makeProgressLog("Build ratio", i, count, 5, &mstone, t0)

		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		fiscal := key.GetString("date_fiscal")
		channelid := key.GetString("report_channelid")
		sales := mr.GetFloat64("PL8A")
		keytotal := toolkit.Sprintf("%s_%s", fiscal, channelid)
		if channelid == "EXP" || channelid == "I4" || channelid == "I6" {
			sgaratio := sgaalloc[channelid]
			sgavalue := sgaratio * sales
			adjustAllocs(&totals1, keytotal, 0, -sgavalue, 0, sales)
			adjustAllocs(&totals2, fiscal, 0, sgavalue, 0, 0)
		} else {
			adjustAllocs(&totals2, fiscal, 0, 0, 0, sales)
		}
	}
}

func processTable() {
	connsave, _ := modules.GetDboxIConnection("db_godrej")
	defer connsave.Close()

	qsave := conn.NewQuery().From(desttablename).Save()
	cursor, _ := conn.NewQuery().From(calctablename).Select().Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	mstone := 0
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		e := cursor.Fetch(&mr, 1, false)
		if e != nil {
			break
		}
		i++
		makeProgressLog("Processing", i, count, 5, &mstone, t0)

		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		fiscal := key.GetString("date_fiscal")
		channelid := key.GetString("report_channelid")
		sales := mr.GetFloat64("PL8A")
		//keytotal := toolkit.Sprintf("%s_%s", fiscal, channelid)
		//total1 := totals1[keytotal]
		total2 := totals2[fiscal]
		if channelid == "EXP" || channelid == "I4" || channelid == "I6" {
			sgaratio := sgaalloc[channelid]
			value := -sgaratio * sales
			mr.Set("PL34_Other", value)
		} else {
			value := sales * total2.Expect / total2.Ref1
			mr.Set("PL34_Other", value)
		}

		gdrj.CalcSum(mr, masters)
		esave := qsave.Exec(toolkit.M{}.Set("data", mr))
		if esave != nil {
			toolkit.Printfn("Error: %s", esave.Error())
			return
		}
	}
}

func adjustAllocs(allocsmap *allocmap, key string, current, expect, absorbed, ref1 float64) {
	allocs := *allocsmap
	alloc := allocs[key]
	if alloc == nil {
		alloc = new(plalloc)
		alloc.Key = key
		alloc.ID = key
	}
	alloc.Current += current
	alloc.Expect += expect
	alloc.Ref1 += ref1
	alloc.Absorbed += absorbed
	allocs[key] = alloc
	*allocsmap = allocs
}

func makeProgressLog(reference string, i, count, step int, current *int, tstart time.Time) int {
	perstep := count * step / 100
	icurrent := *current
	if icurrent == 0 {
		icurrent = perstep
	}
	pct := i * 100 / count
	if i >= icurrent {
		toolkit.Printfn("Processing %s, %d of %d [%d pct] in %s",
			reference, i, count, pct, time.Since(tstart).String())
		icurrent += perstep
	}
	*current = icurrent
	return icurrent
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, ecrx := gdrj.Find(fnModel(), filter, nil)
	if ecrx != nil {
		toolkit.Printfn("Cursor Error: %s", ecrx.Error())
		os.Exit(100)
	}
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmastercalc() {
	toolkit.Println("--> PL MODEL")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}
<commit_msg>updsga<commit_after>package main

import (
	"eaciit/gdrj/model"
	"eaciit/gdrj/modules"
	"os"

	"time"

	"github.com/eaciit/dbox"
	"github.com/eaciit/orm/v1"
	"github.com/eaciit/toolkit"
)

var conn dbox.IConnection
var count int
var ratioTableName string

var (
	sourcetablename = "salespls-summary"
	calctablename   = "salespls-summary"
	desttablename   = "salespls-summary"
	t0              time.Time
	masters         = toolkit.M{}
	sgaalloc        = map[string]float64{
		"EXP": 0.08,
		"I4":  0.08,
		"I6":  0.105,
	}
)

type plalloc struct {
	ID               string `bson:"_id" json:"_id"`
	Key              string
	Key1, Key2, Key3 string
	Ref1             float64
	Current          float64
	Expect           float64
	Absorbed         float64
}

type allocmap map[string]*plalloc

var (
	plallocs = allocmap{}
	totals1  = allocmap{}
	totals2  = allocmap{}
)

func main() {
	setinitialconnection()
	prepmastercalc()
	buildratio()
	processTable()
}

func buildratio() {
	cursor, _ := conn.NewQuery().From(calctablename).Select().Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	mstone := 0
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		e := cursor.Fetch(&mr, 1, false)
		if e != nil {
			break
		}
		i++
		makeProgressLog("Build ratio", i, count, 5, &mstone, t0)

		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		fiscal := key.GetString("date_fiscal")
		channelid := key.GetString("report_channelid")
		sales := mr.GetFloat64("PL8A")
		keytotal := toolkit.Sprintf("%s_%s", fiscal, channelid)
		if channelid == "EXP" || channelid == "I4" || channelid == "I6" {
			sgaratio := sgaalloc[channelid]
			sgavalue := sgaratio * sales
			adjustAllocs(&totals1, keytotal, 0, -sgavalue, 0, sales)
			adjustAllocs(&totals2, fiscal, 0, sgavalue, 0, 0)
		} else {
			adjustAllocs(&totals2, fiscal, 0, 0, 0, sales)
		}
	}
}

func processTable() {
	connsave, _ := modules.GetDboxIConnection("db_godrej")
	defer connsave.Close()

	qsave := conn.NewQuery().From(desttablename).Save()
	cursor, _ := conn.NewQuery().From(calctablename).Select().Cursor(nil)
	defer cursor.Close()

	i := 0
	count := cursor.Count()
	mstone := 0
	t0 = time.Now()
	for {
		mr := toolkit.M{}
		e := cursor.Fetch(&mr, 1, false)
		if e != nil {
			break
		}
		i++
		makeProgressLog("Processing", i, count, 5, &mstone, t0)

		key := mr.Get("key", toolkit.M{}).(toolkit.M)
		fiscal := key.GetString("date_fiscal")
		channelid := key.GetString("report_channelid")
		sales := mr.GetFloat64("PL8A")
		//keytotal := toolkit.Sprintf("%s_%s", fiscal, channelid)
		//total1 := totals1[keytotal]
		total2 := totals2[fiscal]
		if channelid == "EXP" || channelid == "I4" || channelid == "I6" {
			sgaratio := sgaalloc[channelid]
			value := -sgaratio * sales
			mr.Set("PL34_Other", value)
		} else {
			value := mr.GetFloat64("PL34_Other")
			value += sales * total2.Expect / total2.Ref1
			mr.Set("PL34_Other", value)
		}

		gdrj.CalcSum(mr, masters)
		esave := qsave.Exec(toolkit.M{}.Set("data", mr))
		if esave != nil {
			toolkit.Printfn("Error: %s", esave.Error())
			return
		}
	}
}

func adjustAllocs(allocsmap *allocmap, key string, current, expect, absorbed, ref1 float64) {
	allocs := *allocsmap
	alloc := allocs[key]
	if alloc == nil {
		alloc = new(plalloc)
		alloc.Key = key
		alloc.ID = key
	}
	alloc.Current += current
	alloc.Expect += expect
	alloc.Ref1 += ref1
	alloc.Absorbed += absorbed
	allocs[key] = alloc
	*allocsmap = allocs
}

func makeProgressLog(reference string, i, count, step int, current *int, tstart time.Time) int {
	perstep := count * step / 100
	icurrent := *current
	if icurrent == 0 {
		icurrent = perstep
	}
	pct := i * 100 / count
	if i >= icurrent {
		toolkit.Printfn("Processing %s, %d of %d [%d pct] in %s",
			reference, i, count, pct, time.Since(tstart).String())
		icurrent += perstep
	}
	*current = icurrent
	return icurrent
}

func buildmap(holder interface{},
	fnModel func() orm.IModel,
	filter *dbox.Filter,
	fnIter func(holder interface{}, obj interface{})) interface{} {
	crx, ecrx := gdrj.Find(fnModel(), filter, nil)
	if ecrx != nil {
		toolkit.Printfn("Cursor Error: %s", ecrx.Error())
		os.Exit(100)
	}
	defer crx.Close()
	for {
		s := fnModel()
		e := crx.Fetch(s, 1, false)
		if e != nil {
			break
		}
		fnIter(holder, s)
	}
	return holder
}

func prepmastercalc() {
	toolkit.Println("--> PL MODEL")
	masters.Set("plmodel", buildmap(map[string]*gdrj.PLModel{},
		func() orm.IModel {
			return new(gdrj.PLModel)
		},
		nil,
		func(holder, obj interface{}) {
			h := holder.(map[string]*gdrj.PLModel)
			o := obj.(*gdrj.PLModel)
			h[o.ID] = o
		}).(map[string]*gdrj.PLModel))
}

func setinitialconnection() {
	var err error
	conn, err = modules.GetDboxIConnection("db_godrej")

	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}

	err = gdrj.SetDb(conn)
	if err != nil {
		toolkit.Println("Initial connection found : ", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package tango

import (
	"net/http"
	"os"
)

const (
	Dev = iota
	Prod
)

var (
	Env = Dev

	modes = []string{
		"Dev",
		"Product",
	}
)

func Version() string {
	return "0.2.7.0119"
}

type Tango struct {
	Router
	Mode       int
	handlers   []Handler
	logger     Logger
	ErrHandler Handler
}

var (
	ClassicHandlers = []Handler{
		Logging(),
		Recovery(true),
		Compresses([]string{".js", ".css", ".html", ".htm"}),
		Static(StaticOptions{Prefix: "public"}),
		Return(),
		Responses(),
		Requests(),
		Param(),
		Contexts(),
	}
)

func (t *Tango) Logger() Logger {
	return t.logger
}

func (t *Tango) Get(url string, c interface{}) {
	t.Route([]string{"GET", "HEAD"}, url, c)
}

func (t *Tango) Post(url string, c interface{}) {
	t.Route([]string{"POST"}, url, c)
}

func (t *Tango) Head(url string, c interface{}) {
	t.Route([]string{"HEAD"}, url, c)
}

func (t *Tango) Options(url string, c interface{}) {
	t.Route([]string{"OPTIONS"}, url, c)
}

func (t *Tango) Trace(url string, c interface{}) {
	t.Route([]string{"TRACE"}, url, c)
}

func (t *Tango) Patch(url string, c interface{}) {
	t.Route([]string{"PATCH"}, url, c)
}

func (t *Tango) Delete(url string, c interface{}) {
	t.Route([]string{"DELETE"}, url, c)
}

func (t *Tango) Put(url string, c interface{}) {
	t.Route([]string{"PUT"}, url, c)
}

func (t *Tango) Any(url string, c interface{}) {
	t.Route(SupportMethods, url, c)
}

func (t *Tango) Use(handlers ...Handler) {
	for _, handler := range handlers {
		t.handlers = append(t.handlers, handler)
	}
}

func (t *Tango) Run(addrs ...string) {
	var addr string
	if len(addrs) == 0 {
		addr = ":8000"
	} else {
		addr = addrs[0]
	}

	t.logger.Info("listening on", addr, modes[t.Mode])

	err := http.ListenAndServe(addr, t)
	if err != nil {
		t.logger.Error(err)
	}
}

func (t *Tango) RunTLS(certFile, keyFile string, addrs ...string) {
	var addr string
	if len(addrs) == 0 {
		addr = ":8000"
	} else {
		addr = addrs[0]
	}

	t.logger.Info("listening on https", addr, modes[t.Mode])

	err := http.ListenAndServeTLS(addr, certFile, keyFile, t)
	if err != nil {
		t.logger.Error(err)
	}
}

type HandlerFunc func(ctx *Context)

func (h HandlerFunc) Handle(ctx *Context) {
	h(ctx)
}

func WrapBefore(handler http.Handler) HandlerFunc {
	return func(ctx *Context) {
		handler.ServeHTTP(ctx.ResponseWriter, ctx.Req())

		ctx.Next()
	}
}

func WrapAfter(handler http.Handler) HandlerFunc {
	return func(ctx *Context) {
		ctx.Next()

		handler.ServeHTTP(ctx.ResponseWriter, ctx.Req())
	}
}

func (t *Tango) UseHandler(handler http.Handler) {
	t.Use(WrapBefore(handler))
}

func (t *Tango) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ctx := NewContext(
		t,
		req,
		NewResponseWriter(w),
		t.logger,
	)

	ctx.Invoke()

	// if there is no logging or error handle, so the last written check.
	if !ctx.Written() {
		if ctx.Result == nil {
			ctx.Result = NotFound()
		}
		ctx.HandleError()
		p := req.URL.Path
		if len(req.URL.RawQuery) > 0 {
			p = p + "?" + req.URL.RawQuery
		}

		t.logger.Error(req.Method, ctx.Status(), p)
	}
}

func NewWithLog(logger Logger, handlers ...Handler) *Tango {
	tango := &Tango{
		Router:     NewRouter(),
		Mode:       Env,
		logger:     logger,
		handlers:   make([]Handler, 0),
		ErrHandler: Errors(),
	}

	tango.Use(handlers...)

	return tango
}

func New(handlers ...Handler) *Tango {
	return NewWithLog(NewLogger(os.Stdout), handlers...)
}

func Classic(l ...Logger) *Tango {
	var logger Logger
	if len(l) == 0 {
		logger = NewLogger(os.Stdout)
	} else {
		logger = l[0]
	}

	return NewWithLog(
		logger,
		ClassicHandlers...,
	)
}
<commit_msg>directly append slice<commit_after>package tango

import (
	"net/http"
	"os"
)

const (
	Dev = iota
	Prod
)

var (
	Env = Dev

	modes = []string{
		"Dev",
		"Product",
	}
)

func Version() string {
	return "0.2.7.0119"
}

type Tango struct {
	Router
	Mode       int
	handlers   []Handler
	logger     Logger
	ErrHandler Handler
}

var (
	ClassicHandlers = []Handler{
		Logging(),
		Recovery(true),
		Compresses([]string{".js", ".css", ".html", ".htm"}),
		Static(StaticOptions{Prefix: "public"}),
		Return(),
		Responses(),
		Requests(),
		Param(),
		Contexts(),
	}
)

func (t *Tango) Logger() Logger {
	return t.logger
}

func (t *Tango) Get(url string, c interface{}) {
	t.Route([]string{"GET", "HEAD"}, url, c)
}

func (t *Tango) Post(url string, c interface{}) {
	t.Route([]string{"POST"}, url, c)
}

func (t *Tango) Head(url string, c interface{}) {
	t.Route([]string{"HEAD"}, url, c)
}

func (t *Tango) Options(url string, c interface{}) {
	t.Route([]string{"OPTIONS"}, url, c)
}

func (t *Tango) Trace(url string, c interface{}) {
	t.Route([]string{"TRACE"}, url, c)
}

func (t *Tango) Patch(url string, c interface{}) {
	t.Route([]string{"PATCH"}, url, c)
}

func (t *Tango) Delete(url string, c interface{}) {
	t.Route([]string{"DELETE"}, url, c)
}

func (t *Tango) Put(url string, c interface{}) {
	t.Route([]string{"PUT"}, url, c)
}

func (t *Tango) Any(url string, c interface{}) {
	t.Route(SupportMethods, url, c)
}

func (t *Tango) Use(handlers ...Handler) {
	t.handlers = append(t.handlers, handlers...)
}

func (t *Tango) Run(addrs ...string) {
	var addr string
	if len(addrs) == 0 {
		addr = ":8000"
	} else {
		addr = addrs[0]
	}

	t.logger.Info("listening on", addr, modes[t.Mode])

	err := http.ListenAndServe(addr, t)
	if err != nil {
		t.logger.Error(err)
	}
}

func (t *Tango) RunTLS(certFile, keyFile string, addrs ...string) {
	var addr string
	if len(addrs) == 0 {
		addr = ":8000"
	} else {
		addr = addrs[0]
	}

	t.logger.Info("listening on https", addr, modes[t.Mode])

	err := http.ListenAndServeTLS(addr, certFile, keyFile, t)
	if err != nil {
		t.logger.Error(err)
	}
}

type HandlerFunc func(ctx *Context)

func (h HandlerFunc) Handle(ctx *Context) {
	h(ctx)
}

func WrapBefore(handler http.Handler) HandlerFunc {
	return func(ctx *Context) {
		handler.ServeHTTP(ctx.ResponseWriter, ctx.Req())

		ctx.Next()
	}
}

func WrapAfter(handler http.Handler) HandlerFunc {
	return func(ctx *Context) {
		ctx.Next()

		handler.ServeHTTP(ctx.ResponseWriter, ctx.Req())
	}
}

func (t *Tango) UseHandler(handler http.Handler) {
	t.Use(WrapBefore(handler))
}

func (t *Tango) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	ctx := NewContext(
		t,
		req,
		NewResponseWriter(w),
		t.logger,
	)

	ctx.Invoke()

	// if there is no logging or error handle, so the last written check.
	if !ctx.Written() {
		if ctx.Result == nil {
			ctx.Result = NotFound()
		}
		ctx.HandleError()
		p := req.URL.Path
		if len(req.URL.RawQuery) > 0 {
			p = p + "?" + req.URL.RawQuery
		}

		t.logger.Error(req.Method, ctx.Status(), p)
	}
}

func NewWithLog(logger Logger, handlers ...Handler) *Tango {
	tango := &Tango{
		Router:     NewRouter(),
		Mode:       Env,
		logger:     logger,
		handlers:   make([]Handler, 0),
		ErrHandler: Errors(),
	}

	tango.Use(handlers...)

	return tango
}

func New(handlers ...Handler) *Tango {
	return NewWithLog(NewLogger(os.Stdout), handlers...)
}

func Classic(l ...Logger) *Tango {
	var logger Logger
	if len(l) == 0 {
		logger = NewLogger(os.Stdout)
	} else {
		logger = l[0]
	}

	return NewWithLog(
		logger,
		ClassicHandlers...,
	)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"koding/tools/config"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"

	"github.com/koding/kite/cmd/build"
)

var (
	profile = flag.String("c", "", "Define config profile to be included")
	region  = flag.String("r", "", "Define region profile to be included")

	// Proxy only
	proxy = flag.String("p", "", "Select user proxy or koding proxy")
)

type pkg struct {
	appName       string
	importPath    string
	files         []string
	version       string
	upstartScript string
}

func main() {
	flag.Parse()
	if *profile == "" || *region == "" {
		fmt.Println("Please define config -c and region -r")
		os.Exit(1)
	}

	fmt.Println(*profile, *region, *proxy)

	err := buildPackages()
	if err != nil {
		fmt.Println(err)
	}
}

func buildPackages() error {
	if err := buildKontrolProxy(); err != nil {
		return err
	}

	if err := buildOsKite(); err != nil {
		return err
	}

	return nil
}

func buildOsKite() error {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return errors.New("GOPATH is not set")
	}

	oskitePath := "koding/kites/os"
	temps := struct {
		Profile string
		Region  string
	}{
		Profile: *profile,
		Region:  *region,
	}

	var files = make([]string, 0)
	files = append(files, filepath.Join(gopath, "src", oskitePath, "files"))

	// change our upstartscript because it's a template
	oskiteUpstart := filepath.Join(gopath, "src", oskitePath, "files/oskite.conf")
	configUpstart, err := prepareUpstart(oskiteUpstart, temps)
	if err != nil {
		return err
	}
	defer os.Remove(configUpstart)

	oskite := pkg{
		appName:       "oskite",
		importPath:    oskitePath,
		files:         files,
		version:       "0.1.1",
		upstartScript: configUpstart,
	}

	return oskite.build()
}

func buildKontrolProxy() error {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return errors.New("GOPATH is not set")
	}

	kdproxyPath := "koding/kontrol/kontrolproxy"

	// include certs
	if *proxy == "" {
		return errors.New("Please define proxy target. Example: -p koding or -p user")
	}

	temps := struct {
		Profile   string
		Region    string
		UserProxy string
	}{
		Profile: *profile,
		Region:  *region,
	}

	var files = make([]string, 0)
	switch *proxy {
	case "koding":
		files = append(files, "certs/koding_com_cert.pem", "certs/koding_com_key.pem")
	case "y":
		files = append(files, "certs/y_koding_com_cert.pem", "certs/y_koding_com_key.pem")
	case "x":
		files = append(files, "certs/x_koding_com_cert.pem", "certs/x_koding_com_key.pem")
	case "user":
		temps.UserProxy = "-v"
		files = append(files, "certs/kd_io_cert.pem", "certs/kd_io_key.pem")
	default:
		return errors.New("-p can accept either user or koding")
	}

	files = append(files, filepath.Join(gopath, "src", kdproxyPath, "files"))

	// change our upstartscript because it's a template
	kdproxyUpstart := filepath.Join(gopath, "src", kdproxyPath, "files/kontrolproxy.conf")
	configUpstart, err := prepareUpstart(kdproxyUpstart, temps)
	if err != nil {
		return err
	}
	defer os.Remove(configUpstart)

	kontrolproxy := pkg{
		appName:       "kontrolproxy",
		importPath:    kdproxyPath,
		files:         files,
		version:       "0.0.3",
		upstartScript: configUpstart,
	}

	return kontrolproxy.build()
}

func (p *pkg) build() error {
	fmt.Printf("building '%s' for config '%s' and region '%s'\n", p.appName, *profile, *region)

	// prepare config folder
	tempDir, err := ioutil.TempDir(".", "gopackage_")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	configDir := filepath.Join(tempDir, "config")
	os.MkdirAll(configDir, 0755)

	// koding-config-manager needs it
	err = ioutil.WriteFile("VERSION", []byte(p.version), 0755)
	if err != nil {
		return err
	}
	defer os.Remove("VERSION")

	c, err := config.ReadConfigManager(*profile)
	if err != nil {
		return err
	}

	configPretty, err := json.MarshalIndent(c, "", "  ")
	if err != nil {
		return err
	}

	configFile := filepath.Join(configDir, fmt.Sprintf("main.%s.json", *profile))
	err = ioutil.WriteFile(configFile, configPretty, 0644)
	if err != nil {
		return err
	}

	p.files = append(p.files, configDir)

	// Now it's time to build
	if runtime.GOOS != "linux" {
		return errors.New("Not supported. Please run on a linux machine.")
	}

	deb := &build.Deb{
		AppName:       p.appName,
		Version:       p.version,
		ImportPath:    p.importPath,
		Files:         strings.Join(p.files, ","),
		InstallPrefix: "opt/kite",
		UpstartScript: p.upstartScript,
	}

	debFile, err := deb.Build()
	if err != nil {
		log.Println("linux:", err)
	}

	// rename file to see for which region and env it is created
	oldname := debFile
	newname := fmt.Sprintf("%s_%s_%s-%s_%s.deb", p.appName, p.version, *profile, *region, deb.Arch)

	if err := os.Rename(oldname, newname); err != nil {
		return err
	}

	fmt.Printf("success '%s' is ready. Some helpful commands for you:\n\n", newname)
	fmt.Printf("  show deb content   : dpkg -c %s\n", newname)
	fmt.Printf("  show basic info    : dpkg -f %s\n", newname)
	fmt.Printf("  install to machine : dpkg -i %s\n\n", newname)

	return nil
}

func prepareUpstart(path string, v interface{}) (string, error) {
	file, err := ioutil.TempFile(".", "gopackage_")
	if err != nil {
		return "", err
	}
	defer file.Close()

	t, err := template.ParseFiles(path)
	if err != nil {
		return "", err
	}

	if err := t.Execute(file, v); err != nil {
		return "", err
	}

	return file.Name(), nil
}
<commit_msg>oskite: v 0.1.3<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"koding/tools/config"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"

	"github.com/koding/kite/cmd/build"
)

var (
	profile = flag.String("c", "", "Define config profile to be included")
	region  = flag.String("r", "", "Define region profile to be included")

	// Proxy only
	proxy = flag.String("p", "", "Select user proxy or koding proxy")
)

type pkg struct {
	appName       string
	importPath    string
	files         []string
	version       string
	upstartScript string
}

func main() {
	flag.Parse()
	if *profile == "" || *region == "" {
		fmt.Println("Please define config -c and region -r")
		os.Exit(1)
	}

	fmt.Println(*profile, *region, *proxy)

	err := buildPackages()
	if err != nil {
		fmt.Println(err)
	}
}

func buildPackages() error {
	if err := buildKontrolProxy(); err != nil {
		return err
	}

	if err := buildOsKite(); err != nil {
		return err
	}

	return nil
}

func buildOsKite() error {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return errors.New("GOPATH is not set")
	}

	oskitePath := "koding/kites/os"
	temps := struct {
		Profile string
		Region  string
	}{
		Profile: *profile,
		Region:  *region,
	}

	var files = make([]string, 0)
	files = append(files, filepath.Join(gopath, "src", oskitePath, "files"))

	// change our upstartscript because it's a template
	oskiteUpstart := filepath.Join(gopath, "src", oskitePath, "files/oskite.conf")
	configUpstart, err := prepareUpstart(oskiteUpstart, temps)
	if err != nil {
		return err
	}
	defer os.Remove(configUpstart)

	oskite := pkg{
		appName:       "oskite",
		importPath:    oskitePath,
		files:         files,
		version:       "0.1.3",
		upstartScript: configUpstart,
	}

	return oskite.build()
}

func buildKontrolProxy() error {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return errors.New("GOPATH is not set")
	}

	kdproxyPath := "koding/kontrol/kontrolproxy"

	// include certs
	if *proxy == "" {
		return errors.New("Please define proxy target. Example: -p koding or -p user")
	}

	temps := struct {
		Profile   string
		Region    string
		UserProxy string
	}{
		Profile: *profile,
		Region:  *region,
	}

	var files = make([]string, 0)
	switch *proxy {
	case "koding":
		files = append(files, "certs/koding_com_cert.pem", "certs/koding_com_key.pem")
	case "y":
		files = append(files, "certs/y_koding_com_cert.pem", "certs/y_koding_com_key.pem")
	case "x":
		files = append(files, "certs/x_koding_com_cert.pem", "certs/x_koding_com_key.pem")
	case "user":
		temps.UserProxy = "-v"
		files = append(files, "certs/kd_io_cert.pem", "certs/kd_io_key.pem")
	default:
		return errors.New("-p can accept either user or koding")
	}

	files = append(files, filepath.Join(gopath, "src", kdproxyPath, "files"))

	// change our upstartscript because it's a template
	kdproxyUpstart := filepath.Join(gopath, "src", kdproxyPath, "files/kontrolproxy.conf")
	configUpstart, err := prepareUpstart(kdproxyUpstart, temps)
	if err != nil {
		return err
	}
	defer os.Remove(configUpstart)

	kontrolproxy := pkg{
		appName:       "kontrolproxy",
		importPath:    kdproxyPath,
		files:         files,
		version:       "0.0.3",
		upstartScript: configUpstart,
	}

	return kontrolproxy.build()
}

func (p *pkg) build() error {
	fmt.Printf("building '%s' for config '%s' and region '%s'\n", p.appName, *profile, *region)

	// prepare config folder
	tempDir, err := ioutil.TempDir(".", "gopackage_")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	configDir := filepath.Join(tempDir, "config")
	os.MkdirAll(configDir, 0755)

	// koding-config-manager needs it
	err = ioutil.WriteFile("VERSION", []byte(p.version), 0755)
	if err != nil {
		return err
	}
	defer os.Remove("VERSION")

	c, err := config.ReadConfigManager(*profile)
	if err != nil {
		return err
	}

	configPretty, err := json.MarshalIndent(c, "", "  ")
	if err != nil {
		return err
	}

	configFile := filepath.Join(configDir, fmt.Sprintf("main.%s.json", *profile))
	err = ioutil.WriteFile(configFile, configPretty, 0644)
	if err != nil {
		return err
	}

	p.files = append(p.files, configDir)

	// Now it's time to build
	if runtime.GOOS != "linux" {
		return errors.New("Not supported. Please run on a linux machine.")
	}

	deb := &build.Deb{
		AppName:       p.appName,
		Version:       p.version,
		ImportPath:    p.importPath,
		Files:         strings.Join(p.files, ","),
		InstallPrefix: "opt/kite",
		UpstartScript: p.upstartScript,
	}

	debFile, err := deb.Build()
	if err != nil {
		log.Println("linux:", err)
	}

	// rename file to see for which region and env it is created
	oldname := debFile
	newname := fmt.Sprintf("%s_%s_%s-%s_%s.deb", p.appName, p.version, *profile, *region, deb.Arch)

	if err := os.Rename(oldname, newname); err != nil {
		return err
	}

	fmt.Printf("success '%s' is ready. Some helpful commands for you:\n\n", newname)
	fmt.Printf("  show deb content   : dpkg -c %s\n", newname)
	fmt.Printf("  show basic info    : dpkg -f %s\n", newname)
	fmt.Printf("  install to machine : dpkg -i %s\n\n", newname)

	return nil
}

func prepareUpstart(path string, v interface{}) (string, error) {
	file, err := ioutil.TempFile(".", "gopackage_")
	if err != nil {
		return "", err
	}
	defer file.Close()

	t, err := template.ParseFiles(path)
	if err != nil {
		return "", err
	}

	if err := t.Execute(file, v); err != nil {
		return "", err
	}

	return file.Name(), nil
}
<|endoftext|>


<commit_before>package godbg

import (
	"fmt"
	"os"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	Convey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	Convey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual, `[func.010:95]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual, `[func.009:81]
  test
`)
		})
	})
}

func prbgtest() {
	Pdbgf("prbgtest content")
}
<commit_msg> Fix Test pdbg print with global instance<commit_after>package godbg

import (
	"fmt"
	"os"
	"testing"
	. "github.com/smartystreets/goconvey/convey"
)

func TestProject(t *testing.T) {
	Convey("Test buffers", t, func() {

		Convey("By Default, equals to std", func() {
			So(Out(), ShouldEqual, os.Stdout)
			So(Err(), ShouldEqual, os.Stderr)
		})
		Convey("When set to buffer, no longer equals to std", func() {
			SetBuffers(nil)
			So(Out(), ShouldNotEqual, os.Stdout)
			So(Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance buffer equals to std", func() {
			apdbg := NewPdbg()
			So(apdbg.Out(), ShouldEqual, os.Stdout)
			So(apdbg.Err(), ShouldEqual, os.Stderr)
		})
		Convey("By Default, a new pdbg instance set to buffer writes no longer equals to std", func() {
			apdbg := NewPdbg(SetBuffers)
			So(apdbg.Out(), ShouldNotEqual, os.Stdout)
			So(apdbg.Err(), ShouldNotEqual, os.Stderr)
		})
		Convey("Test custom buffer on global pdbg", func() {
			pdbg.bout = nil
			pdbg.sout = nil
			pdbg.berr = nil
			pdbg.serr = nil
			fmt.Fprintln(Out(), "test0 content0")
			So(OutString(), ShouldEqual, ``)
			fmt.Fprintln(Err(), "err0 content0")
			So(ErrString(), ShouldEqual, ``)
			SetBuffers(nil)
			fmt.Fprintln(Out(), "test content")
			fmt.Fprintln(Err(), "err1 cerr")
			fmt.Fprintln(Err(), "err2 cerr2")
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test content
test2 content2`)
			So(ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})

		Convey("Test custom buffer reset on global pdbg", func() {
			SetBuffers(nil)
			fmt.Fprint(Out(), "test content")
			So(OutString(), ShouldEqual, `test content`)
			fmt.Fprint(Err(), "err1 cerr")
			So(ErrString(), ShouldEqual, `err1 cerr`)
			ResetIOs()
			fmt.Fprint(Out(), "test2 content2")
			So(OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(Err(), "err2 cerr2")
			So(ErrString(), ShouldEqual, `err2 cerr2`)
		})

		Convey("Test custom buffer on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprintln(apdbg.Out(), "test content")
			fmt.Fprintln(apdbg.Err(), "err1 cerr")
			fmt.Fprintln(apdbg.Err(), "err2 cerr2")
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test content
test2 content2`)
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr
err2 cerr2
`)
		})
		Convey("Test custom buffer reset on custom pdbg", func() {
			apdbg := NewPdbg(SetBuffers)
			fmt.Fprint(apdbg.Out(), "test content")
			So(apdbg.OutString(), ShouldEqual, `test content`)
			fmt.Fprint(apdbg.Err(), "err1 cerr")
			So(apdbg.ErrString(), ShouldEqual, `err1 cerr`)
			apdbg.ResetIOs()
			fmt.Fprint(apdbg.Out(), "test2 content2")
			So(apdbg.OutString(), ShouldEqual, `test2 content2`)
			fmt.Fprint(apdbg.Err(), "err2 cerr2")
			So(apdbg.ErrString(), ShouldEqual, `err2 cerr2`)
		})
	})

	Convey("Test pdbg print functions", t, func() {
		Convey("Test pdbg print with global instance", func() {
			SetBuffers(nil)
			Pdbgf("test")
			So(ErrString(), ShouldEqual,
				`[func.010:95]
  test
`)
			ResetIOs()
			prbgtest()
			So(ErrString(), ShouldEqual,
				` [prbgtest:111] (func.010:101)
   prbgtest content
`)
		})
	})
}

func prbgtest() {
	Pdbgf("prbgtest content")
}
<|endoftext|>


<commit_before>package parser

import (
	"fmt"
	"testing"
)

var repoMap = map[string]string{
	"df":  "zerowidth/dotfiles",
	"df2": "zerowidth/dotfiles2", // prefix collision
}

func TestParse(t *testing.T) {
	repoTests := []struct {
		input string // the input
		repo  string // the expected repo match or expansion
		match string // the matched repo shorthand
		desc  string // description of the test case
	}{
		{"", "", "", "no input, no repo match"},
		{"df", "zerowidth/dotfiles", "df", "match match"},
		{" df", "", "", "no match, leading space"},
		{"foo/bar", "foo/bar", "", "fully qualified repo name"},
	}

	repoIssueTests := []struct {
		input string // the input
		repo  string // the expected repo match or expansion
		issue string // the expected issue match
		match string // the matched repo shorthand
		query string // the remaining query text after parsing/expansion
		desc  string // description of the test case
	}{
		{
			input: "",
			desc:  "no issue, no repo",
		},
		{
			input: "df 123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "normal expansion",
		},
		{
			input: "df#123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "expansion with #",
		},
		{
			input: "df #123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "space and # both",
		},
		{
			input: "df123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "prefix match",
		},
		{
			input: "df 1",
			repo:  "zerowidth/dotfiles",
			issue: "1",
			match: "df",
			desc:  "single digit issue",
		},
		{
			input: "df2 34",
			repo:  "zerowidth/dotfiles2",
			issue: "34",
			match: "df2",
			desc:  "numeric suffix on match",
		},
		{
			input: "df234",
			repo:  "zerowidth/dotfiles2",
			issue: "34",
			match: "df2",
			desc:  "numerix suffix with no space",
		},
		{
			input: "foo/bar 123",
			repo:  "foo/bar",
			issue: "123",
			match: "",
			desc:  "fully qualified repo",
		},
		{
			input: "df 0123",
			repo:  "zerowidth/dotfiles",
			issue: "",
			match: "df",
			query: "0123",
			desc:  "invalid issue",
		},
		{
			input: "df foo",
			repo:  "zerowidth/dotfiles",
			match: "df",
			query: "foo",
			desc:  "retrieve query after expansion",
		},
		{
			input: "df 123 foo",
			repo:  "zerowidth/dotfiles",
			issue: "",
			match: "df",
			query: "123 foo",
			desc:  "treats unparsed issue as query",
		},
		{
			input: "123 foo",
			issue: "",
			query: "123 foo",
			desc:  "treats issue with any other text as a query",
		},
		{
			input: "foo bar",
			query: "foo bar",
			desc:  "retrieve query",
		},
		{
			input: "df ",
			repo: "zerowidth/dotfiles",
			match: "df",
			query: "",
			desc: "ignores whitespace after shorthand",
		},
		{
			input: "foo/bar ",
			repo: "foo/bar",
			query: "",
			desc: "ignores whitespace after repo",
		},
	}

	for _, tc := range repoTests {
		t.Run(fmt.Sprintf("Parse(%#v): %s", tc.input, tc.desc), func(t *testing.T) {
			result := Parse(repoMap, tc.input)
			if result.Repo != tc.repo {
				t.Errorf("expected repo %#v, got %#v", tc.repo, result.Repo)
			}
			if result.Match != tc.match {
				t.Errorf("expected match %#v, got %#v", tc.match, result.Match)
			}
		})
	}

	for _, tc := range repoIssueTests {
		t.Run(fmt.Sprintf("Parse(%#v): %s", tc.input, tc.desc), func(t *testing.T) {
			result := Parse(repoMap, tc.input)
			if result.Repo != tc.repo {
				t.Errorf("expected repo %#v, got %#v", tc.repo, result.Repo)
			}
			if result.Issue != tc.issue {
				t.Errorf("expected issue %#v, got %#v", tc.issue, result.Issue)
			}
			if result.Match != tc.match {
				t.Errorf("expected match %#v, got %#v", tc.match, result.Match)
			}
			if result.Query != tc.query {
				t.Errorf("expected query %#v, got %#v", tc.query, result.Query)
			}

		})
	}

}
<commit_msg>Consolidate test cases in parser test<commit_after>package parser

import (
	"fmt"
	"testing"
)

var repoMap = map[string]string{
	"df":  "zerowidth/dotfiles",
	"df2": "zerowidth/dotfiles2", // prefix collision
}

func TestParse(t *testing.T) {
	repoTests := []struct {
		input string // the input
		repo  string // the expected repo match or expansion
		issue string // the expected issue match
		match string // the matched repo shorthand
		query string // the remaining query text after parsing/expansion
		desc  string // description of the test case
	}{
		{
			input: "",
			desc:  "no issue, no repo",
		},
		{
			input: "df",
			repo:  "zerowidth/dotfiles",
			match: "df",
			desc:  "shorthand match",
		},
		{
			input: " df",
			repo:  "",
			match: "",
			query: " df",
			desc:  "no match, leading space",
		},
		{
			input: "foo/bar",
			repo:  "foo/bar",
			match: "",
			desc:  "fully qualified repo name",
		},
		{
			input: "df 123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "normal expansion",
		},
		{
			input: "df#123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "expansion with #",
		},
		{
			input: "df #123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "space and # both",
		},
		{
			input: "df123",
			repo:  "zerowidth/dotfiles",
			issue: "123",
			match: "df",
			desc:  "prefix match",
		},
		{
			input: "df 1",
			repo:  "zerowidth/dotfiles",
			issue: "1",
			match: "df",
			desc:  "single digit issue",
		},
		{
			input: "df2 34",
			repo:  "zerowidth/dotfiles2",
			issue: "34",
			match: "df2",
			desc:  "numeric suffix on match",
		},
		{
			input: "df234",
			repo:  "zerowidth/dotfiles2",
			issue: "34",
			match: "df2",
			desc:  "numerix suffix with no space",
		},
		{
			input: "foo/bar 123",
			repo:  "foo/bar",
			issue: "123",
			match: "",
			desc:  "fully qualified repo",
		},
		{
			input: "df 0123",
			repo:  "zerowidth/dotfiles",
			issue: "",
			match: "df",
			query: "0123",
			desc:  "invalid issue",
		},
		{
			input: "df foo",
			repo:  "zerowidth/dotfiles",
			match: "df",
			query: "foo",
			desc:  "retrieve query after expansion",
		},
		{
			input: "df 123 foo",
			repo:  "zerowidth/dotfiles",
			issue: "",
			match: "df",
			query: "123 foo",
			desc:  "treats unparsed issue as query",
		},
		{
			input: "123 foo",
			issue: "",
			query: "123 foo",
			desc:  "treats issue with any other text as a query",
		},
		{
			input: "foo bar",
			query: "foo bar",
			desc:  "retrieve query",
		},
		{
			input: "df ",
			repo:  "zerowidth/dotfiles",
			match: "df",
			query: "",
			desc:  "ignores whitespace after shorthand",
		},
		{
			input: "foo/bar ",
			repo:  "foo/bar",
			query: "",
			desc:  "ignores whitespace after repo",
		},
	}

	for _, tc := range repoTests {
		t.Run(fmt.Sprintf("Parse(%#v): %s", tc.input, tc.desc), func(t *testing.T) {
			result := Parse(repoMap, tc.input)
			if result.Repo != tc.repo {
				t.Errorf("expected repo %#v, got %#v", tc.repo, result.Repo)
			}
			if result.Issue != tc.issue {
				t.Errorf("expected issue %#v, got %#v", tc.issue, result.Issue)
			}
			if result.Match != tc.match {
				t.Errorf("expected match %#v, got %#v", tc.match, result.Match)
			}
			if result.Query != tc.query {
				t.Errorf("expected query %#v, got %#v", tc.query, result.Query)
			}

		})
	}

}
<|endoftext|>