<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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 cmd

import (
	"fmt"
	"io"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/certificates"
	"k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/printers"

	"github.com/spf13/cobra"
)

func NewCmdCertificate(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	cmd := &cobra.Command{
		Use: "certificate SUBCOMMAND",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Modify certificate resources."),
		Long:  "Modify certificate resources.",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(NewCmdCertificateApprove(f, ioStreams))
	cmd.AddCommand(NewCmdCertificateDeny(f, ioStreams))

	return cmd
}

type CertificateOptions struct {
	resource.FilenameOptions

	PrintFlags *printers.PrintFlags
	PrintObj   printers.ResourcePrinterFunc

	csrNames    []string
	outputStyle string

	clientSet internalclientset.Interface
	builder   *resource.Builder

	genericclioptions.IOStreams
}

func (o *CertificateOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	o.csrNames = args
	o.outputStyle = cmdutil.GetFlagString(cmd, "output")

	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}

	o.PrintObj = func(obj runtime.Object, out io.Writer) error {
		return printer.PrintObj(obj, out)
	}

	o.builder = f.NewBuilder()
	o.clientSet, err = f.ClientSet()
	if err != nil {
		return err
	}

	return nil
}

func (o *CertificateOptions) Validate() error {
	if len(o.csrNames) < 1 && cmdutil.IsFilenameSliceEmpty(o.Filenames) {
		return fmt.Errorf("one or more CSRs must be specified as <name> or -f <filename>")
	}
	return nil
}

func NewCmdCertificateApprove(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	options := CertificateOptions{
		PrintFlags: printers.NewPrintFlags("approved"),
		IOStreams:  ioStreams,
	}
	cmd := &cobra.Command{
		Use: "approve (-f FILENAME | NAME)",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Approve a certificate signing request"),
		Long: templates.LongDesc(`
		Approve a certificate signing request.

		kubectl certificate approve allows a cluster admin to approve a certificate
		signing request (CSR). This action tells a certificate signing controller to
		issue a certificate to the requestor with the attributes requested in the CSR.

		SECURITY NOTICE: Depending on the requested attributes, the issued certificate
		can potentially grant a requester access to cluster resources or to authenticate
		as a requested identity. Before approving a CSR, ensure you understand what the
		signed certificate can do.
		`),
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.RunCertificateApprove(cmdutil.GetFlagBool(cmd, "force")))
		},
	}
	cmd.Flags().Bool("force", false, "Update the CSR even if it is already approved.")
	cmdutil.AddOutputFlagsForMutation(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &options.FilenameOptions, "identifying the resource to update")

	return cmd
}

func (o *CertificateOptions) RunCertificateApprove(force bool) error {
	return o.modifyCertificateCondition(o.builder, o.clientSet, force, func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool) {
		var alreadyApproved bool
		for _, c := range csr.Status.Conditions {
			if c.Type == certificates.CertificateApproved {
				alreadyApproved = true
			}
		}
		if alreadyApproved {
			return csr, true
		}
		csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
			Type:           certificates.CertificateApproved,
			Reason:         "KubectlApprove",
			Message:        "This CSR was approved by kubectl certificate approve.",
			LastUpdateTime: metav1.Now(),
		})
		return csr, false
	})
}

func NewCmdCertificateDeny(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	options := CertificateOptions{
		PrintFlags: printers.NewPrintFlags("denied"),
		IOStreams:  ioStreams,
	}
	cmd := &cobra.Command{
		Use: "deny (-f FILENAME | NAME)",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Deny a certificate signing request"),
		Long: templates.LongDesc(`
		Deny a certificate signing request.

		kubectl certificate deny allows a cluster admin to deny a certificate
		signing request (CSR). This action tells a certificate signing controller to
		not to issue a certificate to the requestor.
		`),
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.RunCertificateDeny(cmdutil.GetFlagBool(cmd, "force")))
		},
	}
	cmd.Flags().Bool("force", false, "Update the CSR even if it is already denied.")
	cmdutil.AddOutputFlagsForMutation(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &options.FilenameOptions, "identifying the resource to update")

	return cmd
}

func (o *CertificateOptions) RunCertificateDeny(force bool) error {
	return o.modifyCertificateCondition(o.builder, o.clientSet, force, func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool) {
		var alreadyDenied bool
		for _, c := range csr.Status.Conditions {
			if c.Type == certificates.CertificateDenied {
				alreadyDenied = true
			}
		}
		if alreadyDenied {
			return csr, true
		}
		csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
			Type:           certificates.CertificateDenied,
			Reason:         "KubectlDeny",
			Message:        "This CSR was approved by kubectl certificate deny.",
			LastUpdateTime: metav1.Now(),
		})
		return csr, false
	})
}

func (options *CertificateOptions) modifyCertificateCondition(builder *resource.Builder, clientSet internalclientset.Interface, force bool, modify func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool)) error {
	var found int
	r := builder.
		Internal(legacyscheme.Scheme).
		ContinueOnError().
		FilenameParam(false, &options.FilenameOptions).
		ResourceNames("certificatesigningrequest", options.csrNames...).
		RequireObject(true).
		Flatten().
		Latest().
		Do()
	err := r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		csr := info.Object.(*certificates.CertificateSigningRequest)
		csr, hasCondition := modify(csr)
		if !hasCondition || force {
			csr, err = clientSet.Certificates().
				CertificateSigningRequests().
				UpdateApproval(csr)
			if err != nil {
				return err
			}
		}
		found++

		return options.PrintObj(cmdutil.AsDefaultVersionedOrOriginal(info.Object, info.Mapping), options.Out)
	})
	if found == 0 {
		fmt.Fprintf(options.Out, "No resources found\n")
	}
	return err
}
<commit_msg>Retry certificate approval on conflict errors<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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 cmd

import (
	"fmt"
	"io"

	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/certificates"
	"k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/kubectl/resource"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/printers"

	"github.com/spf13/cobra"
)

func NewCmdCertificate(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	cmd := &cobra.Command{
		Use: "certificate SUBCOMMAND",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Modify certificate resources."),
		Long:  "Modify certificate resources.",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	cmd.AddCommand(NewCmdCertificateApprove(f, ioStreams))
	cmd.AddCommand(NewCmdCertificateDeny(f, ioStreams))

	return cmd
}

type CertificateOptions struct {
	resource.FilenameOptions

	PrintFlags *printers.PrintFlags
	PrintObj   printers.ResourcePrinterFunc

	csrNames    []string
	outputStyle string

	clientSet internalclientset.Interface
	builder   *resource.Builder

	genericclioptions.IOStreams
}

func (o *CertificateOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	o.csrNames = args
	o.outputStyle = cmdutil.GetFlagString(cmd, "output")

	printer, err := o.PrintFlags.ToPrinter()
	if err != nil {
		return err
	}

	o.PrintObj = func(obj runtime.Object, out io.Writer) error {
		return printer.PrintObj(obj, out)
	}

	o.builder = f.NewBuilder()
	o.clientSet, err = f.ClientSet()
	if err != nil {
		return err
	}

	return nil
}

func (o *CertificateOptions) Validate() error {
	if len(o.csrNames) < 1 && cmdutil.IsFilenameSliceEmpty(o.Filenames) {
		return fmt.Errorf("one or more CSRs must be specified as <name> or -f <filename>")
	}
	return nil
}

func NewCmdCertificateApprove(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	options := CertificateOptions{
		PrintFlags: printers.NewPrintFlags("approved"),
		IOStreams:  ioStreams,
	}
	cmd := &cobra.Command{
		Use: "approve (-f FILENAME | NAME)",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Approve a certificate signing request"),
		Long: templates.LongDesc(`
		Approve a certificate signing request.

		kubectl certificate approve allows a cluster admin to approve a certificate
		signing request (CSR). This action tells a certificate signing controller to
		issue a certificate to the requestor with the attributes requested in the CSR.

		SECURITY NOTICE: Depending on the requested attributes, the issued certificate
		can potentially grant a requester access to cluster resources or to authenticate
		as a requested identity. Before approving a CSR, ensure you understand what the
		signed certificate can do.
		`),
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.RunCertificateApprove(cmdutil.GetFlagBool(cmd, "force")))
		},
	}
	cmd.Flags().Bool("force", false, "Update the CSR even if it is already approved.")
	cmdutil.AddOutputFlagsForMutation(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &options.FilenameOptions, "identifying the resource to update")

	return cmd
}

