<commit_before>// Package bytetree provides a Radix tree that stores []byte keys and values.
// See https://en.wikipedia.org/wiki/Radix_tree
package bytetree

import (
	"sync"
	"time"

	"github.com/getlantern/bytemap"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/sql"
)

type Tree struct {
	root   *node
	bytes  int
	length int
	mx     sync.RWMutex
}

type node struct {
	key        []byte
	edges      edges
	data       []encoding.Sequence
	removedFor []int64
}

type edge struct {
	label  []byte
	target *node
}

// New constructs a new Tree.
func New() *Tree {
	return &Tree{root: &node{}}
}

// Bytes returns the number of bytes stored in this Tree (not including
// overhead).
func (bt *Tree) Bytes() int {
	return bt.bytes
}

// Length returns the number of nodes in this Tree.
func (bt *Tree) Length() int {
	return bt.length
}

// Walk walks this Tree, calling the given fn with each node's key and data. If
// the fn returns false, the node will be removed from the Tree as viewed with
// the given ctx. Subsequent walks of the Tree using that same ctx will not see
// removed nodes, but walks using a different context will still see them.
func (bt *Tree) Walk(ctx int64, fn func(key []byte, data []encoding.Sequence) bool) {
	nodes := make([]*node, 0, bt.length)
	nodes = append(nodes, bt.root)
	for {
		if len(nodes) == 0 {
			break
		}
		n := nodes[0]
		nodes = nodes[1:]
		if n.data != nil {
			alreadyRemoved := n.wasRemovedFor(bt, ctx)
			if !alreadyRemoved {
				keep := fn(n.key, n.data)
				if !keep {
					n.doRemoveFor(bt, ctx)
				}
			}
		}
		for _, e := range n.edges {
			nodes = append(nodes, e.target)
		}
	}
}

// Remove removes the given key from this Tree under the given ctx. When viewed
// from this ctx, the key will appear to be gone, but from other contexts it
// will remain visible.
func (bt *Tree) Remove(ctx int64, fullKey []byte) []encoding.Sequence {
	// TODO: basic shape of this is very similar to update, dry violation
	n := bt.root
	key := fullKey
	// Try to update on existing edge
nodeLoop:
	for {
		for _, edge := range n.edges {
			labelLength := len(edge.label)
			keyLength := len(key)
			i := 0
			for ; i < keyLength && i < labelLength; i++ {
				if edge.label[i] != key[i] {
					break
				}
			}
			if i == keyLength && keyLength == labelLength {
				// found it
				alreadyRemoved := edge.target.wasRemovedFor(bt, ctx)
				if alreadyRemoved {
					return nil
				}
				edge.target.doRemoveFor(bt, ctx)
				return edge.target.data
			} else if i == labelLength && labelLength < keyLength {
				// descend
				n = edge.target
				key = key[labelLength:]
				continue nodeLoop
			}
		}

		// not found
		return nil
	}
}

// Copy makes a copy of this Tree.
func (bt *Tree) Copy() *Tree {
	cp := &Tree{bytes: bt.bytes, length: bt.length, root: &node{}}
	nodes := make([]*node, 0, bt.Length())
	nodeCopies := make([]*node, 0, bt.Length())
	nodes = append(nodes, bt.root)
	nodeCopies = append(nodeCopies, cp.root)

	for {
		if len(nodes) == 0 {
			break
		}
		n := nodes[0]
		cpn := nodeCopies[0]
		nodes = nodes[1:]
		nodeCopies = nodeCopies[1:]
		for _, e := range n.edges {
			cpt := &node{key: e.target.key, data: e.target.data}
			cpn.edges = append(cpn.edges, &edge{label: e.label, target: cpt})
			nodes = append(nodes, e.target)
			nodeCopies = append(nodeCopies, cpt)
		}
	}

	return cp
}

// Update updates all of the fields at the given timestamp with the given
// parameters.
func (bt *Tree) Update(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, key []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	bytesAdded, newNode := bt.doUpdate(fields, resolution, truncateBefore, key, vals, metadata)
	bt.bytes += bytesAdded
	if newNode {
		bt.length++
	}
	return bytesAdded
}

func (bt *Tree) doUpdate(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, fullKey []byte, vals encoding.TSParams, metadata bytemap.ByteMap) (int, bool) {
	n := bt.root
	key := fullKey
	// Try to update on existing edge
nodeLoop:
	for {
		for _, edge := range n.edges {
			labelLength := len(edge.label)
			keyLength := len(key)
			i := 0
			for ; i < keyLength && i < labelLength; i++ {
				if edge.label[i] != key[i] {
					break
				}
			}
			if i == keyLength && keyLength == labelLength {
				// update existing node
				return edge.target.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata), false
			} else if i == labelLength && labelLength < keyLength {
				// descend
				n = edge.target
				key = key[labelLength:]
				continue nodeLoop
			} else if i > 0 {
				// common substring, split on that
				return edge.split(bt, fields, resolution, truncateBefore, i, fullKey, key, vals, metadata), true
			}
		}

		// Create new edge
		target := &node{key: fullKey}
		n.edges = append(n.edges, &edge{key, target})
		return target.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata) + len(key), true
	}
}

func (n *node) doUpdate(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, fullKey []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	bytesAdded := 0
	// Grow encoding.Sequences to match number of fields in table
	for i := len(n.data); i < len(fields); i++ {
		n.data = append(n.data, nil)
	}
	for i, field := range fields {
		current := n.data[i]
		previousSize := cap(current)
		updated := current.Update(vals, metadata, field.Expr, resolution, truncateBefore)
		n.data[i] = updated
		bytesAdded += cap(updated) - previousSize
	}
	return bytesAdded
}

func (n *node) wasRemovedFor(bt *Tree, ctx int64) bool {
	if ctx == 0 {
		return false
	}
	bt.mx.RLock()
	for _, _ctx := range n.removedFor {
		if _ctx == ctx {
			bt.mx.RUnlock()
			return true
		}
	}
	bt.mx.RUnlock()
	return false
}

func (n *node) doRemoveFor(bt *Tree, ctx int64) {
	if ctx == 0 {
		return
	}
	bt.mx.Lock()
	n.removedFor = append(n.removedFor, ctx)
	bt.mx.Unlock()
}

func (e *edge) split(bt *Tree, fields []sql.Field, resolution time.Duration, truncateBefore time.Time, splitOn int, fullKey []byte, key []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	newNode := &node{edges: edges{&edge{e.label[splitOn:], e.target}}}
	newLeaf := newNode
	if splitOn != len(key) {
		newLeaf = &node{key: fullKey}
		newNode.edges = append(newNode.edges, &edge{key[splitOn:], newLeaf})
	}
	e.label = e.label[:splitOn]
	e.target = newNode
	return len(key) - splitOn + newLeaf.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata)
}

type edges []*edge
<commit_msg>Reduced allocations in bytetree<commit_after>// Package bytetree provides a Radix tree that stores []byte keys and values.
// See https://en.wikipedia.org/wiki/Radix_tree
package bytetree

import (
	"sync"
	"time"

	"github.com/getlantern/bytemap"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/sql"
)

type Tree struct {
	root   *node
	bytes  int
	length int
	mx     sync.RWMutex
}

type node struct {
	key        []byte
	edges      edges
	data       []encoding.Sequence
	removedFor []int64
}

type edge struct {
	label  []byte
	target *node
}

// New constructs a new Tree.
func New() *Tree {
	return &Tree{root: &node{}}
}

// Bytes returns an estimate of the number of bytes stored in this Tree.
func (bt *Tree) Bytes() int {
	return bt.bytes * 2
}

// Length returns the number of nodes in this Tree.
func (bt *Tree) Length() int {
	return bt.length
}

// Walk walks this Tree, calling the given fn with each node's key and data. If
// the fn returns false, the node will be removed from the Tree as viewed with
// the given ctx. Subsequent walks of the Tree using that same ctx will not see
// removed nodes, but walks using a different context will still see them.
func (bt *Tree) Walk(ctx int64, fn func(key []byte, data []encoding.Sequence) bool) {
	nodes := make([]*node, 0, bt.length)
	nodes = append(nodes, bt.root)
	for {
		if len(nodes) == 0 {
			break
		}
		n := nodes[0]
		nodes = nodes[1:]
		if n.data != nil {
			alreadyRemoved := n.wasRemovedFor(bt, ctx)
			if !alreadyRemoved {
				keep := fn(n.key, n.data)
				if !keep {
					n.doRemoveFor(bt, ctx)
				}
			}
		}
		for _, e := range n.edges {
			nodes = append(nodes, e.target)
		}
	}
}

// Remove removes the given key from this Tree under the given ctx. When viewed
// from this ctx, the key will appear to be gone, but from other contexts it
// will remain visible.
func (bt *Tree) Remove(ctx int64, fullKey []byte) []encoding.Sequence {
	// TODO: basic shape of this is very similar to update, dry violation
	n := bt.root
	key := fullKey
	// Try to update on existing edge
nodeLoop:
	for {
		for _, edge := range n.edges {
			labelLength := len(edge.label)
			keyLength := len(key)
			i := 0
			for ; i < keyLength && i < labelLength; i++ {
				if edge.label[i] != key[i] {
					break
				}
			}
			if i == keyLength && keyLength == labelLength {
				// found it
				alreadyRemoved := edge.target.wasRemovedFor(bt, ctx)
				if alreadyRemoved {
					return nil
				}
				edge.target.doRemoveFor(bt, ctx)
				return edge.target.data
			} else if i == labelLength && labelLength < keyLength {
				// descend
				n = edge.target
				key = key[labelLength:]
				continue nodeLoop
			}
		}

		// not found
		return nil
	}
}