func (o *CertificateOptions) RunCertificateApprove(force bool) error {
	return o.modifyCertificateCondition(o.builder, o.clientSet, force, func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool) {
		var alreadyApproved bool
		for _, c := range csr.Status.Conditions {
			if c.Type == certificates.CertificateApproved {
				alreadyApproved = true
			}
		}
		if alreadyApproved {
			return csr, true
		}
		csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
			Type:           certificates.CertificateApproved,
			Reason:         "KubectlApprove",
			Message:        "This CSR was approved by kubectl certificate approve.",
			LastUpdateTime: metav1.Now(),
		})
		return csr, false
	})
}

func NewCmdCertificateDeny(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	options := CertificateOptions{
		PrintFlags: printers.NewPrintFlags("denied"),
		IOStreams:  ioStreams,
	}
	cmd := &cobra.Command{
		Use: "deny (-f FILENAME | NAME)",
		DisableFlagsInUseLine: true,
		Short: i18n.T("Deny a certificate signing request"),
		Long: templates.LongDesc(`
		Deny a certificate signing request.

		kubectl certificate deny allows a cluster admin to deny a certificate
		signing request (CSR). This action tells a certificate signing controller to
		not to issue a certificate to the requestor.
		`),
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.RunCertificateDeny(cmdutil.GetFlagBool(cmd, "force")))
		},
	}
	cmd.Flags().Bool("force", false, "Update the CSR even if it is already denied.")
	cmdutil.AddOutputFlagsForMutation(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, &options.FilenameOptions, "identifying the resource to update")

	return cmd
}

func (o *CertificateOptions) RunCertificateDeny(force bool) error {
	return o.modifyCertificateCondition(o.builder, o.clientSet, force, func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool) {
		var alreadyDenied bool
		for _, c := range csr.Status.Conditions {
			if c.Type == certificates.CertificateDenied {
				alreadyDenied = true
			}
		}
		if alreadyDenied {
			return csr, true
		}
		csr.Status.Conditions = append(csr.Status.Conditions, certificates.CertificateSigningRequestCondition{
			Type:           certificates.CertificateDenied,
			Reason:         "KubectlDeny",
			Message:        "This CSR was approved by kubectl certificate deny.",
			LastUpdateTime: metav1.Now(),
		})
		return csr, false
	})
}

func (options *CertificateOptions) modifyCertificateCondition(builder *resource.Builder, clientSet internalclientset.Interface, force bool, modify func(csr *certificates.CertificateSigningRequest) (*certificates.CertificateSigningRequest, bool)) error {
	var found int
	r := builder.
		Internal(legacyscheme.Scheme).
		ContinueOnError().
		FilenameParam(false, &options.FilenameOptions).
		ResourceNames("certificatesigningrequest", options.csrNames...).
		RequireObject(true).
		Flatten().
		Latest().
		Do()
	err := r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}
		for i := 0; ; i++ {
			csr := info.Object.(*certificates.CertificateSigningRequest)
			csr, hasCondition := modify(csr)
			if !hasCondition || force {
				csr, err = clientSet.Certificates().
					CertificateSigningRequests().
					UpdateApproval(csr)
				if errors.IsConflict(err) && i < 10 {
					if err := info.Get(); err != nil {
						return err
					}
					continue
				}
				if err != nil {
					return err
				}
			}
			break
		}
		found++

		return options.PrintObj(cmdutil.AsDefaultVersionedOrOriginal(info.Object, info.Mapping), options.Out)
	})
	if found == 0 {
		fmt.Fprintf(options.Out, "No resources found\n")
	}
	return err
}
<|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"
	"strconv"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/latest"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	utilerrors "github.com/GoogleCloudPlatform/kubernetes/pkg/util/errors"
	"github.com/evanphx/json-patch"

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

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

// 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) {
	if err == nil {
		return
	}

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

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

// 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 {
	f := getFlag(cmd, flag)
	return f.Value.String()
}

func GetFlagBool(cmd *cobra.Command, flag string) bool {
	f := getFlag(cmd, flag)
	result, err := strconv.ParseBool(f.Value.String())
	if err != nil {
		glog.Fatalf("Invalid value for a boolean flag: %s", f.Value.String())
	}
	return result
}

// Assumes the flag has a default value.
func GetFlagInt(cmd *cobra.Command, flag string) int {
	f := getFlag(cmd, flag)
	v, err := strconv.Atoi(f.Value.String())
	// This is likely not a sufficiently friendly error message, but cobra
	// should prevent non-integer values from reaching here.
	if err != nil {
		glog.Fatalf("unable to convert flag value to int: %v", err)
	}
	return v
}

func GetFlagDuration(cmd *cobra.Command, flag string) time.Duration {
	f := getFlag(cmd, flag)
	v, err := time.ParseDuration(f.Value.String())
	if err != nil {
		glog.Fatalf("unable to convert flag value to Duration: %v", err)
	}
	return v
}

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
}
<commit_msg>Print multiple validation errors<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"
	"strconv"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/latest"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	utilerrors "github.com/GoogleCloudPlatform/kubernetes/pkg/util/errors"
	"github.com/evanphx/json-patch"

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

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

// 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) {
	if err == nil {
		return
	}

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

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

	msg, ok := StandardErrorMessage(err)
	if !ok {
		msg = fmt.Sprintf("error: %s\n", err.Error())
	}
	fatal(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 {
	f := getFlag(cmd, flag)
	return f.Value.String()
}

func GetFlagBool(cmd *cobra.Command, flag string) bool {
	f := getFlag(cmd, flag)
	result, err := strconv.ParseBool(f.Value.String())
	if err != nil {
		glog.Fatalf("Invalid value for a boolean flag: %s", f.Value.String())
	}
	return result
}

// Assumes the flag has a default value.
func GetFlagInt(cmd *cobra.Command, flag string) int {
	f := getFlag(cmd, flag)
	v, err := strconv.Atoi(f.Value.String())
	// This is likely not a sufficiently friendly error message, but cobra
	// should prevent non-integer values from reaching here.
	if err != nil {
		glog.Fatalf("unable to convert flag value to int: %v", err)
	}
	return v
}

func GetFlagDuration(cmd *cobra.Command, flag string) time.Duration {
	f := getFlag(cmd, flag)
	v, err := time.ParseDuration(f.Value.String())
	if err != nil {
		glog.Fatalf("unable to convert flag value to Duration: %v", err)
	}
	return v
}

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
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package reader

import (
	"bufio"
	"bytes"
	"io"
	"time"

	"github.com/Jeffail/benthos/v3/lib/message"
	"github.com/Jeffail/benthos/v3/lib/types"
)

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

// Lines is a reader implementation that continuously reads line delimited
// messages from an io.Reader type.
type Lines struct {
	handleCtor func() (io.Reader, error)
	onClose    func()

	handle  io.Reader
	scanner *bufio.Scanner

	messageBuffer      *bytes.Buffer
	messageBufferIndex int

	maxBuffer int
	multipart bool
	delimiter []byte
}

// NewLines creates a new reader input type.
//
// Callers must provide a constructor function for the target io.Reader, which
// is called on start up and again each time a reader is exhausted. If the
// constructor is called but there is no more content to create a Reader for
// then the error `io.EOF` should be returned and the Lines will close.
//
// Callers must also provide an onClose function, which will be called if the
// Lines has been instructed to shut down. This function should unblock any
// blocked Read calls.
func NewLines(
	handleCtor func() (io.Reader, error),
	onClose func(),
	options ...func(r *Lines),
) (*Lines, error) {
	r := Lines{
		handleCtor:    handleCtor,
		onClose:       onClose,
		messageBuffer: &bytes.Buffer{},
		maxBuffer:     bufio.MaxScanTokenSize,
		multipart:     false,
		delimiter:     []byte("\n"),
	}

	for _, opt := range options {
		opt(&r)
	}

	return &r, nil
}

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

// OptLinesSetMaxBuffer is a option func that sets the maximum size of the
// line parsing buffers.
func OptLinesSetMaxBuffer(maxBuffer int) func(r *Lines) {
	return func(r *Lines) {
		r.maxBuffer = maxBuffer
	}
}

// OptLinesSetMultipart is a option func that sets the boolean flag
// indicating whether lines should be parsed as multipart or not.
func OptLinesSetMultipart(multipart bool) func(r *Lines) {
	return func(r *Lines) {
		r.multipart = multipart
	}
}

// OptLinesSetDelimiter is a option func that sets the delimiter (default
// '\n') used to divide lines (message parts) in the stream of data.
func OptLinesSetDelimiter(delimiter string) func(r *Lines) {
	return func(r *Lines) {
		r.delimiter = []byte(delimiter)
	}
}

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

func (r *Lines) closeHandle() {
	if r.handle != nil {
		if closer, ok := r.handle.(io.ReadCloser); ok {
			closer.Close()
		}
		r.handle = nil
	}
	r.scanner = nil
}

// Connect attempts to establish a new scanner for an io.Reader.
func (r *Lines) Connect() error {
	if r.scanner != nil {
		return nil
	}
	r.closeHandle() // Just incase we have an open handle without a scanner.

	var err error
	r.handle, err = r.handleCtor()
	if err != nil {
		if err == io.EOF {
			return types.ErrTypeClosed
		}
		return err
	}

	r.scanner = bufio.NewScanner(r.handle)
	if r.maxBuffer != bufio.MaxScanTokenSize {
		r.scanner.Buffer([]byte{}, r.maxBuffer)
	}

	r.scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		if atEOF && len(data) == 0 {
			return 0, nil, nil
		}

		if i := bytes.Index(data, r.delimiter); i >= 0 {
			// We have a full terminated line.
			return i + len(r.delimiter), data[0:i], nil
		}

		// If we're at EOF, we have a final, non-terminated line. Return it.
		if atEOF {
			return len(data), data, nil
		}

		// Request more data.
		return 0, nil, nil
	})

	return nil
}

// Read attempts to read a new line from the io.Reader.
func (r *Lines) Read() (types.Message, error) {
	if r.scanner == nil {
		return nil, types.ErrNotConnected
	}

	msg := message.New(nil)

	for r.scanner.Scan() {
		partSize, err := r.messageBuffer.Write(r.scanner.Bytes())
		rIndex := r.messageBufferIndex
		r.messageBufferIndex += partSize
		if err != nil {
			return nil, err
		}

		// WARNING: According to https://golang.org/pkg/bytes/#Buffer.Bytes the
		// slice returned by Bytes is only correct until the next call to Write.
		// Since we call Write for multiple part messages, and could potentially
		// call it on a consecutive Read call before the next Acknowledge, we
		// are passing slices through messages that are "invalid".
		//
		// However, in practice the calls to Write do not overwrite the memory
		// within the returned slice even if it results in the buffer
		// re-allocating memory. Since we also ensure that Reset is only called
		// once messages are no longer in use we should be fine here.
		//
		// Regardless, we should regularly revisit this code in order to ensure
		// that this remains the case. I can't foresee any case where discarded
		// slices within bytes.Buffer would be wiped or modified during Write,
		// but since the library does not guarantee this:
		//
		// TODO: Have another cheeky gander at
		// https://golang.org/src/bytes/buffer.go to make sure Write never
		// mutates a discarded slice during re-allocation. If it does then we
		// should stop using bytes.Buffer and either eat the allocations or do
		// some buffer rotations of our own.
		if partSize > 0 {
			msg.Append(message.NewPart(r.messageBuffer.Bytes()[rIndex : rIndex+partSize : rIndex+partSize]))
			if !r.multipart {
				return msg, nil
			}
		} else if r.multipart && msg.Len() > 0 {
			// Empty line means we're finished reading parts for this
			// message.
			return msg, nil
		}
	}

	if err := r.scanner.Err(); err != nil {
		r.closeHandle()
		return nil, err
	}

	r.closeHandle()

	if msg.Len() > 0 {
		return msg, nil
	}
	return nil, types.ErrNotConnected
}

// Acknowledge confirms whether or not our unacknowledged messages have been
// successfully propagated or not.
func (r *Lines) Acknowledge(err error) error {
	if err == nil && r.messageBuffer != nil {
		r.messageBuffer.Reset()
		r.messageBufferIndex = 0
	}
	return nil
}

// CloseAsync shuts down the reader input and stops processing requests.
func (r *Lines) CloseAsync() {
	r.onClose()
}

// WaitForClose blocks until the reader input has closed down.
func (r *Lines) WaitForClose(timeout time.Duration) error {
	r.closeHandle()
	return nil
}

//------------------------------------------------------------------------------
<commit_msg>Add context methods to lines reader<commit_after>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package reader

import (
	"bufio"
	"bytes"
	"context"
	"io"
	"time"

	"github.com/Jeffail/benthos/v3/lib/message"
	"github.com/Jeffail/benthos/v3/lib/types"
)

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

// Lines is a reader implementation that continuously reads line delimited
// messages from an io.Reader type.
type Lines struct {
	handleCtor func() (io.Reader, error)
	onClose    func()

	handle  io.Reader
	scanner *bufio.Scanner

	messageBuffer      *bytes.Buffer
	messageBufferIndex int

	maxBuffer int
	multipart bool
	delimiter []byte
}

// NewLines creates a new reader input type.
//
// Callers must provide a constructor function for the target io.Reader, which
// is called on start up and again each time a reader is exhausted. If the
// constructor is called but there is no more content to create a Reader for
// then the error `io.EOF` should be returned and the Lines will close.
//
// Callers must also provide an onClose function, which will be called if the
// Lines has been instructed to shut down. This function should unblock any
// blocked Read calls.
func NewLines(
	handleCtor func() (io.Reader, error),
	onClose func(),
	options ...func(r *Lines),
) (*Lines, error) {
	r := Lines{
		handleCtor:    handleCtor,
		onClose:       onClose,
		messageBuffer: &bytes.Buffer{},
		maxBuffer:     bufio.MaxScanTokenSize,
		multipart:     false,
		delimiter:     []byte("\n"),
	}

	for _, opt := range options {
		opt(&r)
	}

	return &r, nil
}

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

// OptLinesSetMaxBuffer is a option func that sets the maximum size of the
// line parsing buffers.
func OptLinesSetMaxBuffer(maxBuffer int) func(r *Lines) {
	return func(r *Lines) {
		r.maxBuffer = maxBuffer
	}
}