// Copy makes a copy of this Tree.
func (bt *Tree) Copy() *Tree {
	cp := &Tree{bytes: bt.bytes, length: bt.length, root: &node{}}
	nodes := make([]*node, 0, bt.Length())
	nodeCopies := make([]*node, 0, bt.Length())
	nodes = append(nodes, bt.root)
	nodeCopies = append(nodeCopies, cp.root)

	for {
		if len(nodes) == 0 {
			break
		}
		n := nodes[0]
		cpn := nodeCopies[0]
		nodes = nodes[1:]
		nodeCopies = nodeCopies[1:]
		for _, e := range n.edges {
			cpt := &node{key: e.target.key, data: e.target.data}
			cpn.edges = append(cpn.edges, &edge{label: e.label, target: cpt})
			nodes = append(nodes, e.target)
			nodeCopies = append(nodeCopies, cpt)
		}
	}

	return cp
}

// Update updates all of the fields at the given timestamp with the given
// parameters.
func (bt *Tree) Update(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, key []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	bytesAdded, newNode := bt.doUpdate(fields, resolution, truncateBefore, key, vals, metadata)
	bt.bytes += bytesAdded
	if newNode {
		bt.length++
	}
	return bytesAdded
}

func (bt *Tree) doUpdate(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, fullKey []byte, vals encoding.TSParams, metadata bytemap.ByteMap) (int, bool) {
	n := bt.root
	key := fullKey
	// Try to update on existing edge
nodeLoop:
	for {
		for _, edge := range n.edges {
			labelLength := len(edge.label)
			keyLength := len(key)
			i := 0
			for ; i < keyLength && i < labelLength; i++ {
				if edge.label[i] != key[i] {
					break
				}
			}
			if i == keyLength && keyLength == labelLength {
				// update existing node
				return edge.target.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata), false
			} else if i == labelLength && labelLength < keyLength {
				// descend
				n = edge.target
				key = key[labelLength:]
				continue nodeLoop
			} else if i > 0 {
				// common substring, split on that
				return edge.split(bt, fields, resolution, truncateBefore, i, fullKey, key, vals, metadata), true
			}
		}

		// Create new edge
		target := &node{key: fullKey}
		n.edges = append(n.edges, &edge{key, target})
		return target.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata) + len(key), true
	}
}

func (n *node) doUpdate(fields []sql.Field, resolution time.Duration, truncateBefore time.Time, fullKey []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	bytesAdded := 0
	// Grow encoding.Sequences to match number of fields in table
	if len(n.data) < len(fields) {
		newData := make([]encoding.Sequence, len(fields))
		copy(newData, n.data)
		n.data = newData
	}
	for i, field := range fields {
		current := n.data[i]
		previousSize := cap(current)
		updated := current.Update(vals, metadata, field.Expr, resolution, truncateBefore)
		n.data[i] = updated
		bytesAdded += cap(updated) - previousSize
	}
	return bytesAdded
}

func (n *node) wasRemovedFor(bt *Tree, ctx int64) bool {
	if ctx == 0 {
		return false
	}
	bt.mx.RLock()
	for _, _ctx := range n.removedFor {
		if _ctx == ctx {
			bt.mx.RUnlock()
			return true
		}
	}
	bt.mx.RUnlock()
	return false
}

func (n *node) doRemoveFor(bt *Tree, ctx int64) {
	if ctx == 0 {
		return
	}
	bt.mx.Lock()
	n.removedFor = append(n.removedFor, ctx)
	bt.mx.Unlock()
}

func (e *edge) split(bt *Tree, fields []sql.Field, resolution time.Duration, truncateBefore time.Time, splitOn int, fullKey []byte, key []byte, vals encoding.TSParams, metadata bytemap.ByteMap) int {
	newNode := &node{edges: edges{&edge{e.label[splitOn:], e.target}}}
	newLeaf := newNode
	if splitOn != len(key) {
		newLeaf = &node{key: fullKey}
		newNode.edges = append(newNode.edges, &edge{key[splitOn:], newLeaf})
	}
	e.label = e.label[:splitOn]
	e.target = newNode
	return len(key) - splitOn + newLeaf.doUpdate(fields, resolution, truncateBefore, fullKey, vals, metadata)
}

type edges []*edge
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

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

	"github.com/evanphx/json-patch"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/latest"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/runtime"
	utilerrors "k8s.io/kubernetes/pkg/util/errors"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type debugError interface {
	DebugError() (msg string, args []interface{})
}

// AddSourceToErr adds handleResourcePrefix and source string to error message.
// verb is the string like "creating", "deleting" etc.
// souce is the filename or URL to the template file(*.json or *.yaml), or stdin to use to handle the resource.
func AddSourceToErr(verb string, source string, err error) error {
	if source != "" {
		if statusError, ok := err.(*errors.StatusError); ok {
			status := statusError.Status()
			status.Message = fmt.Sprintf("error when %s %q: %v", verb, source, status.Message)
			return &errors.StatusError{ErrStatus: status}
		}
		return fmt.Errorf("error when %s %q: %v", verb, source, err)
	}
	return err
}

// CheckErr prints a user friendly error to STDERR and exits with a non-zero
// exit code. Unrecognized errors will be printed with an "error: " prefix.
//
// This method is generic to the command in use and may be used by non-Kubectl
// commands.
func CheckErr(err error) {
	checkErr(err, fatal)
}

func checkErr(err error, handleErr func(string)) {
	if err == nil {
		return
	}

	if errors.IsInvalid(err) {
		details := err.(*errors.StatusError).Status().Details
		prefix := fmt.Sprintf("The %s %q is invalid.\n", details.Kind, details.Name)
		errs := statusCausesToAggrError(details.Causes)
		handleErr(MultilineError(prefix, errs))
	}

	// handle multiline errors
	if clientcmd.IsConfigurationInvalid(err) {
		handleErr(MultilineError("Error in configuration: ", err))
	}
	if agg, ok := err.(utilerrors.Aggregate); ok && len(agg.Errors()) > 0 {
		handleErr(MultipleErrors("", agg.Errors()))
	}

	msg, ok := StandardErrorMessage(err)
	if !ok {
		msg = fmt.Sprintf("error: %s", err.Error())
	}
	handleErr(msg)
}

func statusCausesToAggrError(scs []api.StatusCause) utilerrors.Aggregate {
	errs := make([]error, len(scs))
	for i, sc := range scs {
		errs[i] = fmt.Errorf("%s: %s", sc.Field, sc.Message)
	}
	return utilerrors.NewAggregate(errs)
}

// StandardErrorMessage translates common errors into a human readable message, or returns
// false if the error is not one of the recognized types. It may also log extended
// information to glog.
//
// This method is generic to the command in use and may be used by non-Kubectl
// commands.
func StandardErrorMessage(err error) (string, bool) {
	if debugErr, ok := err.(debugError); ok {
		glog.V(4).Infof(debugErr.DebugError())
	}
	_, isStatus := err.(client.APIStatus)
	switch {
	case isStatus:
		return fmt.Sprintf("Error from server: %s", err.Error()), true
	case errors.IsUnexpectedObjectError(err):
		return fmt.Sprintf("Server returned an unexpected response: %s", err.Error()), true
	}
	switch t := err.(type) {
	case *url.Error:
		glog.V(4).Infof("Connection error: %s %s: %v", t.Op, t.URL, t.Err)
		switch {
		case strings.Contains(t.Err.Error(), "connection refused"):
			host := t.URL
			if server, err := url.Parse(t.URL); err == nil {
				host = server.Host
			}
			return fmt.Sprintf("The connection to the server %s was refused - did you specify the right host or port?", host), true
		}
		return fmt.Sprintf("Unable to connect to the server: %v", t.Err), true
	}
	return "", false
}

// MultilineError returns a string representing an error that splits sub errors into their own
// lines. The returned string will end with a newline.
func MultilineError(prefix string, err error) string {
	if agg, ok := err.(utilerrors.Aggregate); ok {
		errs := utilerrors.Flatten(agg).Errors()
		buf := &bytes.Buffer{}
		switch len(errs) {
		case 0:
			return fmt.Sprintf("%s%v\n", prefix, err)
		case 1:
			return fmt.Sprintf("%s%v\n", prefix, messageForError(errs[0]))
		default:
			fmt.Fprintln(buf, prefix)
			for _, err := range errs {
				fmt.Fprintf(buf, "* %v\n", messageForError(err))
			}
			return buf.String()
		}
	}
	return fmt.Sprintf("%s%s\n", prefix, err)
}

// MultipleErrors returns a newline delimited string containing
// the prefix and referenced errors in standard form.
func MultipleErrors(prefix string, errs []error) string {
	buf := &bytes.Buffer{}
	for _, err := range errs {
		fmt.Fprintf(buf, "%s%v\n", prefix, messageForError(err))
	}
	return buf.String()
}

// messageForError returns the string representing the error.
func messageForError(err error) string {
	msg, ok := StandardErrorMessage(err)
	if !ok {
		msg = err.Error()
	}
	return msg
}