// OptLinesSetMultipart is a option func that sets the boolean flag
// indicating whether lines should be parsed as multipart or not.
func OptLinesSetMultipart(multipart bool) func(r *Lines) {
	return func(r *Lines) {
		r.multipart = multipart
	}
}

// OptLinesSetDelimiter is a option func that sets the delimiter (default
// '\n') used to divide lines (message parts) in the stream of data.
func OptLinesSetDelimiter(delimiter string) func(r *Lines) {
	return func(r *Lines) {
		r.delimiter = []byte(delimiter)
	}
}

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

func (r *Lines) closeHandle() {
	if r.handle != nil {
		if closer, ok := r.handle.(io.ReadCloser); ok {
			closer.Close()
		}
		r.handle = nil
	}
	r.scanner = nil
}

// Connect attempts to establish a new scanner for an io.Reader.
func (r *Lines) Connect() error {
	return r.ConnectWithContext(context.Background())
}

// ConnectWithContext attempts to establish a new scanner for an io.Reader.
func (r *Lines) ConnectWithContext(ctx context.Context) error {
	if r.scanner != nil {
		return nil
	}
	r.closeHandle() // Just incase we have an open handle without a scanner.

	var err error
	r.handle, err = r.handleCtor()
	if err != nil {
		if err == io.EOF {
			return types.ErrTypeClosed
		}
		return err
	}

	r.scanner = bufio.NewScanner(r.handle)
	if r.maxBuffer != bufio.MaxScanTokenSize {
		r.scanner.Buffer([]byte{}, r.maxBuffer)
	}

	r.scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		if atEOF && len(data) == 0 {
			return 0, nil, nil
		}

		if i := bytes.Index(data, r.delimiter); i >= 0 {
			// We have a full terminated line.
			return i + len(r.delimiter), data[0:i], nil
		}

		// If we're at EOF, we have a final, non-terminated line. Return it.
		if atEOF {
			return len(data), data, nil
		}

		// Request more data.
		return 0, nil, nil
	})

	return nil
}

// ReadWithContext attempts to read a new line from the io.Reader.
func (r *Lines) ReadWithContext(ctx context.Context) (types.Message, AsyncAckFn, error) {
	if r.scanner == nil {
		return nil, nil, types.ErrNotConnected
	}

	msg := message.New(nil)

	for r.scanner.Scan() {
		partBytes := make([]byte, len(r.scanner.Bytes()))
		partSize := copy(partBytes, r.scanner.Bytes())

		if partSize > 0 {
			msg.Append(message.NewPart(partBytes))
			if !r.multipart {
				return msg, noopAsyncAckFn, nil
			}
		} else if r.multipart && msg.Len() > 0 {
			// Empty line means we're finished reading parts for this
			// message.
			return msg, noopAsyncAckFn, nil
		}
	}

	if err := r.scanner.Err(); err != nil {
		r.closeHandle()
		return nil, nil, err
	}

	r.closeHandle()

	if msg.Len() > 0 {
		return msg, noopAsyncAckFn, nil
	}
	return nil, nil, types.ErrNotConnected
}

// Read attempts to read a new line from the io.Reader.
func (r *Lines) Read() (types.Message, error) {
	if r.scanner == nil {
		return nil, types.ErrNotConnected
	}

	msg := message.New(nil)

	for r.scanner.Scan() {
		partSize, err := r.messageBuffer.Write(r.scanner.Bytes())
		rIndex := r.messageBufferIndex
		r.messageBufferIndex += partSize
		if err != nil {
			return nil, err
		}

		// WARNING: According to https://golang.org/pkg/bytes/#Buffer.Bytes the
		// slice returned by Bytes is only correct until the next call to Write.
		// Since we call Write for multiple part messages, and could potentially
		// call it on a consecutive Read call before the next Acknowledge, we
		// are passing slices through messages that are "invalid".
		//
		// However, in practice the calls to Write do not overwrite the memory
		// within the returned slice even if it results in the buffer
		// re-allocating memory. Since we also ensure that Reset is only called
		// once messages are no longer in use we should be fine here.
		//
		// Regardless, we should regularly revisit this code in order to ensure
		// that this remains the case. I can't foresee any case where discarded
		// slices within bytes.Buffer would be wiped or modified during Write,
		// but since the library does not guarantee this:
		//
		// TODO: Have another cheeky gander at
		// https://golang.org/src/bytes/buffer.go to make sure Write never
		// mutates a discarded slice during re-allocation. If it does then we
		// should stop using bytes.Buffer and either eat the allocations or do
		// some buffer rotations of our own.
		if partSize > 0 {
			msg.Append(message.NewPart(r.messageBuffer.Bytes()[rIndex : rIndex+partSize : rIndex+partSize]))
			if !r.multipart {
				return msg, nil
			}
		} else if r.multipart && msg.Len() > 0 {
			// Empty line means we're finished reading parts for this
			// message.
			return msg, nil
		}
	}

	if err := r.scanner.Err(); err != nil {
		r.closeHandle()
		return nil, err
	}

	r.closeHandle()

	if msg.Len() > 0 {
		return msg, nil
	}
	return nil, types.ErrNotConnected
}

// Acknowledge confirms whether or not our unacknowledged messages have been
// successfully propagated or not.
func (r *Lines) Acknowledge(err error) error {
	if err == nil && r.messageBuffer != nil {
		r.messageBuffer.Reset()
		r.messageBufferIndex = 0
	}
	return nil
}

// CloseAsync shuts down the reader input and stops processing requests.
func (r *Lines) CloseAsync() {
	r.onClose()
}

// WaitForClose blocks until the reader input has closed down.
func (r *Lines) WaitForClose(timeout time.Duration) error {
	r.closeHandle()
	return nil
}

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>// Command safesql is a tool for performing static analysis on programs to
// ensure that SQL injection attacks are not possible. It does this by ensuring
// package database/sql is only used with compile-time constant queries.
package main

import (
	"flag"
	"fmt"
	"os"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
	"golang.org/x/tools/go/types"
)

func main() {
	var verbose, quiet bool
	flag.BoolVar(&verbose, "v", false, "Verbose mode")
	flag.BoolVar(&quiet, "q", false, "Only print on failure")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-q] [-v] package1 [package2 ...]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()
	pkgs := flag.Args()
	if len(pkgs) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	c := loader.Config{}
	c.Import("database/sql")
	for _, pkg := range pkgs {
		c.Import(pkg)
	}
	p, err := c.Load()
	if err != nil {
		fmt.Printf("error loading packages %v: %v\n", pkgs, err)
		os.Exit(2)
	}
	s := ssautil.CreateProgram(p, 0)
	s.BuildAll()

	qms := FindQueryMethods(p.Package("database/sql").Pkg, s)
	if verbose {
		fmt.Println("database/sql functions that accept queries:")
		for _, m := range qms {
			fmt.Printf("- %s (param %d)\n", m.Func, m.Param)
		}
		fmt.Println()
	}

	mains := FindMains(p, s)
	if len(mains) == 0 {
		fmt.Println("Did not find any commands (i.e., main functions).")
		os.Exit(2)
	}

	res, err := pointer.Analyze(&pointer.Config{
		Mains:          mains,
		BuildCallGraph: true,
	})
	if err != nil {
		fmt.Printf("error performing pointer analysis: %v\n", err)
		os.Exit(2)
	}

	bad := FindNonConstCalls(res.CallGraph, qms)
	if len(bad) == 0 {
		if !quiet {
			fmt.Println(`You're safe from SQL injection! Yay \o/`)
		}
		return
	}

	fmt.Printf("Found %d potentially unsafe SQL statements:\n", len(bad))
	for _, ci := range bad {
		pos := p.Fset.Position(ci.Pos())
		fmt.Printf("- %s\n", pos)
	}
	fmt.Println("Please ensure that all SQL queries you use are compile-time constants.")
	fmt.Println("You should always use parameterized queries or prepared statements")
	fmt.Println("instead of building queries from strings.")
	os.Exit(1)
}

// QueryMethod represents a method on a type which has a string parameter named
// "query".
type QueryMethod struct {
	Func     *types.Func
	SSA      *ssa.Function
	ArgCount int
	Param    int
}

// FindQueryMethods locates all methods in the given package (assumed to be
// package database/sql) with a string parameter named "query".
func FindQueryMethods(sql *types.Package, ssa *ssa.Program) []*QueryMethod {
	methods := make([]*QueryMethod, 0)
	scope := sql.Scope()
	for _, name := range scope.Names() {
		o := scope.Lookup(name)
		if !o.Exported() {
			continue
		}
		if _, ok := o.(*types.TypeName); !ok {
			continue
		}
		n := o.Type().(*types.Named)
		for i := 0; i < n.NumMethods(); i++ {
			m := n.Method(i)
			if !m.Exported() {
				continue
			}
			s := m.Type().(*types.Signature)
			if num, ok := FuncHasQuery(s); ok {
				methods = append(methods, &QueryMethod{
					Func:     m,
					SSA:      ssa.FuncValue(m),
					ArgCount: s.Params().Len(),
					Param:    num,
				})
			}
		}
	}
	return methods
}

var stringType types.Type = types.Typ[types.String]

// FuncHasQuery returns the offset of the string parameter named "query", or
// none if no such parameter exists.
func FuncHasQuery(s *types.Signature) (offset int, ok bool) {
	params := s.Params()
	for i := 0; i < params.Len(); i++ {
		v := params.At(i)
		if v.Name() == "query" && v.Type() == stringType {
			return i, true
		}
	}
	return 0, false
}

// FindMains returns the set of all packages loaded into the given
// loader.Program which contain main functions
func FindMains(p *loader.Program, s *ssa.Program) []*ssa.Package {
	ips := p.InitialPackages()
	mains := make([]*ssa.Package, 0, len(ips))
	for _, info := range ips {
		ssaPkg := s.Package(info.Pkg)
		if ssaPkg.Func("main") != nil {
			mains = append(mains, ssaPkg)
		}
	}
	return mains
}

// FindNonConstCalls returns the set of callsites of the given set of methods
// for which the "query" parameter is not a compile-time constant.
func FindNonConstCalls(cg *callgraph.Graph, qms []*QueryMethod) []ssa.CallInstruction {
	cg.DeleteSyntheticNodes()

	// package database/sql has a couple helper functions which are thin
	// wrappers around other sensitive functions. Instead of handling the
	// general case by tracing down callsites of wrapper functions
	// recursively, let's just whitelist the functions we're already
	// tracking, since it happens to be good enough for our use case.
	okFuncs := make(map[*ssa.Function]struct{}, len(qms))
	for _, m := range qms {
		okFuncs[m.SSA] = struct{}{}
	}

	bad := make([]ssa.CallInstruction, 0)
	for _, m := range qms {
		node := cg.CreateNode(m.SSA)
		for _, edge := range node.In {
			if _, ok := okFuncs[edge.Site.Parent()]; ok {
				continue
			}
			cc := edge.Site.Common()
			args := cc.Args
			// The first parameter is occasionally the receiver.
			if len(args) == m.ArgCount+1 {
				args = args[1:]
			} else if len(args) != m.ArgCount {
				panic("arg count mismatch")
			}
			v := args[m.Param]
			if _, ok := v.(*ssa.Const); !ok {
				bad = append(bad, edge.Site)
			}
		}
	}

	return bad
}
<commit_msg>Update usage of go tools ssa In particular, use Build, renamed from BuildAll in golang/tools@afcda55b39c7af49a70f0b6bdb9c381a95205816<commit_after>// Command safesql is a tool for performing static analysis on programs to
// ensure that SQL injection attacks are not possible. It does this by ensuring
// package database/sql is only used with compile-time constant queries.
package main

import (
	"flag"
	"fmt"
	"os"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
	"golang.org/x/tools/go/types"
)

func main() {
	var verbose, quiet bool
	flag.BoolVar(&verbose, "v", false, "Verbose mode")
	flag.BoolVar(&quiet, "q", false, "Only print on failure")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [-q] [-v] package1 [package2 ...]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.Parse()
	pkgs := flag.Args()
	if len(pkgs) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	c := loader.Config{}
	c.Import("database/sql")
	for _, pkg := range pkgs {
		c.Import(pkg)
	}
	p, err := c.Load()
	if err != nil {
		fmt.Printf("error loading packages %v: %v\n", pkgs, err)
		os.Exit(2)
	}
	s := ssautil.CreateProgram(p, 0)
	s.Build()

	qms := FindQueryMethods(p.Package("database/sql").Pkg, s)
	if verbose {
		fmt.Println("database/sql functions that accept queries:")
		for _, m := range qms {
			fmt.Printf("- %s (param %d)\n", m.Func, m.Param)
		}
		fmt.Println()
	}

	mains := FindMains(p, s)
	if len(mains) == 0 {
		fmt.Println("Did not find any commands (i.e., main functions).")
		os.Exit(2)
	}

	res, err := pointer.Analyze(&pointer.Config{
		Mains:          mains,
		BuildCallGraph: true,
	})
	if err != nil {
		fmt.Printf("error performing pointer analysis: %v\n", err)
		os.Exit(2)
	}

	bad := FindNonConstCalls(res.CallGraph, qms)
	if len(bad) == 0 {
		if !quiet {
			fmt.Println(`You're safe from SQL injection! Yay \o/`)
		}
		return
	}

	fmt.Printf("Found %d potentially unsafe SQL statements:\n", len(bad))
	for _, ci := range bad {
		pos := p.Fset.Position(ci.Pos())
		fmt.Printf("- %s\n", pos)
	}
	fmt.Println("Please ensure that all SQL queries you use are compile-time constants.")
	fmt.Println("You should always use parameterized queries or prepared statements")
	fmt.Println("instead of building queries from strings.")
	os.Exit(1)
}

// QueryMethod represents a method on a type which has a string parameter named
// "query".
type QueryMethod struct {
	Func     *types.Func
	SSA      *ssa.Function
	ArgCount int
	Param    int
}

// FindQueryMethods locates all methods in the given package (assumed to be
// package database/sql) with a string parameter named "query".
func FindQueryMethods(sql *types.Package, ssa *ssa.Program) []*QueryMethod {
	methods := make([]*QueryMethod, 0)
	scope := sql.Scope()
	for _, name := range scope.Names() {
		o := scope.Lookup(name)
		if !o.Exported() {
			continue
		}
		if _, ok := o.(*types.TypeName); !ok {
			continue
		}
		n := o.Type().(*types.Named)
		for i := 0; i < n.NumMethods(); i++ {
			m := n.Method(i)
			if !m.Exported() {
				continue
			}
			s := m.Type().(*types.Signature)
			if num, ok := FuncHasQuery(s); ok {
				methods = append(methods, &QueryMethod{
					Func:     m,
					SSA:      ssa.FuncValue(m),
					ArgCount: s.Params().Len(),
					Param:    num,
				})
			}
		}
	}
	return methods
}

var stringType types.Type = types.Typ[types.String]

// FuncHasQuery returns the offset of the string parameter named "query", or
// none if no such parameter exists.
func FuncHasQuery(s *types.Signature) (offset int, ok bool) {
	params := s.Params()
	for i := 0; i < params.Len(); i++ {
		v := params.At(i)
		if v.Name() == "query" && v.Type() == stringType {
			return i, true
		}
	}
	return 0, false
}