// fatal prints the message and then exits. If V(2) or greater, glog.Fatal
// is invoked for extended information.
func fatal(msg string) {
	// add newline if needed
	if !strings.HasSuffix(msg, "\n") {
		msg += "\n"
	}

	if glog.V(2) {
		glog.FatalDepth(2, msg)
	}
	fmt.Fprint(os.Stderr, msg)
	os.Exit(1)
}

func UsageError(cmd *cobra.Command, format string, args ...interface{}) error {
	msg := fmt.Sprintf(format, args...)
	return fmt.Errorf("%s\nsee '%s -h' for help.", msg, cmd.CommandPath())
}

func getFlag(cmd *cobra.Command, flag string) *pflag.Flag {
	f := cmd.Flags().Lookup(flag)
	if f == nil {
		glog.Fatalf("flag accessed but not defined for command %s: %s", cmd.Name(), flag)
	}
	return f
}

func GetFlagString(cmd *cobra.Command, flag string) string {
	s, err := cmd.Flags().GetString(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return s
}

// GetFlagStringList can be used to accept multiple argument with flag repetition (e.g. -f arg1 -f arg2 ...)
func GetFlagStringSlice(cmd *cobra.Command, flag string) []string {
	s, err := cmd.Flags().GetStringSlice(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return s
}

// GetWideFlag is used to determine if "-o wide" is used
func GetWideFlag(cmd *cobra.Command) bool {
	f := cmd.Flags().Lookup("output")
	if f.Value.String() == "wide" {
		return true
	}
	return false
}

func GetFlagBool(cmd *cobra.Command, flag string) bool {
	b, err := cmd.Flags().GetBool(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return b
}

// Assumes the flag has a default value.
func GetFlagInt(cmd *cobra.Command, flag string) int {
	i, err := cmd.Flags().GetInt(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return i
}

func GetFlagDuration(cmd *cobra.Command, flag string) time.Duration {
	d, err := cmd.Flags().GetDuration(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return d
}

func AddValidateFlag(cmd *cobra.Command) {
	cmd.Flags().Bool("validate", true, "If true, use a schema to validate the input before sending it")
}

func ReadConfigDataFromReader(reader io.Reader, source string) ([]byte, error) {
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	if len(data) == 0 {
		return nil, fmt.Errorf(`Read from %s but no data found`, source)
	}

	return data, nil
}

// ReadConfigData reads the bytes from the specified filesytem or network
// location or from stdin if location == "-".
// TODO: replace with resource.Builder
func ReadConfigData(location string) ([]byte, error) {
	if len(location) == 0 {
		return nil, fmt.Errorf("location given but empty")
	}

	if location == "-" {
		// Read from stdin.
		return ReadConfigDataFromReader(os.Stdin, "stdin ('-')")
	}

	// Use the location as a file path or URL.
	return ReadConfigDataFromLocation(location)
}

// TODO: replace with resource.Builder
func ReadConfigDataFromLocation(location string) ([]byte, error) {
	// we look for http:// or https:// to determine if valid URL, otherwise do normal file IO
	if strings.Index(location, "http://") == 0 || strings.Index(location, "https://") == 0 {
		resp, err := http.Get(location)
		if err != nil {
			return nil, fmt.Errorf("unable to access URL %s: %v\n", location, err)
		}
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			return nil, fmt.Errorf("unable to read URL, server reported %d %s", resp.StatusCode, resp.Status)
		}
		return ReadConfigDataFromReader(resp.Body, location)
	} else {
		file, err := os.Open(location)
		if err != nil {
			return nil, fmt.Errorf("unable to read %s: %v\n", location, err)
		}
		return ReadConfigDataFromReader(file, location)
	}
}

func Merge(dst runtime.Object, fragment, kind string) (runtime.Object, error) {
	// Ok, this is a little hairy, we'd rather not force the user to specify a kind for their JSON
	// So we pull it into a map, add the Kind field, and then reserialize.
	// We also pull the apiVersion for proper parsing
	var intermediate interface{}
	if err := json.Unmarshal([]byte(fragment), &intermediate); err != nil {
		return nil, err
	}
	dataMap, ok := intermediate.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("Expected a map, found something else: %s", fragment)
	}
	version, found := dataMap["apiVersion"]
	if !found {
		return nil, fmt.Errorf("Inline JSON requires an apiVersion field")
	}
	versionString, ok := version.(string)
	if !ok {
		return nil, fmt.Errorf("apiVersion must be a string")
	}
	i, err := latest.InterfacesFor(versionString)
	if err != nil {
		return nil, err
	}

	// encode dst into versioned json and apply fragment directly too it
	target, err := i.Codec.Encode(dst)
	if err != nil {
		return nil, err
	}
	patched, err := jsonpatch.MergePatch(target, []byte(fragment))
	if err != nil {
		return nil, err
	}
	out, err := i.Codec.Decode(patched)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// DumpReaderToFile writes all data from the given io.Reader to the specified file
// (usually for temporary use).
func DumpReaderToFile(reader io.Reader, filename string) error {
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	defer f.Close()
	if err != nil {
		return err
	}
	buffer := make([]byte, 1024)
	for {
		count, err := reader.Read(buffer)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		_, err = f.Write(buffer[:count])
		if err != nil {
			return err
		}
	}
	return nil
}

// UpdateObject updates resource object with updateFn
func UpdateObject(info *resource.Info, updateFn func(runtime.Object) error) (runtime.Object, error) {
	helper := resource.NewHelper(info.Client, info.Mapping)

	err := updateFn(info.Object)
	if err != nil {
		return nil, err
	}
	data, err := helper.Codec.Encode(info.Object)
	if err != nil {
		return nil, err
	}

	_, err = helper.Replace(info.Namespace, info.Name, true, data)
	if err != nil {
		return nil, err
	}
	return info.Object, nil
}
<commit_msg>Allow to replace os.Exit() with panic when CLI command fatal error<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

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

	"github.com/evanphx/json-patch"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/errors"
	"k8s.io/kubernetes/pkg/api/latest"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/runtime"
	utilerrors "k8s.io/kubernetes/pkg/util/errors"

	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

type debugError interface {
	DebugError() (msg string, args []interface{})
}

// AddSourceToErr adds handleResourcePrefix and source string to error message.
// verb is the string like "creating", "deleting" etc.
// souce is the filename or URL to the template file(*.json or *.yaml), or stdin to use to handle the resource.
func AddSourceToErr(verb string, source string, err error) error {
	if source != "" {
		if statusError, ok := err.(*errors.StatusError); ok {
			status := statusError.Status()
			status.Message = fmt.Sprintf("error when %s %q: %v", verb, source, status.Message)
			return &errors.StatusError{ErrStatus: status}
		}
		return fmt.Errorf("error when %s %q: %v", verb, source, err)
	}
	return err
}

var fatalErrHandler = fatal

// BehaviorOnFatal allows you to override the default behavior when a fatal
// error occurs, which is call os.Exit(1). You can pass 'panic' as a function
// here if you prefer the panic() over os.Exit(1).
func BehaviorOnFatal(f func(string)) {
	fatalErrHandler = f
}

// fatal prints the message and then exits. If V(2) or greater, glog.Fatal
// is invoked for extended information.
func fatal(msg string) {
	// add newline if needed
	if !strings.HasSuffix(msg, "\n") {
		msg += "\n"
	}

	if glog.V(2) {
		glog.FatalDepth(2, msg)
	}
	fmt.Fprint(os.Stderr, msg)
	os.Exit(1)
}

// CheckErr prints a user friendly error to STDERR and exits with a non-zero
// exit code. Unrecognized errors will be printed with an "error: " prefix.
//
// This method is generic to the command in use and may be used by non-Kubectl
// commands.
func CheckErr(err error) {
	checkErr(err, fatalErrHandler)
}

func checkErr(err error, handleErr func(string)) {
	if err == nil {
		return
	}

	if errors.IsInvalid(err) {
		details := err.(*errors.StatusError).Status().Details
		prefix := fmt.Sprintf("The %s %q is invalid.\n", details.Kind, details.Name)
		errs := statusCausesToAggrError(details.Causes)
		handleErr(MultilineError(prefix, errs))
	}

	// handle multiline errors
	if clientcmd.IsConfigurationInvalid(err) {
		handleErr(MultilineError("Error in configuration: ", err))
	}
	if agg, ok := err.(utilerrors.Aggregate); ok && len(agg.Errors()) > 0 {
		handleErr(MultipleErrors("", agg.Errors()))
	}

	msg, ok := StandardErrorMessage(err)
	if !ok {
		msg = fmt.Sprintf("error: %s", err.Error())
	}
	handleErr(msg)
}

func statusCausesToAggrError(scs []api.StatusCause) utilerrors.Aggregate {
	errs := make([]error, len(scs))
	for i, sc := range scs {
		errs[i] = fmt.Errorf("%s: %s", sc.Field, sc.Message)
	}
	return utilerrors.NewAggregate(errs)
}

// StandardErrorMessage translates common errors into a human readable message, or returns
// false if the error is not one of the recognized types. It may also log extended
// information to glog.
//
// This method is generic to the command in use and may be used by non-Kubectl
// commands.
func StandardErrorMessage(err error) (string, bool) {
	if debugErr, ok := err.(debugError); ok {
		glog.V(4).Infof(debugErr.DebugError())
	}
	_, isStatus := err.(client.APIStatus)
	switch {
	case isStatus:
		return fmt.Sprintf("Error from server: %s", err.Error()), true
	case errors.IsUnexpectedObjectError(err):
		return fmt.Sprintf("Server returned an unexpected response: %s", err.Error()), true
	}
	switch t := err.(type) {
	case *url.Error:
		glog.V(4).Infof("Connection error: %s %s: %v", t.Op, t.URL, t.Err)
		switch {
		case strings.Contains(t.Err.Error(), "connection refused"):
			host := t.URL
			if server, err := url.Parse(t.URL); err == nil {
				host = server.Host
			}
			return fmt.Sprintf("The connection to the server %s was refused - did you specify the right host or port?", host), true
		}
		return fmt.Sprintf("Unable to connect to the server: %v", t.Err), true
	}
	return "", false
}

// MultilineError returns a string representing an error that splits sub errors into their own
// lines. The returned string will end with a newline.
func MultilineError(prefix string, err error) string {
	if agg, ok := err.(utilerrors.Aggregate); ok {
		errs := utilerrors.Flatten(agg).Errors()
		buf := &bytes.Buffer{}
		switch len(errs) {
		case 0:
			return fmt.Sprintf("%s%v\n", prefix, err)
		case 1:
			return fmt.Sprintf("%s%v\n", prefix, messageForError(errs[0]))
		default:
			fmt.Fprintln(buf, prefix)
			for _, err := range errs {
				fmt.Fprintf(buf, "* %v\n", messageForError(err))
			}
			return buf.String()
		}
	}
	return fmt.Sprintf("%s%s\n", prefix, err)
}

// MultipleErrors returns a newline delimited string containing
// the prefix and referenced errors in standard form.
func MultipleErrors(prefix string, errs []error) string {
	buf := &bytes.Buffer{}
	for _, err := range errs {
		fmt.Fprintf(buf, "%s%v\n", prefix, messageForError(err))
	}
	return buf.String()
}

// messageForError returns the string representing the error.
func messageForError(err error) string {
	msg, ok := StandardErrorMessage(err)
	if !ok {
		msg = err.Error()
	}
	return msg
}

func UsageError(cmd *cobra.Command, format string, args ...interface{}) error {
	msg := fmt.Sprintf(format, args...)
	return fmt.Errorf("%s\nsee '%s -h' for help.", msg, cmd.CommandPath())
}

func getFlag(cmd *cobra.Command, flag string) *pflag.Flag {
	f := cmd.Flags().Lookup(flag)
	if f == nil {
		glog.Fatalf("flag accessed but not defined for command %s: %s", cmd.Name(), flag)
	}
	return f
}

func GetFlagString(cmd *cobra.Command, flag string) string {
	s, err := cmd.Flags().GetString(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return s
}

// GetFlagStringList can be used to accept multiple argument with flag repetition (e.g. -f arg1 -f arg2 ...)
func GetFlagStringSlice(cmd *cobra.Command, flag string) []string {
	s, err := cmd.Flags().GetStringSlice(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return s
}

// GetWideFlag is used to determine if "-o wide" is used
func GetWideFlag(cmd *cobra.Command) bool {
	f := cmd.Flags().Lookup("output")
	if f.Value.String() == "wide" {
		return true
	}
	return false
}

func GetFlagBool(cmd *cobra.Command, flag string) bool {
	b, err := cmd.Flags().GetBool(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return b
}

// Assumes the flag has a default value.
func GetFlagInt(cmd *cobra.Command, flag string) int {
	i, err := cmd.Flags().GetInt(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return i
}

func GetFlagDuration(cmd *cobra.Command, flag string) time.Duration {
	d, err := cmd.Flags().GetDuration(flag)
	if err != nil {
		glog.Fatalf("err accessing flag %s for command %s: %v", flag, cmd.Name(), err)
	}
	return d
}

func AddValidateFlag(cmd *cobra.Command) {
	cmd.Flags().Bool("validate", true, "If true, use a schema to validate the input before sending it")
}

func ReadConfigDataFromReader(reader io.Reader, source string) ([]byte, error) {
	data, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	if len(data) == 0 {
		return nil, fmt.Errorf(`Read from %s but no data found`, source)
	}

	return data, nil
}

// ReadConfigData reads the bytes from the specified filesytem or network
// location or from stdin if location == "-".
// TODO: replace with resource.Builder
func ReadConfigData(location string) ([]byte, error) {
	if len(location) == 0 {
		return nil, fmt.Errorf("location given but empty")
	}

	if location == "-" {
		// Read from stdin.
		return ReadConfigDataFromReader(os.Stdin, "stdin ('-')")
	}

	// Use the location as a file path or URL.
	return ReadConfigDataFromLocation(location)
}

// TODO: replace with resource.Builder
func ReadConfigDataFromLocation(location string) ([]byte, error) {
	// we look for http:// or https:// to determine if valid URL, otherwise do normal file IO
	if strings.Index(location, "http://") == 0 || strings.Index(location, "https://") == 0 {
		resp, err := http.Get(location)
		if err != nil {
			return nil, fmt.Errorf("unable to access URL %s: %v\n", location, err)
		}
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			return nil, fmt.Errorf("unable to read URL, server reported %d %s", resp.StatusCode, resp.Status)
		}
		return ReadConfigDataFromReader(resp.Body, location)
	} else {
		file, err := os.Open(location)
		if err != nil {
			return nil, fmt.Errorf("unable to read %s: %v\n", location, err)
		}
		return ReadConfigDataFromReader(file, location)
	}
}

func Merge(dst runtime.Object, fragment, kind string) (runtime.Object, error) {
	// Ok, this is a little hairy, we'd rather not force the user to specify a kind for their JSON
	// So we pull it into a map, add the Kind field, and then reserialize.
	// We also pull the apiVersion for proper parsing
	var intermediate interface{}
	if err := json.Unmarshal([]byte(fragment), &intermediate); err != nil {
		return nil, err
	}
	dataMap, ok := intermediate.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("Expected a map, found something else: %s", fragment)
	}
	version, found := dataMap["apiVersion"]
	if !found {
		return nil, fmt.Errorf("Inline JSON requires an apiVersion field")
	}
	versionString, ok := version.(string)
	if !ok {
		return nil, fmt.Errorf("apiVersion must be a string")
	}
	i, err := latest.InterfacesFor(versionString)
	if err != nil {
		return nil, err
	}

	// encode dst into versioned json and apply fragment directly too it
	target, err := i.Codec.Encode(dst)
	if err != nil {
		return nil, err
	}
	patched, err := jsonpatch.MergePatch(target, []byte(fragment))
	if err != nil {
		return nil, err
	}
	out, err := i.Codec.Decode(patched)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// DumpReaderToFile writes all data from the given io.Reader to the specified file
// (usually for temporary use).
func DumpReaderToFile(reader io.Reader, filename string) error {
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	defer f.Close()
	if err != nil {
		return err
	}
	buffer := make([]byte, 1024)
	for {
		count, err := reader.Read(buffer)
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		_, err = f.Write(buffer[:count])
		if err != nil {
			return err
		}
	}
	return nil
}

// UpdateObject updates resource object with updateFn
func UpdateObject(info *resource.Info, updateFn func(runtime.Object) error) (runtime.Object, error) {
	helper := resource.NewHelper(info.Client, info.Mapping)

	err := updateFn(info.Object)
	if err != nil {
		return nil, err
	}
	data, err := helper.Codec.Encode(info.Object)
	if err != nil {
		return nil, err
	}

	_, err = helper.Replace(info.Namespace, info.Name, true, data)
	if err != nil {
		return nil, err
	}
	return info.Object, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/miekg/dns"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

const (
	testQuery    = "_test.consul."
	consulDomain = "consul."
)

// DNSServer is used to wrap an Agent and expose various
// service discovery endpoints using a DNS interface.
type DNSServer struct {
	agent        *Agent
	dnsHandler   *dns.ServeMux
	dnsServer    *dns.Server
	dnsServerTCP *dns.Server
	domain       string
	recursor     string
	logger       *log.Logger
}

// NewDNSServer starts a new DNS server to provide an agent interface
func NewDNSServer(agent *Agent, logOutput io.Writer, domain, bind, recursor string) (*DNSServer, error) {
	// Make sure domain is FQDN
	domain = dns.Fqdn(domain)

	// Construct the DNS components
	mux := dns.NewServeMux()

	// Setup the servers
	server := &dns.Server{
		Addr:    bind,
		Net:     "udp",
		Handler: mux,
		UDPSize: 65535,
	}
	serverTCP := &dns.Server{
		Addr:    bind,
		Net:     "tcp",
		Handler: mux,
	}

	// Create the server
	srv := &DNSServer{
		agent:        agent,
		dnsHandler:   mux,
		dnsServer:    server,
		dnsServerTCP: serverTCP,
		domain:       domain,
		recursor:     recursor,
		logger:       log.New(logOutput, "", log.LstdFlags),
	}

	// Register mux handlers, always handle "consul."
	mux.HandleFunc(domain, srv.handleQuery)
	if domain != consulDomain {
		mux.HandleFunc(consulDomain, srv.handleTest)
	}
	if recursor != "" {
		mux.HandleFunc(".", srv.handleRecurse)
	}

	// Async start the DNS Servers, handle a potential error
	errCh := make(chan error, 1)
	go func() {
		err := server.ListenAndServe()
		srv.logger.Printf("[ERR] dns: error starting udp server: %v", err)
		errCh <- err
	}()

	errChTCP := make(chan error, 1)
	go func() {
		err := serverTCP.ListenAndServe()
		srv.logger.Printf("[ERR] dns: error starting tcp server: %v", err)
		errChTCP <- err
	}()

	// Check the server is running, do a test lookup
	checkCh := make(chan error, 1)
	go func() {
		// This is jank, but we have no way to edge trigger on
		// the start of our server, so we just wait and hope it is up.
		time.Sleep(50 * time.Millisecond)

		m := new(dns.Msg)
		m.SetQuestion(testQuery, dns.TypeANY)

		c := new(dns.Client)
		in, _, err := c.Exchange(m, bind)
		if err != nil {
			checkCh <- err
			return
		}

		if len(in.Answer) == 0 {
			checkCh <- fmt.Errorf("no response to test message")
			return
		}
		close(checkCh)
	}()

	// Wait for either the check, listen error, or timeout
	select {
	case e := <-errCh:
		return srv, e
	case e := <-errChTCP:
		return srv, e
	case e := <-checkCh:
		return srv, e
	case <-time.After(time.Second):
		return srv, fmt.Errorf("timeout setting up DNS server")
	}
	return srv, nil
}

// handleQUery is used to handle DNS queries in the configured domain
func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%v)", q, time.Now().Sub(s))
	}(time.Now())

	// Check if this is potentially a test query
	if q.Name == testQuery {
		d.handleTest(resp, req)
		return
	}

	// Setup the message response
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true
	d.addSOA(d.domain, m)

	// Dispatch the correct handler
	d.dispatch(req, m)

	// Write out the complete response
	if err := resp.WriteMsg(m); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}

// handleTest is used to handle DNS queries in the ".consul." domain
func (d *DNSServer) handleTest(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%v)", q, time.Now().Sub(s))
	}(time.Now())

	if !(q.Qtype == dns.TypeANY || q.Qtype == dns.TypeTXT) {
		return
	}
	if q.Name != testQuery {
		return
	}

	// Always respond with TXT "ok"
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true
	header := dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
	txt := &dns.TXT{header, []string{"ok"}}
	m.Answer = append(m.Answer, txt)
	d.addSOA(consulDomain, m)
	if err := resp.WriteMsg(m); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}

// addSOA is used to add an SOA record to a message for the given domain
func (d *DNSServer) addSOA(domain string, msg *dns.Msg) {
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   domain,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		Ns:      "ns." + domain,
		Mbox:    "postmaster." + domain,
		Serial:  uint32(time.Now().Unix()),
		Refresh: 3600,
		Retry:   600,
		Expire:  86400,
		Minttl:  0,
	}
	msg.Ns = append(msg.Ns, soa)
}

// dispatch is used to parse a request and invoke the correct handler
func (d *DNSServer) dispatch(req, resp *dns.Msg) {
	// By default the query is in the default datacenter
	datacenter := d.agent.config.Datacenter

	// Get the QName without the domain suffix
	qName := dns.Fqdn(req.Question[0].Name)
	qName = strings.TrimSuffix(qName, d.domain)

	// Split into the label parts
	labels := dns.SplitDomainName(qName)

	// The last label is either "node", "service" or a datacenter name
PARSE:
	if len(labels) == 0 {
		goto INVALID
	}
	switch labels[len(labels)-1] {
	case "service":
		// Handle lookup with and without tag
		switch len(labels) {
		case 2:
			d.serviceLookup(datacenter, labels[0], "", req, resp)
		case 3:
			d.serviceLookup(datacenter, labels[1], labels[0], req, resp)
		default:
			goto INVALID
		}

	case "node":
		if len(labels) != 2 {
			goto INVALID
		}
		d.nodeLookup(datacenter, labels[0], req, resp)

	default:
		// Store the DC, and re-parse
		datacenter = labels[len(labels)-1]
		labels = labels[:len(labels)-1]
		goto PARSE
	}
	return
INVALID:
	d.logger.Printf("[WARN] dns: QName invalid: %s", qName)
	resp.SetRcode(req, dns.RcodeNameError)
}

// nodeLookup is used to handle a node query
func (d *DNSServer) nodeLookup(datacenter, node string, req, resp *dns.Msg) {
	// Only handle ANY and A type requests
	qType := req.Question[0].Qtype
	if qType != dns.TypeANY && qType != dns.TypeA {
		return
	}

	// Make an RPC request
	args := structs.NodeServicesRequest{
		Datacenter: datacenter,
		Node:       node,
	}
	var out structs.NodeServices
	if err := d.agent.RPC("Catalog.NodeServices", &args, &out); err != nil {
		d.logger.Printf("[ERR] dns: rpc error: %v", err)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// If we have no address, return not found!
	if out.Address == "" {
		resp.SetRcode(req, dns.RcodeNameError)
		return
	}

	// Parse the IP
	ip := net.ParseIP(out.Address)
	if ip == nil {
		d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", out.Address, node)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// Format A record
	aRec := &dns.A{
		Hdr: dns.RR_Header{
			Name:   req.Question[0].Name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		A: ip,
	}

	// Add the response
	resp.Answer = append(resp.Answer, aRec)
}

// serviceLookup is used to handle a service query
func (d *DNSServer) serviceLookup(datacenter, service, tag string, req, resp *dns.Msg) {
	// Make an RPC request
	args := structs.ServiceNodesRequest{
		Datacenter:  datacenter,
		ServiceName: service,
		ServiceTag:  tag,
		TagFilter:   tag != "",
	}
	var out structs.ServiceNodes
	if err := d.agent.RPC("Catalog.ServiceNodes", &args, &out); err != nil {
		d.logger.Printf("[ERR] dns: rpc error: %v", err)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// If we have no nodes, return not found!
	if len(out) == 0 {
		resp.SetRcode(req, dns.RcodeNameError)
		return
	}

	// Add various responses depending on the request
	qType := req.Question[0].Qtype
	if qType == dns.TypeANY || qType == dns.TypeA {
		d.serviceARecords(out, req, resp)
	}
	if qType == dns.TypeANY || qType == dns.TypeSRV {
		d.serviceSRVRecords(datacenter, out, req, resp)
	}
}

// serviceARecords is used to add the A records for a service lookup
func (d *DNSServer) serviceARecords(nodes structs.ServiceNodes, req, resp *dns.Msg) {
	for _, node := range nodes {
		ip := net.ParseIP(node.Address)
		if ip == nil {
			d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", node.Address, node.Node)
			continue
		}
		aRec := &dns.A{
			Hdr: dns.RR_Header{
				Name:   req.Question[0].Name,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			A: ip,
		}
		resp.Answer = append(resp.Answer, aRec)
	}
}

// serviceARecords is used to add the SRV records for a service lookup
func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.ServiceNodes, req, resp *dns.Msg) {
	for _, node := range nodes {
		srvRec := &dns.SRV{
			Hdr: dns.RR_Header{
				Name:   req.Question[0].Name,
				Rrtype: dns.TypeSRV,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			Priority: 1,
			Weight:   1,
			Port:     uint16(node.ServicePort),
			Target:   fmt.Sprintf("%s.node.%s.%s", node.Node, dc, d.domain),
		}
		resp.Answer = append(resp.Answer, srvRec)

		ip := net.ParseIP(node.Address)
		if ip == nil {
			d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", node.Address, node.Node)
			continue
		}
		aRec := &dns.A{
			Hdr: dns.RR_Header{
				Name:   srvRec.Target,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			A: ip,
		}
		resp.Extra = append(resp.Extra, aRec)
	}
}

// handleRecurse is used to handle recursive DNS queries
func (d *DNSServer) handleRecurse(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	network := "udp"
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%s) (%v)", q, network, time.Now().Sub(s))
	}(time.Now())

	// Switch to TCP if the client is
	if _, ok := resp.RemoteAddr().(*net.TCPAddr); ok {
		network = "tcp"
	}

	// Recursively resolve
	c := &dns.Client{Net: network}
	r, rtt, err := c.Exchange(req, d.recursor)

	// On failure, return a SERVFAIL message
	if err != nil {
		d.logger.Printf("[ERR] dns: recurse failed: %v", err)
		m := &dns.Msg{}
		m.SetReply(req)
		m.SetRcode(req, dns.RcodeServerFailure)
		resp.WriteMsg(m)
		return
	}
	d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v)", q, rtt)

	// Seems to be a bug that forcing compression fixes...
	r.Compress = true

	// Forward the response
	if err := resp.WriteMsg(r); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}
<commit_msg>Remove fix for upstream DNS bug<commit_after>package agent

import (
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/miekg/dns"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

const (
	testQuery    = "_test.consul."
	consulDomain = "consul."
)

// DNSServer is used to wrap an Agent and expose various
// service discovery endpoints using a DNS interface.
type DNSServer struct {
	agent        *Agent
	dnsHandler   *dns.ServeMux
	dnsServer    *dns.Server
	dnsServerTCP *dns.Server
	domain       string
	recursor     string
	logger       *log.Logger
}

// NewDNSServer starts a new DNS server to provide an agent interface
func NewDNSServer(agent *Agent, logOutput io.Writer, domain, bind, recursor string) (*DNSServer, error) {
	// Make sure domain is FQDN
	domain = dns.Fqdn(domain)

	// Construct the DNS components
	mux := dns.NewServeMux()

	// Setup the servers
	server := &dns.Server{
		Addr:    bind,
		Net:     "udp",
		Handler: mux,
		UDPSize: 65535,
	}
	serverTCP := &dns.Server{
		Addr:    bind,
		Net:     "tcp",
		Handler: mux,
	}

	// Create the server
	srv := &DNSServer{
		agent:        agent,
		dnsHandler:   mux,
		dnsServer:    server,
		dnsServerTCP: serverTCP,
		domain:       domain,
		recursor:     recursor,
		logger:       log.New(logOutput, "", log.LstdFlags),
	}

	// Register mux handlers, always handle "consul."
	mux.HandleFunc(domain, srv.handleQuery)
	if domain != consulDomain {
		mux.HandleFunc(consulDomain, srv.handleTest)
	}
	if recursor != "" {
		mux.HandleFunc(".", srv.handleRecurse)
	}

	// Async start the DNS Servers, handle a potential error
	errCh := make(chan error, 1)
	go func() {
		err := server.ListenAndServe()
		srv.logger.Printf("[ERR] dns: error starting udp server: %v", err)
		errCh <- err
	}()

	errChTCP := make(chan error, 1)
	go func() {
		err := serverTCP.ListenAndServe()
		srv.logger.Printf("[ERR] dns: error starting tcp server: %v", err)
		errChTCP <- err
	}()

	// Check the server is running, do a test lookup
	checkCh := make(chan error, 1)
	go func() {
		// This is jank, but we have no way to edge trigger on
		// the start of our server, so we just wait and hope it is up.
		time.Sleep(50 * time.Millisecond)

		m := new(dns.Msg)
		m.SetQuestion(testQuery, dns.TypeANY)

		c := new(dns.Client)
		in, _, err := c.Exchange(m, bind)
		if err != nil {
			checkCh <- err
			return
		}

		if len(in.Answer) == 0 {
			checkCh <- fmt.Errorf("no response to test message")
			return
		}
		close(checkCh)
	}()

	// Wait for either the check, listen error, or timeout
	select {
	case e := <-errCh:
		return srv, e
	case e := <-errChTCP:
		return srv, e
	case e := <-checkCh:
		return srv, e
	case <-time.After(time.Second):
		return srv, fmt.Errorf("timeout setting up DNS server")
	}
	return srv, nil
}

// handleQUery is used to handle DNS queries in the configured domain
func (d *DNSServer) handleQuery(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%v)", q, time.Now().Sub(s))
	}(time.Now())

	// Check if this is potentially a test query
	if q.Name == testQuery {
		d.handleTest(resp, req)
		return
	}

	// Setup the message response
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true
	d.addSOA(d.domain, m)

	// Dispatch the correct handler
	d.dispatch(req, m)

	// Write out the complete response
	if err := resp.WriteMsg(m); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}

// handleTest is used to handle DNS queries in the ".consul." domain
func (d *DNSServer) handleTest(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%v)", q, time.Now().Sub(s))
	}(time.Now())

	if !(q.Qtype == dns.TypeANY || q.Qtype == dns.TypeTXT) {
		return
	}
	if q.Name != testQuery {
		return
	}

	// Always respond with TXT "ok"
	m := new(dns.Msg)
	m.SetReply(req)
	m.Authoritative = true
	header := dns.RR_Header{Name: q.Name, Rrtype: dns.TypeTXT, Class: dns.ClassINET, Ttl: 0}
	txt := &dns.TXT{header, []string{"ok"}}
	m.Answer = append(m.Answer, txt)
	d.addSOA(consulDomain, m)
	if err := resp.WriteMsg(m); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}

// addSOA is used to add an SOA record to a message for the given domain
func (d *DNSServer) addSOA(domain string, msg *dns.Msg) {
	soa := &dns.SOA{
		Hdr: dns.RR_Header{
			Name:   domain,
			Rrtype: dns.TypeSOA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		Ns:      "ns." + domain,
		Mbox:    "postmaster." + domain,
		Serial:  uint32(time.Now().Unix()),
		Refresh: 3600,
		Retry:   600,
		Expire:  86400,
		Minttl:  0,
	}
	msg.Ns = append(msg.Ns, soa)
}

// dispatch is used to parse a request and invoke the correct handler
func (d *DNSServer) dispatch(req, resp *dns.Msg) {
	// By default the query is in the default datacenter
	datacenter := d.agent.config.Datacenter

	// Get the QName without the domain suffix
	qName := dns.Fqdn(req.Question[0].Name)
	qName = strings.TrimSuffix(qName, d.domain)

	// Split into the label parts
	labels := dns.SplitDomainName(qName)

	// The last label is either "node", "service" or a datacenter name
PARSE:
	if len(labels) == 0 {
		goto INVALID
	}
	switch labels[len(labels)-1] {
	case "service":
		// Handle lookup with and without tag
		switch len(labels) {
		case 2:
			d.serviceLookup(datacenter, labels[0], "", req, resp)
		case 3:
			d.serviceLookup(datacenter, labels[1], labels[0], req, resp)
		default:
			goto INVALID
		}

	case "node":
		if len(labels) != 2 {
			goto INVALID
		}
		d.nodeLookup(datacenter, labels[0], req, resp)

	default:
		// Store the DC, and re-parse
		datacenter = labels[len(labels)-1]
		labels = labels[:len(labels)-1]
		goto PARSE
	}
	return
INVALID:
	d.logger.Printf("[WARN] dns: QName invalid: %s", qName)
	resp.SetRcode(req, dns.RcodeNameError)
}

// nodeLookup is used to handle a node query
func (d *DNSServer) nodeLookup(datacenter, node string, req, resp *dns.Msg) {
	// Only handle ANY and A type requests
	qType := req.Question[0].Qtype
	if qType != dns.TypeANY && qType != dns.TypeA {
		return
	}

	// Make an RPC request
	args := structs.NodeServicesRequest{
		Datacenter: datacenter,
		Node:       node,
	}
	var out structs.NodeServices
	if err := d.agent.RPC("Catalog.NodeServices", &args, &out); err != nil {
		d.logger.Printf("[ERR] dns: rpc error: %v", err)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// If we have no address, return not found!
	if out.Address == "" {
		resp.SetRcode(req, dns.RcodeNameError)
		return
	}

	// Parse the IP
	ip := net.ParseIP(out.Address)
	if ip == nil {
		d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", out.Address, node)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// Format A record
	aRec := &dns.A{
		Hdr: dns.RR_Header{
			Name:   req.Question[0].Name,
			Rrtype: dns.TypeA,
			Class:  dns.ClassINET,
			Ttl:    0,
		},
		A: ip,
	}

	// Add the response
	resp.Answer = append(resp.Answer, aRec)
}

// serviceLookup is used to handle a service query
func (d *DNSServer) serviceLookup(datacenter, service, tag string, req, resp *dns.Msg) {
	// Make an RPC request
	args := structs.ServiceNodesRequest{
		Datacenter:  datacenter,
		ServiceName: service,
		ServiceTag:  tag,
		TagFilter:   tag != "",
	}
	var out structs.ServiceNodes
	if err := d.agent.RPC("Catalog.ServiceNodes", &args, &out); err != nil {
		d.logger.Printf("[ERR] dns: rpc error: %v", err)
		resp.SetRcode(req, dns.RcodeServerFailure)
		return
	}

	// If we have no nodes, return not found!
	if len(out) == 0 {
		resp.SetRcode(req, dns.RcodeNameError)
		return
	}

	// Add various responses depending on the request
	qType := req.Question[0].Qtype
	if qType == dns.TypeANY || qType == dns.TypeA {
		d.serviceARecords(out, req, resp)
	}
	if qType == dns.TypeANY || qType == dns.TypeSRV {
		d.serviceSRVRecords(datacenter, out, req, resp)
	}
}

// serviceARecords is used to add the A records for a service lookup
func (d *DNSServer) serviceARecords(nodes structs.ServiceNodes, req, resp *dns.Msg) {
	for _, node := range nodes {
		ip := net.ParseIP(node.Address)
		if ip == nil {
			d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", node.Address, node.Node)
			continue
		}
		aRec := &dns.A{
			Hdr: dns.RR_Header{
				Name:   req.Question[0].Name,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			A: ip,
		}
		resp.Answer = append(resp.Answer, aRec)
	}
}

// serviceARecords is used to add the SRV records for a service lookup
func (d *DNSServer) serviceSRVRecords(dc string, nodes structs.ServiceNodes, req, resp *dns.Msg) {
	for _, node := range nodes {
		srvRec := &dns.SRV{
			Hdr: dns.RR_Header{
				Name:   req.Question[0].Name,
				Rrtype: dns.TypeSRV,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			Priority: 1,
			Weight:   1,
			Port:     uint16(node.ServicePort),
			Target:   fmt.Sprintf("%s.node.%s.%s", node.Node, dc, d.domain),
		}
		resp.Answer = append(resp.Answer, srvRec)

		ip := net.ParseIP(node.Address)
		if ip == nil {
			d.logger.Printf("[ERR] dns: failed to parse IP %v for %v", node.Address, node.Node)
			continue
		}
		aRec := &dns.A{
			Hdr: dns.RR_Header{
				Name:   srvRec.Target,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			A: ip,
		}
		resp.Extra = append(resp.Extra, aRec)
	}
}

// handleRecurse is used to handle recursive DNS queries
func (d *DNSServer) handleRecurse(resp dns.ResponseWriter, req *dns.Msg) {
	q := req.Question[0]
	network := "udp"
	defer func(s time.Time) {
		d.logger.Printf("[DEBUG] dns: request for %v (%s) (%v)", q, network, time.Now().Sub(s))
	}(time.Now())

	// Switch to TCP if the client is
	if _, ok := resp.RemoteAddr().(*net.TCPAddr); ok {
		network = "tcp"
	}

	// Recursively resolve
	c := &dns.Client{Net: network}
	r, rtt, err := c.Exchange(req, d.recursor)

	// On failure, return a SERVFAIL message
	if err != nil {
		d.logger.Printf("[ERR] dns: recurse failed: %v", err)
		m := &dns.Msg{}
		m.SetReply(req)
		m.SetRcode(req, dns.RcodeServerFailure)
		resp.WriteMsg(m)
		return
	}
	d.logger.Printf("[DEBUG] dns: recurse RTT for %v (%v)", q, rtt)

	// Forward the response
	if err := resp.WriteMsg(r); err != nil {
		d.logger.Printf("[WARN] dns: failed to respond: %v", err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Rook Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

	http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package k8sutil for Kubernetes helpers.
package k8sutil

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/pkg/capnslog"
	"github.com/rook/rook/pkg/clusterd"
	rookversion "github.com/rook/rook/pkg/version"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/validation"
	"k8s.io/apimachinery/pkg/util/version"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-k8sutil")

const (
	// Namespace for rook
	Namespace = "rook"
	// DefaultNamespace for the cluster
	DefaultNamespace = "default"
	// DataDirVolume data dir volume
	DataDirVolume = "rook-data"
	// DataDir folder
	DataDir = "/var/lib/rook"
	// RookType for the CRD
	RookType = "kubernetes.io/rook"
	// PodNameEnvVar is the env variable for getting the pod name via downward api
	PodNameEnvVar = "POD_NAME"
	// PodNamespaceEnvVar is the env variable for getting the pod namespace via downward api
	PodNamespaceEnvVar = "POD_NAMESPACE"
	// NodeNameEnvVar is the env variable for getting the node via downward api
	NodeNameEnvVar = "NODE_NAME"

	// RookVersionLabelKey is the key used for reporting the Rook version which last created or
	// modified a resource.
	RookVersionLabelKey = "rook-version"
)

// GetK8SVersion gets the version of the running K8S cluster
func GetK8SVersion(clientset kubernetes.Interface) (*version.Version, error) {
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		return nil, fmt.Errorf("error getting server version: %v", err)
	}

	// make sure the kubernetes version is parseable
	index := strings.Index(serverVersion.GitVersion, "+")
	if index != -1 {
		newVersion := serverVersion.GitVersion[:index]
		logger.Infof("returning version %s instead of %s", newVersion, serverVersion.GitVersion)
		serverVersion.GitVersion = newVersion
	}
	return version.MustParseSemantic(serverVersion.GitVersion), nil
}

// Hash stableName computes a stable pseudorandom string suitable for inclusion in a Kubernetes object name from the given seed string.
func Hash(s string) string {
	h := sha256.Sum256([]byte(s))
	return hex.EncodeToString(h[:16])
}

// TruncateNodeName hashes the nodeName in case it would case the name to be longer than 63 characters
// WARNING If your format and nodeName as a hash, are longer than 63 chars it won't be truncated!
// Your format alone should only be 31 chars at max because of MD5 hash being 32 chars.
// For more information, see the following resources:
// https://stackoverflow.com/a/50451893
// https://stackoverflow.com/a/32294443
func TruncateNodeName(format, nodeName string) string {
	if len(nodeName)+len(fmt.Sprintf(format, "")) > validation.DNS1035LabelMaxLength {
		hashed := Hash(nodeName)
		logger.Infof("format and nodeName longer than %d chars, nodeName %s will be %s", validation.DNS1035LabelMaxLength, nodeName, hashed)
		nodeName = hashed
	}
	return fmt.Sprintf(format, nodeName)
}

// deleteResourceAndWait will delete a resource, then wait for it to be purged from the system
func deleteResourceAndWait(namespace, name, resourceType string,
	deleteAction func(*metav1.DeleteOptions) error,
	getAction func() error,
) error {
	var gracePeriod int64
	propagation := metav1.DeletePropagationForeground
	options := &metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod, PropagationPolicy: &propagation}

	// Delete the resource if it exists
	logger.Infof("removing %s %s if it exists", resourceType, name)
	err := deleteAction(options)
	if err != nil {
		if !errors.IsNotFound(err) {
			return fmt.Errorf("failed to delete %s. %+v", name, err)
		}
		return nil
	}
	logger.Infof("Removed %s %s", resourceType, name)

	// wait for the resource to be deleted
	sleepTime := 2 * time.Second
	for i := 0; i < 45; i++ {
		// check for the existence of the resource
		err = getAction()
		if err != nil {
			if errors.IsNotFound(err) {
				logger.Infof("confirmed %s does not exist", name)
				return nil
			}
			return fmt.Errorf("failed to get %s. %+v", name, err)
		}

		if i%5 == 0 {
			// occasionally print a message
			logger.Infof("%q still found. waiting...", name)
		}
		time.Sleep(sleepTime)
	}

	return fmt.Errorf("gave up waiting for %s pods to be terminated", name)
}

// Add the rook version to the labels. This should *not* be used on pod specifications, because this
// will result in the deployment/daemonset/ect. recreating all of its pods even if an update
// wouldn't otherwise be required. Upgrading unnecessarily increases risk for loss of data
// reliability, even if only briefly.
//
// Note that the label may not match the version string exactly, since some characters used
// in version strings are illegal in pod labels.
func addRookVersionLabel(labels map[string]string) {
	value := validateLabelValue(rookversion.Version)
	labels[RookVersionLabelKey] = value
}

// validateLabelValue replaces any invalid characters
// in the input string with a replacement character,
// and enforces other limitations for k8s label values.
//
// See https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
func validateLabelValue(value string) string {
	repl := "-"
	maxlen := validation.LabelValueMaxLength
	re := regexp.MustCompile("[^a-z0-9A-Z._-]")
	// restrict label value to valid character set
	sanitized := re.ReplaceAllLiteralString(value, repl)
	// ensure value begins and ends with a valid character
	sanitized = strings.TrimRight(strings.TrimLeft(sanitized, ".-_"), ".-_")
	// limit length
	if len(sanitized) > maxlen {
		sanitized = sanitized[:maxlen]
	}
	return sanitized
}

// StartOperatorSettingsWatch starts the watch for Operator Settings ConfigMap
func StartOperatorSettingsWatch(context *clusterd.Context, operatorNamespace, operatorSettingConfigMapName string,
	addFunc func(obj interface{}), updateFunc func(oldObj, newObj interface{}), deleteFunc func(obj interface{}), stopCh chan struct{}) {
	_, cacheController := cache.NewInformer(cache.NewFilteredListWatchFromClient(context.Clientset.CoreV1().RESTClient(),
		"configmaps", operatorNamespace, func(options *metav1.ListOptions) {
			options.FieldSelector = fmt.Sprintf("%s=%s", "metadata.name", operatorSettingConfigMapName)
		}), &v1.ConfigMap{},
		0,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    addFunc,
			UpdateFunc: updateFunc,
			DeleteFunc: deleteFunc,
		})
	go cacheController.Run(stopCh)
}
<commit_msg>core: add comment warning to keep k8sutil hashing consistent<commit_after>/*
Copyright 2016 The Rook Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

	http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package k8sutil for Kubernetes helpers.
package k8sutil

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/coreos/pkg/capnslog"
	"github.com/rook/rook/pkg/clusterd"
	rookversion "github.com/rook/rook/pkg/version"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/validation"
	"k8s.io/apimachinery/pkg/util/version"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
)

var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-k8sutil")

const (
	// Namespace for rook
	Namespace = "rook"
	// DefaultNamespace for the cluster
	DefaultNamespace = "default"
	// DataDirVolume data dir volume
	DataDirVolume = "rook-data"
	// DataDir folder
	DataDir = "/var/lib/rook"
	// RookType for the CRD
	RookType = "kubernetes.io/rook"
	// PodNameEnvVar is the env variable for getting the pod name via downward api
	PodNameEnvVar = "POD_NAME"
	// PodNamespaceEnvVar is the env variable for getting the pod namespace via downward api
	PodNamespaceEnvVar = "POD_NAMESPACE"
	// NodeNameEnvVar is the env variable for getting the node via downward api
	NodeNameEnvVar = "NODE_NAME"

	// RookVersionLabelKey is the key used for reporting the Rook version which last created or
	// modified a resource.
	RookVersionLabelKey = "rook-version"
)

// GetK8SVersion gets the version of the running K8S cluster
func GetK8SVersion(clientset kubernetes.Interface) (*version.Version, error) {
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		return nil, fmt.Errorf("error getting server version: %v", err)
	}

	// make sure the kubernetes version is parseable
	index := strings.Index(serverVersion.GitVersion, "+")
	if index != -1 {
		newVersion := serverVersion.GitVersion[:index]
		logger.Infof("returning version %s instead of %s", newVersion, serverVersion.GitVersion)
		serverVersion.GitVersion = newVersion
	}
	return version.MustParseSemantic(serverVersion.GitVersion), nil
}

// Hash stableName computes a stable pseudorandom string suitable for inclusion in a Kubernetes object name from the given seed string.
// Do **NOT** edit this function in a way that would change its output as it needs to
// provide consistent mappings from string to hash accross versions of rook.
func Hash(s string) string {
	h := sha256.Sum256([]byte(s))
	return hex.EncodeToString(h[:16])
}

// TruncateNodeName hashes the nodeName in case it would case the name to be longer than 63 characters
// WARNING If your format and nodeName as a hash, are longer than 63 chars it won't be truncated!
// Your format alone should only be 31 chars at max because of MD5 hash being 32 chars.
// For more information, see the following resources:
// https://stackoverflow.com/a/50451893
// https://stackoverflow.com/a/32294443
// Do **NOT** edit this function in a way that would change its output as it needs to
// provide consistent mappings from string to hash accross versions of rook.
func TruncateNodeName(format, nodeName string) string {
	if len(nodeName)+len(fmt.Sprintf(format, "")) > validation.DNS1035LabelMaxLength {
		hashed := Hash(nodeName)
		logger.Infof("format and nodeName longer than %d chars, nodeName %s will be %s", validation.DNS1035LabelMaxLength, nodeName, hashed)
		nodeName = hashed
	}
	return fmt.Sprintf(format, nodeName)
}

// deleteResourceAndWait will delete a resource, then wait for it to be purged from the system
func deleteResourceAndWait(namespace, name, resourceType string,
	deleteAction func(*metav1.DeleteOptions) error,
	getAction func() error,
) error {
	var gracePeriod int64
	propagation := metav1.DeletePropagationForeground
	options := &metav1.DeleteOptions{GracePeriodSeconds: &gracePeriod, PropagationPolicy: &propagation}

	// Delete the resource if it exists
	logger.Infof("removing %s %s if it exists", resourceType, name)
	err := deleteAction(options)
	if err != nil {
		if !errors.IsNotFound(err) {
			return fmt.Errorf("failed to delete %s. %+v", name, err)
		}
		return nil
	}
	logger.Infof("Removed %s %s", resourceType, name)

	// wait for the resource to be deleted
	sleepTime := 2 * time.Second
	for i := 0; i < 45; i++ {
		// check for the existence of the resource
		err = getAction()
		if err != nil {
			if errors.IsNotFound(err) {
				logger.Infof("confirmed %s does not exist", name)
				return nil
			}
			return fmt.Errorf("failed to get %s. %+v", name, err)
		}

		if i%5 == 0 {
			// occasionally print a message
			logger.Infof("%q still found. waiting...", name)
		}
		time.Sleep(sleepTime)
	}

	return fmt.Errorf("gave up waiting for %s pods to be terminated", name)
}

// Add the rook version to the labels. This should *not* be used on pod specifications, because this
// will result in the deployment/daemonset/ect. recreating all of its pods even if an update
// wouldn't otherwise be required. Upgrading unnecessarily increases risk for loss of data
// reliability, even if only briefly.
//
// Note that the label may not match the version string exactly, since some characters used
// in version strings are illegal in pod labels.
func addRookVersionLabel(labels map[string]string) {
	value := validateLabelValue(rookversion.Version)
	labels[RookVersionLabelKey] = value
}

// validateLabelValue replaces any invalid characters
// in the input string with a replacement character,
// and enforces other limitations for k8s label values.
//
// See https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
func validateLabelValue(value string) string {
	repl := "-"
	maxlen := validation.LabelValueMaxLength
	re := regexp.MustCompile("[^a-z0-9A-Z._-]")
	// restrict label value to valid character set
	sanitized := re.ReplaceAllLiteralString(value, repl)
	// ensure value begins and ends with a valid character
	sanitized = strings.TrimRight(strings.TrimLeft(sanitized, ".-_"), ".-_")
	// limit length
	if len(sanitized) > maxlen {
		sanitized = sanitized[:maxlen]
	}
	return sanitized
}

// StartOperatorSettingsWatch starts the watch for Operator Settings ConfigMap
func StartOperatorSettingsWatch(context *clusterd.Context, operatorNamespace, operatorSettingConfigMapName string,
	addFunc func(obj interface{}), updateFunc func(oldObj, newObj interface{}), deleteFunc func(obj interface{}), stopCh chan struct{}) {
	_, cacheController := cache.NewInformer(cache.NewFilteredListWatchFromClient(context.Clientset.CoreV1().RESTClient(),
		"configmaps", operatorNamespace, func(options *metav1.ListOptions) {
			options.FieldSelector = fmt.Sprintf("%s=%s", "metadata.name", operatorSettingConfigMapName)
		}), &v1.ConfigMap{},
		0,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    addFunc,
			UpdateFunc: updateFunc,
			DeleteFunc: deleteFunc,
		})
	go cacheController.Run(stopCh)
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"url"
	"xml"
)

const (
	version   = "2.1"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 5, "pages to prime at once")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show additional information about the priming process")
	nowarn      *bool   = flag.Bool("no-warn", false, "do not warn about pages that can't be loaded")
)

type Sitemap struct {
	Loc string
	// Lastmod string
}

type Url struct {
	Loc string
	// Lastmod string
	// Changefreq string
	Priority float64
}

type Urlset struct {
	XMLName xml.Name "urlset"
	Sitemap []Sitemap
	Url     []Url
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (r *http.Response, err os.Error) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("User-Agent", useragent)
	if err != nil {
		return nil, err
	}
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, os.Error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    os.Error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		if *verbose {
			log.Println("Downloading", path)
		}
		res, err = Get(path)
		if res.Status != "200 OK" {
			return nil, os.NewError("Could not fetch sitemap: " + res.Status)
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		if *verbose {
			log.Println("Extracting compressed data")
		}
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, os.NewError("Gzip decompression failed")
		}
	}
	err = xml.Unmarshal(f, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan *Urlset, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding URLs from child sitemap %s\n", v)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
				} else {
					ch <- ourlset
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			urlset.Url = append(urlset.Url, childUrlset.Url...)
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		log.Println("URLs to prime:", urlset.Len())
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) os.Error {
	var (
		err   os.Error
		found bool = false
	)
	if *verbose {
		log.Printf("Get (weight %d) %s\n", int(u.Priority*100), u.Loc)
	}
	if *localDir != "" {
		parsed, err := url.ParseWithReference(u.Loc)
		joined := path.Join(*localDir, parsed.Path, *localSuffix)
		if _, err = os.Lstat(joined); err == nil {
			found = true
		}
	}
	if !found {
		res, err := Get(u.Loc)
		if (err != nil || res.Status != "200 OK") && !*nowarn {
			var errmsg string
			if err != nil {
				errmsg = err.String()
			} else {
				errmsg = res.Status
			}
			log.Printf("Error priming %s: %s\n", u.Loc, errmsg)
		}
	}
	wg.Done()
	<-sem
	return err
}

func main() {
	var (
		urlset *Urlset
		err    os.Error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		PrimeUrlset(urlset)
	}
}
<commit_msg>Increment version; revert concurrency to 1--let user decide; additional terminology<commit_after>package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"http"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
	"url"
	"xml"
)

const (
	version   = "2.2"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 1, "URLs to prime at once")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show additional information about the priming process")
	nowarn      *bool   = flag.Bool("no-warn", false, "do not warn about pages that were not primed successfully")
)

type Sitemap struct {
	Loc string
	// Lastmod string
}

type Url struct {
	Loc string
	// Lastmod string
	// Changefreq string
	Priority float64
}

type Urlset struct {
	XMLName xml.Name "urlset"
	Sitemap []Sitemap
	Url     []Url
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (r *http.Response, err os.Error) {
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("User-Agent", useragent)
	if err != nil {
		return nil, err
	}
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, os.Error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    os.Error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		if *verbose {
			log.Println("Downloading", path)
		}
		res, err = Get(path)
		if res.Status != "200 OK" {
			return nil, os.NewError("Could not fetch sitemap: " + res.Status)
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		if *verbose {
			log.Println("Extracting compressed data")
		}
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, os.NewError("Gzip decompression failed")
		}
	}
	err = xml.Unmarshal(f, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan *Urlset, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding URLs from child sitemap %s\n", v)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
				} else {
					ch <- ourlset
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			urlset.Url = append(urlset.Url, childUrlset.Url...)
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		log.Println("URLs to prime:", urlset.Len())
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) os.Error {
	var (
		err   os.Error
		found bool = false
	)
	if *verbose {
		log.Printf("Get (weight %d) %s\n", int(u.Priority*100), u.Loc)
	}
	if *localDir != "" {
		parsed, err := url.ParseWithReference(u.Loc)
		joined := path.Join(*localDir, parsed.Path, *localSuffix)
		if _, err = os.Lstat(joined); err == nil {
			found = true
		}
	}
	if !found {
		res, err := Get(u.Loc)
		if (err != nil || res.Status != "200 OK") && !*nowarn {
			var errmsg string
			if err != nil {
				errmsg = err.String()
			} else {
				errmsg = res.Status
			}
			log.Printf("Error priming %s: %s\n", u.Loc, errmsg)
		}
	}
	wg.Done()
	<-sem
	return err
}

func main() {
	var (
		urlset *Urlset
		err    os.Error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		PrimeUrlset(urlset)
	}
}
<|endoftext|>