// FindMains returns the set of all packages loaded into the given
// loader.Program which contain main functions
func FindMains(p *loader.Program, s *ssa.Program) []*ssa.Package {
	ips := p.InitialPackages()
	mains := make([]*ssa.Package, 0, len(ips))
	for _, info := range ips {
		ssaPkg := s.Package(info.Pkg)
		if ssaPkg.Func("main") != nil {
			mains = append(mains, ssaPkg)
		}
	}
	return mains
}

// FindNonConstCalls returns the set of callsites of the given set of methods
// for which the "query" parameter is not a compile-time constant.
func FindNonConstCalls(cg *callgraph.Graph, qms []*QueryMethod) []ssa.CallInstruction {
	cg.DeleteSyntheticNodes()

	// package database/sql has a couple helper functions which are thin
	// wrappers around other sensitive functions. Instead of handling the
	// general case by tracing down callsites of wrapper functions
	// recursively, let's just whitelist the functions we're already
	// tracking, since it happens to be good enough for our use case.
	okFuncs := make(map[*ssa.Function]struct{}, len(qms))
	for _, m := range qms {
		okFuncs[m.SSA] = struct{}{}
	}

	bad := make([]ssa.CallInstruction, 0)
	for _, m := range qms {
		node := cg.CreateNode(m.SSA)
		for _, edge := range node.In {
			if _, ok := okFuncs[edge.Site.Parent()]; ok {
				continue
			}
			cc := edge.Site.Common()
			args := cc.Args
			// The first parameter is occasionally the receiver.
			if len(args) == m.ArgCount+1 {
				args = args[1:]
			} else if len(args) != m.ArgCount {
				panic("arg count mismatch")
			}
			v := args[m.Param]
			if _, ok := v.(*ssa.Const); !ok {
				bad = append(bad, edge.Site)
			}
		}
	}

	return bad
}
<|endoftext|>


<commit_before>/*
Copyright 2016 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 sshutil

import (
	"bufio"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/docker/machine/libmachine/drivers"
	machinessh "github.com/docker/machine/libmachine/ssh"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
	"k8s.io/klog/v2"

	"k8s.io/minikube/pkg/util/retry"
)

// NewSSHClient returns an SSH client object for running commands.
func NewSSHClient(d drivers.Driver) (*ssh.Client, error) {
	h, err := newSSHHost(d)
	if err != nil {
		return nil, errors.Wrap(err, "Error creating new ssh host from driver")

	}
	auth := &machinessh.Auth{}
	if h.SSHKeyPath != "" {
		auth.Keys = []string{h.SSHKeyPath}
	}

	klog.Infof("new ssh client: %+v", h)

	config, err := machinessh.NewNativeConfig(h.Username, auth)
	if err != nil {
		return nil, errors.Wrapf(err, "Error creating new native config from ssh using: %s, %s", h.Username, auth)
	}

	var client *ssh.Client
	getSSH := func() (err error) {
		client, err = ssh.Dial("tcp", net.JoinHostPort(h.IP, strconv.Itoa(h.Port)), &config)
		if err != nil {
			klog.Warningf("dial failure (will retry): %v", err)
		}
		return err
	}

	if err := retry.Expo(getSSH, 250*time.Millisecond, 2*time.Second); err != nil {
		return nil, err
	}

	return client, nil
}

type sshHost struct {
	IP         string
	Port       int
	SSHKeyPath string
	Username   string
}

func newSSHHost(d drivers.Driver) (*sshHost, error) {

	ip, err := d.GetSSHHostname()
	if err != nil {
		return nil, errors.Wrap(err, "Error getting ssh host name for driver")
	}
	port, err := d.GetSSHPort()
	if err != nil {
		return nil, errors.Wrap(err, "Error getting ssh port for driver")
	}
	return &sshHost{
		IP:         ip,
		Port:       port,
		SSHKeyPath: d.GetSSHKeyPath(),
		Username:   d.GetSSHUsername(),
	}, nil
}

func KnownHost(host string, knownHosts string) bool {
	fd, err := os.Open(knownHosts)
	if err != nil {
		return false
	}
	defer fd.Close()

	hashhost := knownhosts.HashHostname(host)
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		_, hosts, _, _, _, err := ssh.ParseKnownHosts(scanner.Bytes())
		if err != nil {
			continue
		}

		for _, h := range hosts {
			if h == host || h == hashhost {
				return true
			}
		}
	}

	return false
}
<commit_msg>ssh: support default key path for built-in ssh<commit_after>/*
Copyright 2016 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 sshutil

import (
	"bufio"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/docker/machine/libmachine/drivers"
	machinessh "github.com/docker/machine/libmachine/ssh"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
	"k8s.io/client-go/util/homedir"
	"k8s.io/klog/v2"

	"k8s.io/minikube/pkg/util/retry"
)

// NewSSHClient returns an SSH client object for running commands.
func NewSSHClient(d drivers.Driver) (*ssh.Client, error) {
	h, err := newSSHHost(d)
	if err != nil {
		return nil, errors.Wrap(err, "Error creating new ssh host from driver")

	}
	defaultKeyPath := filepath.Join(homedir.HomeDir(), ".ssh", "id_rsa")
	auth := &machinessh.Auth{}
	if h.SSHKeyPath != "" {
		auth.Keys = []string{h.SSHKeyPath}
	} else {
		auth.Keys = []string{defaultKeyPath}
	}

	klog.Infof("new ssh client: %+v", h)

	config, err := machinessh.NewNativeConfig(h.Username, auth)
	if err != nil {
		return nil, errors.Wrapf(err, "Error creating new native config from ssh using: %s, %s", h.Username, auth)
	}

	var client *ssh.Client
	getSSH := func() (err error) {
		client, err = ssh.Dial("tcp", net.JoinHostPort(h.IP, strconv.Itoa(h.Port)), &config)
		if err != nil {
			klog.Warningf("dial failure (will retry): %v", err)
		}
		return err
	}

	if err := retry.Expo(getSSH, 250*time.Millisecond, 2*time.Second); err != nil {
		return nil, err
	}

	return client, nil
}

type sshHost struct {
	IP         string
	Port       int
	SSHKeyPath string
	Username   string
}

func newSSHHost(d drivers.Driver) (*sshHost, error) {

	ip, err := d.GetSSHHostname()
	if err != nil {
		return nil, errors.Wrap(err, "Error getting ssh host name for driver")
	}
	port, err := d.GetSSHPort()
	if err != nil {
		return nil, errors.Wrap(err, "Error getting ssh port for driver")
	}
	return &sshHost{
		IP:         ip,
		Port:       port,
		SSHKeyPath: d.GetSSHKeyPath(),
		Username:   d.GetSSHUsername(),
	}, nil
}

func KnownHost(host string, knownHosts string) bool {
	fd, err := os.Open(knownHosts)
	if err != nil {
		return false
	}
	defer fd.Close()

	hashhost := knownhosts.HashHostname(host)
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		_, hosts, _, _, _, err := ssh.ParseKnownHosts(scanner.Bytes())
		if err != nil {
			continue
		}

		for _, h := range hosts {
			if h == host || h == hashhost {
				return true
			}
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright 2021 The Prometheus Authors
// 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 tsdb

import (
	"sync"

	"github.com/go-kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"go.uber.org/atomic"

	"github.com/grafana/loki/pkg/storage/stores/tsdb/index"
)

/*
Disclaimer: This is largely inspired from Prometheus' TSDB Head, albeit
with significant changes (generally reductions rather than additions) to accommodate Loki
*/

const (
	// Note, this is significantly less than the stripe values used by Prometheus' stripeSeries.
	// This is for two reasons.
	// 1) Heads are per-tenant in Loki
	// 2) Loki tends to have a few orders of magnitude less series per node than
	// Prometheus|Cortex|Mimir.
	// Do not specify without bit shifting. This allows us to
	// do shard index calcuations via bitwise & rather than modulos.
	defaultStripeSize = 64
)

/*
Head is a per-tenant accumulator for index entries in memory.
It can be queried as an IndexReader and consumed to generate a TSDB index.
These are written to on the ingester component when chunks are flushed,
then written to disk as per tenant TSDB indices at the end of the WAL checkpointing cycle.
Every n cycles, they are compacted together and written to object storage.

In turn, many `Head`s may be wrapped into a multi-tenant head.
This allows Loki to serve `GetChunkRefs` requests for _chunks_ which have been flushed
whereas the corresponding index has not yet been uploaded to object storage,
guaranteeing we maintain querying consistency for the entire data lifecycle.
*/

// TODO(owen-d)
type Metrics struct {
	seriesNotFound                prometheus.Counter
	tsdbCreationsTotal            prometheus.Counter
	tsdbCreationFailures          prometheus.Counter
	tsdbManagerUpdatesTotal       prometheus.Counter
	tsdbManagerUpdatesFailedTotal prometheus.Counter
}

func NewMetrics(r prometheus.Registerer) *Metrics {
	return &Metrics{
		seriesNotFound: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_head_series_not_found_total",
			Help: "Total number of requests for series that were not found.",
		}),
		tsdbCreationsTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_creations_total",
			Help: "Total number of tsdb creations attempted",
		}),
		tsdbCreationFailures: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_creations_failed_total",
			Help: "Total number of tsdb creations failed",
		}),
		tsdbManagerUpdatesTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_manager_updates_total",
			Help: "Total number of tsdb manager updates (loading/rotating tsdbs in mem)",
		}),
		tsdbManagerUpdatesFailedTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_manager_updates_failed_total",
			Help: "Total number of tsdb manager update failures (loading/rotating tsdbs in mem)",
		}),
	}
}

type Head struct {
	tenant           string
	numSeries        atomic.Uint64
	minTime, maxTime atomic.Int64 // Current min and max of the samples included in the head.

	// auto incrementing counter to uniquely identify series. This is also used
	// in the MemPostings, but is eventually discarded when we create a real TSDB index.
	lastSeriesID atomic.Uint64

	metrics *Metrics
	logger  log.Logger

	series *stripeSeries

	postings *index.MemPostings // Postings lists for terms.
}

func NewHead(tenant string, metrics *Metrics, logger log.Logger) *Head {
	return &Head{
		tenant:   tenant,
		metrics:  metrics,
		logger:   logger,
		series:   newStripeSeries(),
		postings: index.NewMemPostings(),
	}
}

// MinTime returns the lowest time bound on visible data in the head.
func (h *Head) MinTime() int64 {
	return h.minTime.Load()
}

// MaxTime returns the highest timestamp seen in data of the head.
func (h *Head) MaxTime() int64 {
	return h.maxTime.Load()
}

// Will CAS until successfully updates bounds or the condition is no longer valid
func updateMintMaxt(mint, maxt int64, mintSrc, maxtSrc *atomic.Int64) {
	for {
		lt := mintSrc.Load()
		if mint >= lt && lt != 0 {
			break
		}
		if mintSrc.CAS(lt, mint) {
			break
		}
	}
	for {
		ht := maxtSrc.Load()
		if maxt <= ht {
			break
		}
		if maxtSrc.CAS(ht, maxt) {
			break
		}
	}
}

// Note: chks must not be nil or zero-length
func (h *Head) Append(ls labels.Labels, chks index.ChunkMetas) (created bool, refID uint64) {
	from, through := chks.Bounds()
	var id uint64
	created, refID = h.series.Append(ls, chks, func() *memSeries {
		id = h.lastSeriesID.Inc()
		return newMemSeries(id, ls)
	})
	updateMintMaxt(int64(from), int64(through), &h.minTime, &h.maxTime)

	if !created {
		return
	}
	h.postings.Add(storage.SeriesRef(id), ls)
	h.numSeries.Inc()
	return
}

// seriesHashmap is a simple hashmap for memSeries by their label set. It is built
// on top of a regular hashmap and holds a slice of series to resolve hash collisions.
// Its methods require the hash to be submitted with it to avoid re-computations throughout
// the code.
type seriesHashmap map[uint64][]*memSeries

func (m seriesHashmap) get(hash uint64, ls labels.Labels) *memSeries {
	for _, s := range m[hash] {
		if labels.Equal(s.ls, ls) {
			return s
		}
	}
	return nil
}

func (m seriesHashmap) set(hash uint64, s *memSeries) {
	l := m[hash]
	for i, prev := range l {
		if labels.Equal(prev.ls, s.ls) {
			l[i] = s
			return
		}
	}
	m[hash] = append(l, s)
}

type stripeSeries struct {
	shards int
	locks  []sync.RWMutex
	hashes []seriesHashmap
	// Sharded by ref. A series ref is the value of `size` when the series was being newly added.
	series []map[uint64]*memSeries
}

func newStripeSeries() *stripeSeries {
	s := &stripeSeries{
		shards: defaultStripeSize,
		locks:  make([]sync.RWMutex, defaultStripeSize),
		hashes: make([]seriesHashmap, defaultStripeSize),
		series: make([]map[uint64]*memSeries, defaultStripeSize),
	}
	for i := range s.hashes {
		s.hashes[i] = seriesHashmap{}
	}
	for i := range s.series {
		s.series[i] = map[uint64]*memSeries{}
	}
	return s
}

func (s *stripeSeries) getByID(id uint64) *memSeries {
	i := id & uint64(s.shards-1)

	s.locks[i].RLock()
	series := s.series[i][id]
	s.locks[i].RUnlock()

	return series
}

// Append adds chunks to the correct series and returns whether a new series was added
func (s *stripeSeries) Append(
	ls labels.Labels,
	chks index.ChunkMetas,
	createFn func() *memSeries,
) (created bool, refID uint64) {
	fp := ls.Hash()
	i := fp & uint64(s.shards-1)
	mtx := &s.locks[i]

	mtx.Lock()
	series := s.hashes[i].get(fp, ls)
	if series == nil {
		series = createFn()
		s.hashes[i].set(fp, series)

		// the series locks are determined by the ref, not fingerprint
		refIdx := series.ref & uint64(s.shards-1)
		s.series[refIdx][series.ref] = series
		created = true
	}
	mtx.Unlock()

	series.Lock()
	series.chks = append(series.chks, chks...)
	refID = series.ref
	series.Unlock()

	return
}

type memSeries struct {
	sync.RWMutex
	ref  uint64 // The unique reference within a *Head
	ls   labels.Labels
	fp   uint64
	chks index.ChunkMetas
}

func newMemSeries(ref uint64, ls labels.Labels) *memSeries {
	return &memSeries{
		ref: ref,
		ls:  ls,
		fp:  ls.Hash(),
	}
}
<commit_msg>[bugfix] fixes issues sharding series by storageref vs fingerprint in stripeSeries (#6882)<commit_after>// Copyright 2021 The Prometheus Authors
// 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 tsdb

import (
	"sync"

	"github.com/go-kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/prometheus/model/labels"
	"github.com/prometheus/prometheus/storage"
	"go.uber.org/atomic"

	"github.com/grafana/loki/pkg/storage/stores/tsdb/index"
)

/*
Disclaimer: This is largely inspired from Prometheus' TSDB Head, albeit
with significant changes (generally reductions rather than additions) to accommodate Loki
*/

const (
	// Note, this is significantly less than the stripe values used by Prometheus' stripeSeries.
	// This is for two reasons.
	// 1) Heads are per-tenant in Loki
	// 2) Loki tends to have a few orders of magnitude less series per node than
	// Prometheus|Cortex|Mimir.
	// Do not specify without bit shifting. This allows us to
	// do shard index calcuations via bitwise & rather than modulos.
	defaultStripeSize = 64
)

/*
Head is a per-tenant accumulator for index entries in memory.
It can be queried as an IndexReader and consumed to generate a TSDB index.
These are written to on the ingester component when chunks are flushed,
then written to disk as per tenant TSDB indices at the end of the WAL checkpointing cycle.
Every n cycles, they are compacted together and written to object storage.

In turn, many `Head`s may be wrapped into a multi-tenant head.
This allows Loki to serve `GetChunkRefs` requests for _chunks_ which have been flushed
whereas the corresponding index has not yet been uploaded to object storage,
guaranteeing we maintain querying consistency for the entire data lifecycle.
*/

// TODO(owen-d)
type Metrics struct {
	seriesNotFound                prometheus.Counter
	tsdbCreationsTotal            prometheus.Counter
	tsdbCreationFailures          prometheus.Counter
	tsdbManagerUpdatesTotal       prometheus.Counter
	tsdbManagerUpdatesFailedTotal prometheus.Counter
}

func NewMetrics(r prometheus.Registerer) *Metrics {
	return &Metrics{
		seriesNotFound: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_head_series_not_found_total",
			Help: "Total number of requests for series that were not found.",
		}),
		tsdbCreationsTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_creations_total",
			Help: "Total number of tsdb creations attempted",
		}),
		tsdbCreationFailures: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_creations_failed_total",
			Help: "Total number of tsdb creations failed",
		}),
		tsdbManagerUpdatesTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_manager_updates_total",
			Help: "Total number of tsdb manager updates (loading/rotating tsdbs in mem)",
		}),
		tsdbManagerUpdatesFailedTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
			Name: "loki_tsdb_manager_updates_failed_total",
			Help: "Total number of tsdb manager update failures (loading/rotating tsdbs in mem)",
		}),
	}
}

type Head struct {
	tenant           string
	numSeries        atomic.Uint64
	minTime, maxTime atomic.Int64 // Current min and max of the samples included in the head.

	// auto incrementing counter to uniquely identify series. This is also used
	// in the MemPostings, but is eventually discarded when we create a real TSDB index.
	lastSeriesID atomic.Uint64

	metrics *Metrics
	logger  log.Logger

	series *stripeSeries

	postings *index.MemPostings // Postings lists for terms.
}

func NewHead(tenant string, metrics *Metrics, logger log.Logger) *Head {
	return &Head{
		tenant:   tenant,
		metrics:  metrics,
		logger:   logger,
		series:   newStripeSeries(),
		postings: index.NewMemPostings(),
	}
}

// MinTime returns the lowest time bound on visible data in the head.
func (h *Head) MinTime() int64 {
	return h.minTime.Load()
}

// MaxTime returns the highest timestamp seen in data of the head.
func (h *Head) MaxTime() int64 {
	return h.maxTime.Load()
}

// Will CAS until successfully updates bounds or the condition is no longer valid
func updateMintMaxt(mint, maxt int64, mintSrc, maxtSrc *atomic.Int64) {
	for {
		lt := mintSrc.Load()
		if mint >= lt && lt != 0 {
			break
		}
		if mintSrc.CAS(lt, mint) {
			break
		}
	}
	for {
		ht := maxtSrc.Load()
		if maxt <= ht {
			break
		}
		if maxtSrc.CAS(ht, maxt) {
			break
		}
	}
}

// Note: chks must not be nil or zero-length
func (h *Head) Append(ls labels.Labels, chks index.ChunkMetas) (created bool, refID uint64) {
	from, through := chks.Bounds()
	var id uint64
	created, refID = h.series.Append(ls, chks, func() *memSeries {
		id = h.lastSeriesID.Inc()
		return newMemSeries(id, ls)
	})
	updateMintMaxt(int64(from), int64(through), &h.minTime, &h.maxTime)

	if !created {
		return
	}
	h.postings.Add(storage.SeriesRef(id), ls)
	h.numSeries.Inc()
	return
}

// seriesHashmap is a simple hashmap for memSeries by their label set. It is built
// on top of a regular hashmap and holds a slice of series to resolve hash collisions.
// Its methods require the hash to be submitted with it to avoid re-computations throughout
// the code.
type seriesHashmap struct {
	sync.RWMutex
	m map[uint64][]*memSeries
}

func newSeriesHashmap() *seriesHashmap {
	return &seriesHashmap{
		m: map[uint64][]*memSeries{},
	}
}

func (m *seriesHashmap) get(hash uint64, ls labels.Labels) *memSeries {
	for _, s := range m.m[hash] {
		if labels.Equal(s.ls, ls) {
			return s
		}
	}
	return nil
}

func (m *seriesHashmap) set(hash uint64, s *memSeries) {
	l := m.m[hash]
	for i, prev := range l {
		if labels.Equal(prev.ls, s.ls) {
			l[i] = s
			return
		}
	}
	m.m[hash] = append(l, s)
}

type memSeriesMap struct {
	sync.RWMutex
	m map[uint64]*memSeries
}

func newMemseriesMap() *memSeriesMap {
	return &memSeriesMap{
		m: map[uint64]*memSeries{},
	}
}

type stripeSeries struct {
	shards int
	// Sharded by fingerprint.
	hashes []*seriesHashmap
	// Sharded by ref. A series ref is the value of `size` when the series was being newly added.
	series []*memSeriesMap
}

func newStripeSeries() *stripeSeries {
	s := &stripeSeries{
		shards: defaultStripeSize,
		hashes: make([]*seriesHashmap, defaultStripeSize),
		series: make([]*memSeriesMap, defaultStripeSize),
	}
	for i := range s.hashes {
		s.hashes[i] = newSeriesHashmap()
	}
	for i := range s.series {
		s.series[i] = newMemseriesMap()
	}
	return s
}

func (s *stripeSeries) getByID(id uint64) *memSeries {
	x := s.series[id&uint64(s.shards-1)]
	x.RLock()
	defer x.RUnlock()
	return x.m[id]
}

// Append adds chunks to the correct series and returns whether a new series was added
func (s *stripeSeries) Append(
	ls labels.Labels,
	chks index.ChunkMetas,
	createFn func() *memSeries,
) (created bool, refID uint64) {
	fp := ls.Hash()
	hashes := s.hashes[fp&uint64(s.shards-1)]
	hashes.Lock()
	series := hashes.get(fp, ls)
	if series == nil {
		series = createFn()
		hashes.set(fp, series)

		// the series locks are determined by the ref, not fingerprint
		seriesMap := s.series[series.ref&uint64(s.shards-1)]
		seriesMap.Lock()
		seriesMap.m[series.ref] = series
		seriesMap.Unlock()

		created = true
	}
	hashes.Unlock()

	series.Lock()
	series.chks = append(series.chks, chks...)
	refID = series.ref
	series.Unlock()

	return
}

type memSeries struct {
	sync.RWMutex
	ref  uint64 // The unique reference within a *Head
	ls   labels.Labels
	fp   uint64
	chks index.ChunkMetas
}

func newMemSeries(ref uint64, ls labels.Labels) *memSeries {
	return &memSeries{
		ref: ref,
		ls:  ls,
		fp:  ls.Hash(),
	}
}
<|endoftext|>