<commit_before>package clustermanager

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	clusterController "github.com/rancher/rancher/pkg/controllers/user"
	"github.com/rancher/rancher/pkg/rbac"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/rancher/types/config/dialer"
	"github.com/sirupsen/logrus"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/client-go/rest"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	"k8s.io/client-go/util/cert"
)

type Manager struct {
	httpsPort     int
	ScaledContext *config.ScaledContext
	clusterLister v3.ClusterLister
	controllers   sync.Map
	accessControl types.AccessControl
	dialer        dialer.Factory
}

type record struct {
	clusterRec    *v3.Cluster
	cluster       *config.UserContext
	accessControl types.AccessControl
	ctx           context.Context
	cancel        context.CancelFunc
}

func NewManager(httpsPort int, context *config.ScaledContext) *Manager {
	return &Manager{
		httpsPort:     httpsPort,
		ScaledContext: context,
		dialer:        context.Dialer,
		accessControl: rbac.NewAccessControl(context.RBAC),
		clusterLister: context.Management.Clusters("").Controller().Lister(),
	}
}

func (m *Manager) Stop(cluster *v3.Cluster) {
	obj, ok := m.controllers.Load(cluster.UID)
	if !ok {
		return
	}
	logrus.Info("Stopping cluster agent for", obj.(*record).cluster.ClusterName)
	obj.(*record).cancel()
	m.controllers.Delete(cluster.UID)
}

func (m *Manager) Start(ctx context.Context, cluster *v3.Cluster) error {
	// reload cluster, always use the cached one
	cluster, err := m.clusterLister.Get("", cluster.Name)
	if err != nil {
		return err
	}
	_, err = m.start(ctx, cluster)
	return err
}

func (m *Manager) RESTConfig(cluster *v3.Cluster) (rest.Config, error) {
	obj, ok := m.controllers.Load(cluster.UID)
	if !ok {
		return rest.Config{}, fmt.Errorf("cluster record not found %s %s", cluster.Name, cluster.UID)
	}

	record := obj.(*record)
	return record.cluster.RESTConfig, nil
}

func (m *Manager) start(ctx context.Context, cluster *v3.Cluster) (*record, error) {
	obj, ok := m.controllers.Load(cluster.UID)
	if ok {
		if !m.changed(obj.(*record), cluster) {
			return obj.(*record), nil
		}
		m.Stop(obj.(*record).clusterRec)
	}

	controller, err := m.toRecord(ctx, cluster)
	if err != nil {
		return nil, err
	}
	if controller == nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, "cluster not found")
	}

	obj, loaded := m.controllers.LoadOrStore(cluster.UID, controller)
	if !loaded {
		go func() {
			if err := m.doStart(obj.(*record)); err != nil {
				m.Stop(cluster)
			}
		}()
	}

	return obj.(*record), nil
}

func (m *Manager) changed(r *record, cluster *v3.Cluster) bool {
	existing := r.clusterRec
	if existing.Status.APIEndpoint != cluster.Status.APIEndpoint ||
		existing.Status.ServiceAccountToken != cluster.Status.ServiceAccountToken ||
		existing.Status.CACert != cluster.Status.CACert {
		return true
	}

	return false
}

func (m *Manager) doStart(rec *record) error {
	logrus.Info("Starting cluster agent for", rec.cluster.ClusterName)
	if err := clusterController.Register(rec.ctx, rec.cluster, m, m); err != nil {
		return err
	}
	return rec.cluster.Start(rec.ctx)
}

func (m *Manager) toRESTConfig(cluster *v3.Cluster) (*rest.Config, error) {
	if cluster == nil {
		return nil, nil
	}

	if cluster.DeletionTimestamp != nil {
		return nil, nil
	}

	if cluster.Spec.Internal {
		return m.ScaledContext.LocalConfig, nil
	}

	if cluster.Status.APIEndpoint == "" || cluster.Status.CACert == "" || cluster.Status.ServiceAccountToken == "" {
		return nil, nil
	}

	if !v3.ClusterConditionProvisioned.IsTrue(cluster) {
		return nil, nil
	}

	u, err := url.Parse(cluster.Status.APIEndpoint)
	if err != nil {
		return nil, err
	}

	caBytes, err := base64.StdEncoding.DecodeString(cluster.Status.CACert)
	if err != nil {
		return nil, err
	}

	clusterDialer, err := m.dialer.ClusterDialer(cluster.Name)
	if err != nil {
		return nil, err
	}

	var tlsDialer dialer.Dialer
	if cluster.Status.Driver == v3.ClusterDriverRKE {
		tlsDialer, err = nameIgnoringTLSDialer(clusterDialer, caBytes)
		if err != nil {
			return nil, err
		}
	}

	rc := &rest.Config{
		Host:        u.Host,
		Prefix:      u.Path,
		BearerToken: cluster.Status.ServiceAccountToken,
		TLSClientConfig: rest.TLSClientConfig{
			CAData: caBytes,
		},
		Timeout: 30 * time.Second,
		WrapTransport: func(rt http.RoundTripper) http.RoundTripper {
			if ht, ok := rt.(*http.Transport); ok {
				ht.DialContext = nil
				ht.DialTLS = tlsDialer
				ht.Dial = clusterDialer
			}
			return rt
		},
	}

	return rc, nil
}

func nameIgnoringTLSDialer(dialer dialer.Dialer, caBytes []byte) (dialer.Dialer, error) {
	rkeVerify, err := VerifyIgnoreDNSName(caBytes)
	if err != nil {
		return nil, err
	}

	tlsConfig := &tls.Config{
		// Use custom TLS validate that validates the cert chain, but not the server.  This should be secure because
		// we use a private per cluster CA always for RKE
		InsecureSkipVerify:    true,
		VerifyPeerCertificate: rkeVerify,
	}

	return func(network, address string) (net.Conn, error) {
		rawConn, err := dialer(network, address)
		if err != nil {
			return nil, err
		}
		tlsConn := tls.Client(rawConn, tlsConfig)
		if err := tlsConn.Handshake(); err != nil {
			rawConn.Close()
			return nil, err
		}
		return tlsConn, err
	}, nil
}

func VerifyIgnoreDNSName(caCertsPEM []byte) (func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error, error) {
	rootCAs := x509.NewCertPool()
	if len(caCertsPEM) > 0 {
		caCerts, err := cert.ParseCertsPEM(caCertsPEM)
		if err != nil {
			return nil, err
		}
		for _, cert := range caCerts {
			rootCAs.AddCert(cert)
		}
	}

	return func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
		certs := make([]*x509.Certificate, len(rawCerts))
		for i, asn1Data := range rawCerts {
			cert, err := x509.ParseCertificate(asn1Data)
			if err != nil {
				return fmt.Errorf("failed to parse cert")
			}
			certs[i] = cert
		}

		opts := x509.VerifyOptions{
			Roots:         rootCAs,
			CurrentTime:   time.Now(),
			DNSName:       "",
			Intermediates: x509.NewCertPool(),
		}

		for i, cert := range certs {
			if i == 0 {
				continue
			}
			opts.Intermediates.AddCert(cert)
		}
		_, err := certs[0].Verify(opts)
		return err
	}, nil
}

func (m *Manager) toRecord(ctx context.Context, cluster *v3.Cluster) (*record, error) {
	kubeConfig, err := m.toRESTConfig(cluster)
	if kubeConfig == nil || err != nil {
		return nil, err
	}

	clusterContext, err := config.NewUserContext(m.ScaledContext, *kubeConfig, cluster.Name)
	if err != nil {
		return nil, err
	}

	s := &record{
		cluster:       clusterContext,
		clusterRec:    cluster,
		accessControl: rbac.NewAccessControl(clusterContext.RBAC),
	}
	s.ctx, s.cancel = context.WithCancel(ctx)

	return s, nil
}

func (m *Manager) AccessControl(apiContext *types.APIContext, storageContext types.StorageContext) (types.AccessControl, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.accessControl, nil
	}
	return record.accessControl, nil
}

func (m *Manager) Config(apiContext *types.APIContext, storageContext types.StorageContext) (rest.Config, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return rest.Config{}, err
	}
	if record == nil {
		return m.ScaledContext.RESTConfig, nil
	}
	return record.cluster.RESTConfig, nil
}

func (m *Manager) UnversionedClient(apiContext *types.APIContext, storageContext types.StorageContext) (rest.Interface, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.ScaledContext.UnversionedClient, nil
	}
	return record.cluster.UnversionedClient, nil
}

func (m *Manager) APIExtClient(apiContext *types.APIContext, storageContext types.StorageContext) (clientset.Interface, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.ScaledContext.APIExtClient, nil
	}
	return record.cluster.APIExtClient, nil
}

func (m *Manager) UserContext(clusterName string) (*config.UserContext, error) {
	cluster, err := m.clusterLister.Get("", clusterName)
	if err != nil {
		return nil, err
	}

	record, err := m.start(context.Background(), cluster)
	if err != nil || record == nil {
		msg := ""
		if err != nil {
			msg = err.Error()
		}
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, msg)
	}

	if record == nil {
		return nil, httperror.NewAPIError(httperror.NotFound, "failed to find cluster")
	}

	return record.cluster, nil
}

func (m *Manager) record(apiContext *types.APIContext, storageContext types.StorageContext) (*record, error) {
	if apiContext == nil {
		return nil, nil
	}
	cluster, err := m.cluster(apiContext, storageContext)
	if err != nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, err.Error())
	}
	if cluster == nil {
		return nil, nil
	}
	record, err := m.start(context.Background(), cluster)
	if err != nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, err.Error())
	}

	return record, nil
}

func (m *Manager) ClusterName(apiContext *types.APIContext) string {
	clusterID := apiContext.SubContext["/v3/schemas/cluster"]
	if clusterID == "" {
		projectID, ok := apiContext.SubContext["/v3/schemas/project"]
		if ok {
			parts := strings.SplitN(projectID, ":", 2)
			if len(parts) == 2 {
				clusterID = parts[0]
			}
		}
	}
	return clusterID
}

func (m *Manager) cluster(apiContext *types.APIContext, context types.StorageContext) (*v3.Cluster, error) {
	switch context {
	case types.DefaultStorageContext:
		return nil, nil
	case config.ManagementStorageContext:
		return nil, nil
	case config.UserStorageContext:
	default:
		return nil, fmt.Errorf("illegal context: %s", context)

	}

	clusterID := m.ClusterName(apiContext)
	if clusterID == "" {
		return nil, nil
	}

	return m.clusterLister.Get("", clusterID)
}

func (m *Manager) KubeConfig(clusterName, token string) *clientcmdapi.Config {
	return &clientcmdapi.Config{
		CurrentContext: "default",
		APIVersion:     "v1",
		Kind:           "Config",
		Clusters: map[string]*clientcmdapi.Cluster{
			"default": {
				Server:                fmt.Sprintf("https://localhost:%d/k8s/clusters/%s", m.httpsPort, clusterName),
				InsecureSkipTLSVerify: true,
			},
		},
		Contexts: map[string]*clientcmdapi.Context{
			"default": {
				AuthInfo: "user",
				Cluster:  "default",
			},
		},
		AuthInfos: map[string]*clientcmdapi.AuthInfo{
			"user": {
				Token: token,
			},
		},
	}
}

func (m *Manager) GetHTTPSPort() int {
	return m.httpsPort
}
<commit_msg>Mark cluster unavailable on start if we can't ping them<commit_after>package clustermanager

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	clusterController "github.com/rancher/rancher/pkg/controllers/user"
	"github.com/rancher/rancher/pkg/rbac"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/rancher/types/config/dialer"
	"github.com/sirupsen/logrus"
	"k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/rest"
	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
	"k8s.io/client-go/util/cert"
)

type Manager struct {
	httpsPort     int
	ScaledContext *config.ScaledContext
	clusterLister v3.ClusterLister
	clusters      v3.ClusterInterface
	controllers   sync.Map
	accessControl types.AccessControl
	dialer        dialer.Factory
}

type record struct {
	clusterRec    *v3.Cluster
	cluster       *config.UserContext
	accessControl types.AccessControl
	ctx           context.Context
	cancel        context.CancelFunc
}

func NewManager(httpsPort int, context *config.ScaledContext) *Manager {
	return &Manager{
		httpsPort:     httpsPort,
		ScaledContext: context,
		dialer:        context.Dialer,
		accessControl: rbac.NewAccessControl(context.RBAC),
		clusterLister: context.Management.Clusters("").Controller().Lister(),
		clusters:      context.Management.Clusters(""),
	}
}

func (m *Manager) Stop(cluster *v3.Cluster) {
	obj, ok := m.controllers.Load(cluster.UID)
	if !ok {
		return
	}
	logrus.Info("Stopping cluster agent for", obj.(*record).cluster.ClusterName)
	obj.(*record).cancel()
	m.controllers.Delete(cluster.UID)
}

func (m *Manager) Start(ctx context.Context, cluster *v3.Cluster) error {
	// reload cluster, always use the cached one
	cluster, err := m.clusterLister.Get("", cluster.Name)
	if err != nil {
		return err
	}
	_, err = m.start(ctx, cluster)
	return err
}

func (m *Manager) RESTConfig(cluster *v3.Cluster) (rest.Config, error) {
	obj, ok := m.controllers.Load(cluster.UID)
	if !ok {
		return rest.Config{}, fmt.Errorf("cluster record not found %s %s", cluster.Name, cluster.UID)
	}

	record := obj.(*record)
	return record.cluster.RESTConfig, nil
}

func (m *Manager) markUnavailable(clusterName string) {
	if cluster, err := m.clusters.Get(clusterName, v1.GetOptions{}); err == nil {
		if !v3.ClusterConditionReady.IsFalse(cluster) {
			v3.ClusterConditionReady.False(cluster)
			m.clusters.Update(cluster)
		}
	}
}

func (m *Manager) start(ctx context.Context, cluster *v3.Cluster) (*record, error) {
	obj, ok := m.controllers.Load(cluster.UID)
	if ok {
		if !m.changed(obj.(*record), cluster) {
			return obj.(*record), nil
		}
		m.Stop(obj.(*record).clusterRec)
	}

	controller, err := m.toRecord(ctx, cluster)
	if err != nil {
		m.markUnavailable(cluster.Name)
		return nil, err
	}
	if controller == nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, "cluster not found")
	}

	obj, loaded := m.controllers.LoadOrStore(cluster.UID, controller)
	if !loaded {
		go func() {
			if err := m.doStart(obj.(*record)); err != nil {
				m.Stop(cluster)
			}
		}()
	}

	return obj.(*record), nil
}

func (m *Manager) changed(r *record, cluster *v3.Cluster) bool {
	existing := r.clusterRec
	if existing.Status.APIEndpoint != cluster.Status.APIEndpoint ||
		existing.Status.ServiceAccountToken != cluster.Status.ServiceAccountToken ||
		existing.Status.CACert != cluster.Status.CACert {
		return true
	}

	return false
}

func (m *Manager) doStart(rec *record) error {
	logrus.Info("Starting cluster agent for", rec.cluster.ClusterName)
	if err := clusterController.Register(rec.ctx, rec.cluster, m, m); err != nil {
		return err
	}
	return rec.cluster.Start(rec.ctx)
}

func (m *Manager) toRESTConfig(cluster *v3.Cluster) (*rest.Config, error) {
	if cluster == nil {
		return nil, nil
	}

	if cluster.DeletionTimestamp != nil {
		return nil, nil
	}

	if cluster.Spec.Internal {
		return m.ScaledContext.LocalConfig, nil
	}

	if cluster.Status.APIEndpoint == "" || cluster.Status.CACert == "" || cluster.Status.ServiceAccountToken == "" {
		return nil, nil
	}

	if !v3.ClusterConditionProvisioned.IsTrue(cluster) {
		return nil, nil
	}

	u, err := url.Parse(cluster.Status.APIEndpoint)
	if err != nil {
		return nil, err
	}

	caBytes, err := base64.StdEncoding.DecodeString(cluster.Status.CACert)
	if err != nil {
		return nil, err
	}

	clusterDialer, err := m.dialer.ClusterDialer(cluster.Name)
	if err != nil {
		return nil, err
	}

	var tlsDialer dialer.Dialer
	if cluster.Status.Driver == v3.ClusterDriverRKE {
		tlsDialer, err = nameIgnoringTLSDialer(clusterDialer, caBytes)
		if err != nil {
			return nil, err
		}
	}

	rc := &rest.Config{
		Host:        u.Host,
		Prefix:      u.Path,
		BearerToken: cluster.Status.ServiceAccountToken,
		TLSClientConfig: rest.TLSClientConfig{
			CAData: caBytes,
		},
		Timeout: 30 * time.Second,
		WrapTransport: func(rt http.RoundTripper) http.RoundTripper {
			if ht, ok := rt.(*http.Transport); ok {
				ht.DialContext = nil
				ht.DialTLS = tlsDialer
				ht.Dial = clusterDialer
			}
			return rt
		},
	}

	return rc, nil
}

func nameIgnoringTLSDialer(dialer dialer.Dialer, caBytes []byte) (dialer.Dialer, error) {
	rkeVerify, err := VerifyIgnoreDNSName(caBytes)
	if err != nil {
		return nil, err
	}

	tlsConfig := &tls.Config{
		// Use custom TLS validate that validates the cert chain, but not the server.  This should be secure because
		// we use a private per cluster CA always for RKE
		InsecureSkipVerify:    true,
		VerifyPeerCertificate: rkeVerify,
	}

	return func(network, address string) (net.Conn, error) {
		rawConn, err := dialer(network, address)
		if err != nil {
			return nil, err
		}
		tlsConn := tls.Client(rawConn, tlsConfig)
		if err := tlsConn.Handshake(); err != nil {
			rawConn.Close()
			return nil, err
		}
		return tlsConn, err
	}, nil
}

func VerifyIgnoreDNSName(caCertsPEM []byte) (func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error, error) {
	rootCAs := x509.NewCertPool()
	if len(caCertsPEM) > 0 {
		caCerts, err := cert.ParseCertsPEM(caCertsPEM)
		if err != nil {
			return nil, err
		}
		for _, cert := range caCerts {
			rootCAs.AddCert(cert)
		}
	}

	return func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
		certs := make([]*x509.Certificate, len(rawCerts))
		for i, asn1Data := range rawCerts {
			cert, err := x509.ParseCertificate(asn1Data)
			if err != nil {
				return fmt.Errorf("failed to parse cert")
			}
			certs[i] = cert
		}

		opts := x509.VerifyOptions{
			Roots:         rootCAs,
			CurrentTime:   time.Now(),
			DNSName:       "",
			Intermediates: x509.NewCertPool(),
		}

		for i, cert := range certs {
			if i == 0 {
				continue
			}
			opts.Intermediates.AddCert(cert)
		}
		_, err := certs[0].Verify(opts)
		return err
	}, nil
}

func (m *Manager) toRecord(ctx context.Context, cluster *v3.Cluster) (*record, error) {
	kubeConfig, err := m.toRESTConfig(cluster)
	if kubeConfig == nil || err != nil {
		return nil, err
	}

	clusterContext, err := config.NewUserContext(m.ScaledContext, *kubeConfig, cluster.Name)
	if err != nil {
		return nil, err
	}

	s := &record{
		cluster:       clusterContext,
		clusterRec:    cluster,
		accessControl: rbac.NewAccessControl(clusterContext.RBAC),
	}
	s.ctx, s.cancel = context.WithCancel(ctx)

	return s, nil
}

func (m *Manager) AccessControl(apiContext *types.APIContext, storageContext types.StorageContext) (types.AccessControl, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.accessControl, nil
	}
	return record.accessControl, nil
}

func (m *Manager) Config(apiContext *types.APIContext, storageContext types.StorageContext) (rest.Config, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return rest.Config{}, err
	}
	if record == nil {
		return m.ScaledContext.RESTConfig, nil
	}
	return record.cluster.RESTConfig, nil
}

func (m *Manager) UnversionedClient(apiContext *types.APIContext, storageContext types.StorageContext) (rest.Interface, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.ScaledContext.UnversionedClient, nil
	}
	return record.cluster.UnversionedClient, nil
}

func (m *Manager) APIExtClient(apiContext *types.APIContext, storageContext types.StorageContext) (clientset.Interface, error) {
	record, err := m.record(apiContext, storageContext)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return m.ScaledContext.APIExtClient, nil
	}
	return record.cluster.APIExtClient, nil
}

func (m *Manager) UserContext(clusterName string) (*config.UserContext, error) {
	cluster, err := m.clusterLister.Get("", clusterName)
	if err != nil {
		return nil, err
	}

	record, err := m.start(context.Background(), cluster)
	if err != nil || record == nil {
		msg := ""
		if err != nil {
			msg = err.Error()
		}
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, msg)
	}

	if record == nil {
		return nil, httperror.NewAPIError(httperror.NotFound, "failed to find cluster")
	}

	return record.cluster, nil
}

func (m *Manager) record(apiContext *types.APIContext, storageContext types.StorageContext) (*record, error) {
	if apiContext == nil {
		return nil, nil
	}
	cluster, err := m.cluster(apiContext, storageContext)
	if err != nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, err.Error())
	}
	if cluster == nil {
		return nil, nil
	}
	record, err := m.start(context.Background(), cluster)
	if err != nil {
		return nil, httperror.NewAPIError(httperror.ClusterUnavailable, err.Error())
	}

	return record, nil
}

func (m *Manager) ClusterName(apiContext *types.APIContext) string {
	clusterID := apiContext.SubContext["/v3/schemas/cluster"]
	if clusterID == "" {
		projectID, ok := apiContext.SubContext["/v3/schemas/project"]
		if ok {
			parts := strings.SplitN(projectID, ":", 2)
			if len(parts) == 2 {
				clusterID = parts[0]
			}
		}
	}
	return clusterID
}

func (m *Manager) cluster(apiContext *types.APIContext, context types.StorageContext) (*v3.Cluster, error) {
	switch context {
	case types.DefaultStorageContext:
		return nil, nil
	case config.ManagementStorageContext:
		return nil, nil
	case config.UserStorageContext:
	default:
		return nil, fmt.Errorf("illegal context: %s", context)

	}

	clusterID := m.ClusterName(apiContext)
	if clusterID == "" {
		return nil, nil
	}

	return m.clusterLister.Get("", clusterID)
}

func (m *Manager) KubeConfig(clusterName, token string) *clientcmdapi.Config {
	return &clientcmdapi.Config{
		CurrentContext: "default",
		APIVersion:     "v1",
		Kind:           "Config",
		Clusters: map[string]*clientcmdapi.Cluster{
			"default": {
				Server:                fmt.Sprintf("https://localhost:%d/k8s/clusters/%s", m.httpsPort, clusterName),
				InsecureSkipTLSVerify: true,
			},
		},
		Contexts: map[string]*clientcmdapi.Context{
			"default": {
				AuthInfo: "user",
				Cluster:  "default",
			},
		},
		AuthInfos: map[string]*clientcmdapi.AuthInfo{
			"user": {
				Token: token,
			},
		},
	}
}

func (m *Manager) GetHTTPSPort() int {
	return m.httpsPort
}
<|endoftext|>


<commit_before>/*
Copyright 2015 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 editor

import (
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"k8s.io/klog/v2"

	"k8s.io/kubectl/pkg/util/term"
)

const (
	// sorry, blame Git
	// TODO: on Windows rely on 'start' to launch the editor associated
	// with the given file type. If we can't because of the need of
	// blocking, use a script with 'ftype' and 'assoc' to detect it.
	defaultEditor = "vi"
	defaultShell  = "/bin/bash"
	windowsEditor = "notepad"
	windowsShell  = "cmd"
)

// Editor holds the command-line args to fire up the editor
type Editor struct {
	Args  []string
	Shell bool
}

// NewDefaultEditor creates a struct Editor that uses the OS environment to
// locate the editor program, looking at EDITOR environment variable to find
// the proper command line. If the provided editor has no spaces, or no quotes,
// it is treated as a bare command to be loaded. Otherwise, the string will
// be passed to the user's shell for execution.
func NewDefaultEditor(envs []string) Editor {
	args, shell := defaultEnvEditor(envs)
	return Editor{
		Args:  args,
		Shell: shell,
	}
}

func defaultEnvShell() []string {
	shell := os.Getenv("SHELL")
	if len(shell) == 0 {
		shell = platformize(defaultShell, windowsShell)
	}
	flag := "-c"
	if shell == windowsShell {
		flag = "/C"
	}
	return []string{shell, flag}
}

func defaultEnvEditor(envs []string) ([]string, bool) {
	var editor string
	for _, env := range envs {
		if len(env) > 0 {
			editor = os.Getenv(env)
		}
		if len(editor) > 0 {
			break
		}
	}
	if len(editor) == 0 {
		editor = platformize(defaultEditor, windowsEditor)
	}
	if !strings.Contains(editor, " ") {
		return []string{editor}, false
	}
	if !strings.ContainsAny(editor, "\"'\\") {
		return strings.Split(editor, " "), false
	}
	// rather than parse the shell arguments ourselves, punt to the shell
	shell := defaultEnvShell()
	return append(shell, editor), true
}

func (e Editor) args(path string) []string {
	args := make([]string, len(e.Args))
	copy(args, e.Args)
	if e.Shell {
		last := args[len(args)-1]
		args[len(args)-1] = fmt.Sprintf("%s %q", last, path)
	} else {
		args = append(args, path)
	}
	return args
}

// Launch opens the described or returns an error. The TTY will be protected, and
// SIGQUIT, SIGTERM, and SIGINT will all be trapped.
func (e Editor) Launch(path string) error {
	if len(e.Args) == 0 {
		return fmt.Errorf("no editor defined, can't open %s", path)
	}
	abs, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	args := e.args(abs)
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	klog.V(5).Infof("Opening file with editor %v", args)
	if err := (term.TTY{In: os.Stdin, TryDev: true}).Safe(cmd.Run); err != nil {
		if err, ok := err.(*exec.Error); ok {
			if err.Err == exec.ErrNotFound {
				return fmt.Errorf("unable to launch the editor %q", strings.Join(e.Args, " "))
			}
		}
		return fmt.Errorf("there was a problem with the editor %q", strings.Join(e.Args, " "))
	}
	return nil
}

// LaunchTempFile reads the provided stream into a temporary file in the given directory
// and file prefix, and then invokes Launch with the path of that file. It will return
// the contents of the file after launch, any errors that occur, and the path of the
// temporary file so the caller can clean it up as needed.
func (e Editor) LaunchTempFile(prefix, suffix string, r io.Reader) ([]byte, string, error) {
	f, err := tempFile(prefix, suffix)
	if err != nil {
		return nil, "", err
	}
	defer f.Close()
	path := f.Name()
	if _, err := io.Copy(f, r); err != nil {
		os.Remove(path)
		return nil, path, err
	}
	// This file descriptor needs to close so the next process (Launch) can claim it.
	f.Close()
	if err := e.Launch(path); err != nil {
		return nil, path, err
	}
	bytes, err := ioutil.ReadFile(path)
	return bytes, path, err
}

func tempFile(prefix, suffix string) (f *os.File, err error) {
	dir := os.TempDir()

	for i := 0; i < 10000; i++ {
		name := filepath.Join(dir, prefix+randSeq(5)+suffix)
		f, err = os.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600)
		if os.IsExist(err) {
			continue
		}
		break
	}
	return
}

var letters = []rune("abcdefghijklmnopqrstuvwxyz0123456789")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

func platformize(linux, windows string) string {
	if runtime.GOOS == "windows" {
		return windows
	}
	return linux
}
<commit_msg>Use os.CreateTemp in kubectl editor (#99921)<commit_after>/*
Copyright 2015 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 editor

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"k8s.io/klog/v2"

	"k8s.io/kubectl/pkg/util/term"
)

const (
	// sorry, blame Git
	// TODO: on Windows rely on 'start' to launch the editor associated
	// with the given file type. If we can't because of the need of
	// blocking, use a script with 'ftype' and 'assoc' to detect it.
	defaultEditor = "vi"
	defaultShell  = "/bin/bash"
	windowsEditor = "notepad"
	windowsShell  = "cmd"
)

// Editor holds the command-line args to fire up the editor
type Editor struct {
	Args  []string
	Shell bool
}

// NewDefaultEditor creates a struct Editor that uses the OS environment to
// locate the editor program, looking at EDITOR environment variable to find
// the proper command line. If the provided editor has no spaces, or no quotes,
// it is treated as a bare command to be loaded. Otherwise, the string will
// be passed to the user's shell for execution.
func NewDefaultEditor(envs []string) Editor {
	args, shell := defaultEnvEditor(envs)
	return Editor{
		Args:  args,
		Shell: shell,
	}
}

func defaultEnvShell() []string {
	shell := os.Getenv("SHELL")
	if len(shell) == 0 {
		shell = platformize(defaultShell, windowsShell)
	}
	flag := "-c"
	if shell == windowsShell {
		flag = "/C"
	}
	return []string{shell, flag}
}

func defaultEnvEditor(envs []string) ([]string, bool) {
	var editor string
	for _, env := range envs {
		if len(env) > 0 {
			editor = os.Getenv(env)
		}
		if len(editor) > 0 {
			break
		}
	}
	if len(editor) == 0 {
		editor = platformize(defaultEditor, windowsEditor)
	}
	if !strings.Contains(editor, " ") {
		return []string{editor}, false
	}
	if !strings.ContainsAny(editor, "\"'\\") {
		return strings.Split(editor, " "), false
	}
	// rather than parse the shell arguments ourselves, punt to the shell
	shell := defaultEnvShell()
	return append(shell, editor), true
}

func (e Editor) args(path string) []string {
	args := make([]string, len(e.Args))
	copy(args, e.Args)
	if e.Shell {
		last := args[len(args)-1]
		args[len(args)-1] = fmt.Sprintf("%s %q", last, path)
	} else {
		args = append(args, path)
	}
	return args
}

// Launch opens the described or returns an error. The TTY will be protected, and
// SIGQUIT, SIGTERM, and SIGINT will all be trapped.
func (e Editor) Launch(path string) error {
	if len(e.Args) == 0 {
		return fmt.Errorf("no editor defined, can't open %s", path)
	}
	abs, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	args := e.args(abs)
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	klog.V(5).Infof("Opening file with editor %v", args)
	if err := (term.TTY{In: os.Stdin, TryDev: true}).Safe(cmd.Run); err != nil {
		if err, ok := err.(*exec.Error); ok {
			if err.Err == exec.ErrNotFound {
				return fmt.Errorf("unable to launch the editor %q", strings.Join(e.Args, " "))
			}
		}
		return fmt.Errorf("there was a problem with the editor %q", strings.Join(e.Args, " "))
	}
	return nil
}

// LaunchTempFile reads the provided stream into a temporary file in the given directory
// and file prefix, and then invokes Launch with the path of that file. It will return
// the contents of the file after launch, any errors that occur, and the path of the
// temporary file so the caller can clean it up as needed.
func (e Editor) LaunchTempFile(prefix, suffix string, r io.Reader) ([]byte, string, error) {
	f, err := os.CreateTemp("", prefix+"*"+suffix)
	if err != nil {
		return nil, "", err
	}
	defer f.Close()
	path := f.Name()
	if _, err := io.Copy(f, r); err != nil {
		os.Remove(path)
		return nil, path, err
	}
	// This file descriptor needs to close so the next process (Launch) can claim it.
	f.Close()
	if err := e.Launch(path); err != nil {
		return nil, path, err
	}
	bytes, err := ioutil.ReadFile(path)
	return bytes, path, err
}

func platformize(linux, windows string) string {
	if runtime.GOOS == "windows" {
		return windows
	}
	return linux
}
<|endoftext|>


<commit_before>// Package evaltest provides a framework for testing Elvish script.
//
// The entry point for the framework is the Test function, which accepts a
// *testing.T and any number of test cases.
//
// Test cases are constructed using the That function, followed by method calls
// that add additional information to it.
//
// Example:
//
//     Test(t,
//         That("put x").Puts("x"),
//         That("echo x").Prints("x\n"))
//
// If some setup is needed, use the TestWithSetup function instead.

package evaltest

import (
	"bytes"
	"os"
	"reflect"
	"strings"
	"testing"

	"src.elv.sh/pkg/eval"
	"src.elv.sh/pkg/eval/vals"
	"src.elv.sh/pkg/parse"
	"src.elv.sh/pkg/testutil"
)

// Case is a test case that can be used in Test.
type Case struct {
	codes []string
	want  result
}

type result struct {
	ValueOut  []interface{}
	BytesOut  []byte
	StderrOut []byte

	CompilationError error
	Exception        error
}

// That returns a new Case with the specified source code. Multiple arguments
// are joined with newlines. To specify multiple pieces of code that are
// executed separately, use the Then method to append code pieces.
//
// When combined with subsequent method calls, a test case reads like English.
// For example, a test for the fact that "put x" puts "x" reads:
//
//     That("put x").Puts("x")
func That(lines ...string) Case {
	return Case{codes: []string{strings.Join(lines, "\n")}}
}

// Then returns a new Case that executes the given code in addition. Multiple
// arguments are joined with newlines.
func (c Case) Then(lines ...string) Case {
	c.codes = append(c.codes, strings.Join(lines, "\n"))
	return c
}

// DoesNothing returns t unchanged. It is useful to mark tests that don't have
// any side effects, for example:
//
//     That("nop").DoesNothing()
func (c Case) DoesNothing() Case {
	return c
}

// Puts returns an altered Case that requires the source code to produce the
// specified values in the value channel when evaluated.
func (c Case) Puts(vs ...interface{}) Case {
	c.want.ValueOut = vs
	return c
}

// Prints returns an altered Case that requires the source code to produce the
// specified output in the byte pipe when evaluated.
func (c Case) Prints(s string) Case {
	c.want.BytesOut = []byte(s)
	return c
}

// PrintsStderrWith returns an altered Case that requires the stderr output to
// contain the given text.
func (c Case) PrintsStderrWith(s string) Case {
	c.want.StderrOut = []byte(s)
	return c
}

// Throws returns an altered Case that requires the source code to throw an
// exception with the given reason. The reason supports special matcher values
// constructed by functions like ErrorWithMessage.
//
// If at least one stacktrace string is given, the exception must also have a
// stacktrace matching the given source fragments, frame by frame (innermost
// frame first). If no stacktrace string is given, the stack trace of the
// exception is not checked.
func (c Case) Throws(reason error, stacks ...string) Case {
	c.want.Exception = exc{reason, stacks}
	return c
}

// DoesNotCompile returns an altered Case that requires the source code to fail
// compilation.
func (c Case) DoesNotCompile() Case {
	c.want.CompilationError = anyError{}
	return c
}

// Test runs test cases. For each test case, a new Evaler is created with
// NewEvaler.
func Test(t *testing.T, tests ...Case) {
	t.Helper()
	TestWithSetup(t, func(*eval.Evaler) {}, tests...)
}

// TestWithSetup runs test cases. For each test case, a new Evaler is created
// with NewEvaler and passed to the setup function.
func TestWithSetup(t *testing.T, setup func(*eval.Evaler), tests ...Case) {
	t.Helper()
	for _, tt := range tests {
		t.Run(strings.Join(tt.codes, "\n"), func(t *testing.T) {
			t.Helper()
			ev := eval.NewEvaler()
			setup(ev)

			r := evalAndCollect(t, ev, tt.codes)

			if !matchOut(tt.want.ValueOut, r.ValueOut) {
				t.Errorf("got value out %v, want %v",
					reprs(r.ValueOut), reprs(tt.want.ValueOut))
			}
			if !bytes.Equal(tt.want.BytesOut, r.BytesOut) {
				t.Errorf("got bytes out %q, want %q", r.BytesOut, tt.want.BytesOut)
			}
			if !bytes.Contains(r.StderrOut, tt.want.StderrOut) {
				t.Errorf("got stderr out %q, want %q", r.StderrOut, tt.want.StderrOut)
			}
			if !matchErr(tt.want.CompilationError, r.CompilationError) {
				t.Errorf("got compilation error %v, want %v",
					r.CompilationError, tt.want.CompilationError)
			}
			if !matchErr(tt.want.Exception, r.Exception) {
				t.Errorf("unexpected exception")
				if exc, ok := r.Exception.(eval.Exception); ok {
					// For an eval.Exception report the type of the underlying error.
					t.Logf("got: %T: %v", exc.Reason(), exc)
					t.Logf("stack trace: %#v", getStackTexts(exc.StackTrace()))
				} else {
					t.Logf("got: %T: %v", r.Exception, r.Exception)
				}
				t.Errorf("want: %v", tt.want.Exception)
			}
		})
	}
}

func evalAndCollect(t *testing.T, ev *eval.Evaler, texts []string) result {
	var r result

	port1, collect1 := capturePort()
	port2, collect2 := capturePort()
	ports := []*eval.Port{eval.DummyInputPort, port1, port2}

	for _, text := range texts {
		err := ev.Eval(parse.Source{Name: "[test]", Code: text},
			eval.EvalCfg{Ports: ports, Interrupt: eval.ListenInterrupts})

		if parse.GetError(err) != nil {
			t.Fatalf("Parse(%q) error: %s", text, err)
		} else if eval.GetCompilationError(err) != nil {
			// NOTE: If multiple code pieces have compilation errors, only the
			// last one compilation error is saved.
			r.CompilationError = err
		} else if err != nil {
			// NOTE: If multiple code pieces throw exceptions, only the last one
			// is saved.
			r.Exception = err
		}
	}

	r.ValueOut, r.BytesOut = collect1()
	_, r.StderrOut = collect2()
	return r
}

// Like eval.CapturePort, but captures values and bytes separately. Also panics
// if it cannot create a pipe.
func capturePort() (*eval.Port, func() ([]interface{}, []byte)) {
	var values []interface{}
	var bytes []byte
	port, done, err := eval.PipePort(
		func(ch <-chan interface{}) {
			for v := range ch {
				values = append(values, v)
			}
		},
		func(r *os.File) {
			bytes = testutil.MustReadAllAndClose(r)
		})
	if err != nil {
		panic(err)
	}
	return port, func() ([]interface{}, []byte) {
		done()
		return values, bytes
	}
}

func matchOut(want, got []interface{}) bool {
	if len(got) != len(want) {
		return false
	}
	for i := range got {
		if !match(got[i], want[i]) {
			return false
		}
	}
	return true
}

func match(got, want interface{}) bool {
	switch got := got.(type) {
	case float64:
		// Special-case float64 to correctly handle NaN and support
		// approximate comparison.
		switch want := want.(type) {
		case float64:
			return matchFloat64(got, want, 0)
		case Approximately:
			return matchFloat64(got, want.F, ApproximatelyThreshold)
		}
	case string:
		switch want := want.(type) {
		case MatchingRegexp:
			return matchRegexp(want.Pattern, got)
		}
	}
	return vals.Equal(got, want)
}

func reprs(values []interface{}) []string {
	s := make([]string, len(values))
	for i, v := range values {
		s[i] = vals.Repr(v, vals.NoPretty)
	}
	return s
}

func matchErr(want, got error) bool {
	if want == nil {
		return got == nil
	}
	if matcher, ok := want.(errorMatcher); ok {
		return matcher.matchError(got)
	}
	return reflect.DeepEqual(want, got)
}
<commit_msg>pkg/eval/evaltest: Fix package godoc.<commit_after>// Package evaltest provides a framework for testing Elvish script.
//
// The entry point for the framework is the Test function, which accepts a
// *testing.T and any number of test cases.
//
// Test cases are constructed using the That function, followed by method calls
// that add additional information to it.
//
// Example:
//
//     Test(t,
//         That("put x").Puts("x"),
//         That("echo x").Prints("x\n"))
//
// If some setup is needed, use the TestWithSetup function instead.
package evaltest

import (
	"bytes"
	"os"
	"reflect"
	"strings"
	"testing"

	"src.elv.sh/pkg/eval"
	"src.elv.sh/pkg/eval/vals"
	"src.elv.sh/pkg/parse"
	"src.elv.sh/pkg/testutil"
)

// Case is a test case that can be used in Test.
type Case struct {
	codes []string
	want  result
}

type result struct {
	ValueOut  []interface{}
	BytesOut  []byte
	StderrOut []byte

	CompilationError error
	Exception        error
}

// That returns a new Case with the specified source code. Multiple arguments
// are joined with newlines. To specify multiple pieces of code that are
// executed separately, use the Then method to append code pieces.
//
// When combined with subsequent method calls, a test case reads like English.
// For example, a test for the fact that "put x" puts "x" reads:
//
//     That("put x").Puts("x")
func That(lines ...string) Case {
	return Case{codes: []string{strings.Join(lines, "\n")}}
}

// Then returns a new Case that executes the given code in addition. Multiple
// arguments are joined with newlines.
func (c Case) Then(lines ...string) Case {
	c.codes = append(c.codes, strings.Join(lines, "\n"))
	return c
}

// DoesNothing returns t unchanged. It is useful to mark tests that don't have
// any side effects, for example:
//
//     That("nop").DoesNothing()
func (c Case) DoesNothing() Case {
	return c
}

// Puts returns an altered Case that requires the source code to produce the
// specified values in the value channel when evaluated.
func (c Case) Puts(vs ...interface{}) Case {
	c.want.ValueOut = vs
	return c
}

// Prints returns an altered Case that requires the source code to produce the
// specified output in the byte pipe when evaluated.
func (c Case) Prints(s string) Case {
	c.want.BytesOut = []byte(s)
	return c
}

// PrintsStderrWith returns an altered Case that requires the stderr output to
// contain the given text.
func (c Case) PrintsStderrWith(s string) Case {
	c.want.StderrOut = []byte(s)
	return c
}

// Throws returns an altered Case that requires the source code to throw an
// exception with the given reason. The reason supports special matcher values
// constructed by functions like ErrorWithMessage.
//
// If at least one stacktrace string is given, the exception must also have a
// stacktrace matching the given source fragments, frame by frame (innermost
// frame first). If no stacktrace string is given, the stack trace of the
// exception is not checked.
func (c Case) Throws(reason error, stacks ...string) Case {
	c.want.Exception = exc{reason, stacks}
	return c
}

// DoesNotCompile returns an altered Case that requires the source code to fail
// compilation.
func (c Case) DoesNotCompile() Case {
	c.want.CompilationError = anyError{}
	return c
}

// Test runs test cases. For each test case, a new Evaler is created with
// NewEvaler.
func Test(t *testing.T, tests ...Case) {
	t.Helper()
	TestWithSetup(t, func(*eval.Evaler) {}, tests...)
}

// TestWithSetup runs test cases. For each test case, a new Evaler is created
// with NewEvaler and passed to the setup function.
func TestWithSetup(t *testing.T, setup func(*eval.Evaler), tests ...Case) {
	t.Helper()
	for _, tt := range tests {
		t.Run(strings.Join(tt.codes, "\n"), func(t *testing.T) {
			t.Helper()
			ev := eval.NewEvaler()
			setup(ev)

			r := evalAndCollect(t, ev, tt.codes)

			if !matchOut(tt.want.ValueOut, r.ValueOut) {
				t.Errorf("got value out %v, want %v",
					reprs(r.ValueOut), reprs(tt.want.ValueOut))
			}
			if !bytes.Equal(tt.want.BytesOut, r.BytesOut) {
				t.Errorf("got bytes out %q, want %q", r.BytesOut, tt.want.BytesOut)
			}
			if !bytes.Contains(r.StderrOut, tt.want.StderrOut) {
				t.Errorf("got stderr out %q, want %q", r.StderrOut, tt.want.StderrOut)
			}
			if !matchErr(tt.want.CompilationError, r.CompilationError) {
				t.Errorf("got compilation error %v, want %v",
					r.CompilationError, tt.want.CompilationError)
			}
			if !matchErr(tt.want.Exception, r.Exception) {
				t.Errorf("unexpected exception")
				if exc, ok := r.Exception.(eval.Exception); ok {
					// For an eval.Exception report the type of the underlying error.
					t.Logf("got: %T: %v", exc.Reason(), exc)
					t.Logf("stack trace: %#v", getStackTexts(exc.StackTrace()))
				} else {
					t.Logf("got: %T: %v", r.Exception, r.Exception)
				}
				t.Errorf("want: %v", tt.want.Exception)
			}
		})
	}
}

func evalAndCollect(t *testing.T, ev *eval.Evaler, texts []string) result {
	var r result

	port1, collect1 := capturePort()
	port2, collect2 := capturePort()
	ports := []*eval.Port{eval.DummyInputPort, port1, port2}

	for _, text := range texts {
		err := ev.Eval(parse.Source{Name: "[test]", Code: text},
			eval.EvalCfg{Ports: ports, Interrupt: eval.ListenInterrupts})

		if parse.GetError(err) != nil {
			t.Fatalf("Parse(%q) error: %s", text, err)
		} else if eval.GetCompilationError(err) != nil {
			// NOTE: If multiple code pieces have compilation errors, only the
			// last one compilation error is saved.
			r.CompilationError = err
		} else if err != nil {
			// NOTE: If multiple code pieces throw exceptions, only the last one
			// is saved.
			r.Exception = err
		}
	}

	r.ValueOut, r.BytesOut = collect1()
	_, r.StderrOut = collect2()
	return r
}

// Like eval.CapturePort, but captures values and bytes separately. Also panics
// if it cannot create a pipe.
func capturePort() (*eval.Port, func() ([]interface{}, []byte)) {
	var values []interface{}
	var bytes []byte
	port, done, err := eval.PipePort(
		func(ch <-chan interface{}) {
			for v := range ch {
				values = append(values, v)
			}
		},
		func(r *os.File) {
			bytes = testutil.MustReadAllAndClose(r)
		})
	if err != nil {
		panic(err)
	}
	return port, func() ([]interface{}, []byte) {
		done()
		return values, bytes
	}
}

func matchOut(want, got []interface{}) bool {
	if len(got) != len(want) {
		return false
	}
	for i := range got {
		if !match(got[i], want[i]) {
			return false
		}
	}
	return true
}

func match(got, want interface{}) bool {
	switch got := got.(type) {
	case float64:
		// Special-case float64 to correctly handle NaN and support
		// approximate comparison.
		switch want := want.(type) {
		case float64:
			return matchFloat64(got, want, 0)
		case Approximately:
			return matchFloat64(got, want.F, ApproximatelyThreshold)
		}
	case string:
		switch want := want.(type) {
		case MatchingRegexp:
			return matchRegexp(want.Pattern, got)
		}
	}
	return vals.Equal(got, want)
}

func reprs(values []interface{}) []string {
	s := make([]string, len(values))
	for i, v := range values {
		s[i] = vals.Repr(v, vals.NoPretty)
	}
	return s
}

func matchErr(want, got error) bool {
	if want == nil {
		return got == nil
	}
	if matcher, ok := want.(errorMatcher); ok {
		return matcher.matchError(got)
	}
	return reflect.DeepEqual(want, got)
}
<|endoftext|>


<commit_before>package status

import (
	"crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/go-git/go-git/v5"

	"github.com/werf/logboek"

	"github.com/werf/werf/pkg/path_matcher"
	"github.com/werf/werf/pkg/util"
)

type Result struct {
	repository             *git.Repository
	repositoryAbsFilepath  string
	repositoryFullFilepath string
	fileStatusList         git.Status
	submoduleResults       []*SubmoduleResult
}

type SubmoduleResult struct {
	*Result
	isNotInitialized bool
	isNotClean       bool
	currentCommit    string
}

func (r *Result) Status(pathMatcher path_matcher.PathMatcher) (*Result, error) {
	res := &Result{
		repository:             r.repository,
		repositoryAbsFilepath:  r.repositoryAbsFilepath,
		repositoryFullFilepath: r.repositoryFullFilepath,
		fileStatusList:         git.Status{},
		submoduleResults:       []*SubmoduleResult{},
	}

	for fileStatusPath, fileStatus := range r.fileStatusList {
		fileStatusFilepath := filepath.FromSlash(fileStatusPath)
		fileStatusFullFilepath := filepath.Join(r.repositoryFullFilepath, fileStatusFilepath)

		if pathMatcher.MatchPath(fileStatusFullFilepath) {
			res.fileStatusList[fileStatusPath] = fileStatus

			if debugProcess() {
				logboek.Debug.LogF(
					"File was added:         %s (worktree: %s, staging: %s)\n",
					fileStatusFullFilepath,
					fileStatusMapping[rune(fileStatus.Worktree)],
					fileStatusMapping[rune(fileStatus.Staging)],
				)
			}
		}
	}

	for _, submoduleResult := range r.submoduleResults {
		isMatched, shouldGoThrough := pathMatcher.ProcessDirOrSubmodulePath(submoduleResult.repositoryFullFilepath)
		if isMatched || shouldGoThrough {
			if debugProcess() {
				logboek.Debug.LogF("Submodule was checking: %s\n", submoduleResult.repositoryFullFilepath)
			}

			if submoduleResult.isNotInitialized {
				res.submoduleResults = append(res.submoduleResults, submoduleResult)

				if debugProcess() {
					logboek.Debug.LogFWithCustomStyle(
						logboek.StyleByName(logboek.FailStyleName),
						"Submodule is not initialized: path %s will be added to checksum\n",
						submoduleResult.repositoryFullFilepath,
					)
				}
				continue
			}

			if submoduleResult.isNotClean {
				if debugProcess() {
					logboek.Debug.LogFWithCustomStyle(
						logboek.StyleByName(logboek.FailStyleName),
						"Submodule is not clean: current commit %s will be added to checksum\n",
						submoduleResult.currentCommit,
					)
				}
			}

			newResult, err := submoduleResult.Status(pathMatcher)
			if err != nil {
				return nil, err
			}

			newSubmoduleResult := &SubmoduleResult{
				Result:           newResult,
				isNotInitialized: false,
				isNotClean:       submoduleResult.isNotClean,
				currentCommit:    submoduleResult.currentCommit,
			}

			if !newSubmoduleResult.isEmpty() {
				res.submoduleResults = append(res.submoduleResults, newSubmoduleResult)
			}
		}
	}

	return res, nil
}

func (r *Result) Checksum() (string, error) {
	if r.IsEmpty() {
		return "", nil
	}

	h := sha256.New()

	var fileStatusPathList []string
	for fileStatusPath := range r.fileStatusList {
		fileStatusPathList = append(fileStatusPathList, fileStatusPath)
	}

	fileModeAndDataFunc := func(fileStatusAbsFilepath string) (string, string, error) {
		stat, err := os.Lstat(fileStatusAbsFilepath)
		if err != nil {
			return "", "", fmt.Errorf("os stat %s failed: %s", fileStatusAbsFilepath, err)
		}

		dataH := sha256.New()
		if stat.Mode()&os.ModeSymlink != 0 {
			linkTo, err := os.Readlink(fileStatusAbsFilepath)
			if err != nil {
				return "", "", fmt.Errorf("os read link %s failed: %s", fileStatusAbsFilepath, err)
			}

			dataH.Write([]byte(linkTo))
		} else {
			data, err := ioutil.ReadFile(fileStatusAbsFilepath)
			if err != nil {
				return "", "", fmt.Errorf("os read file %s failed: %s", fileStatusAbsFilepath, err)
			}

			dataH.Write(data)
		}

		return stat.Mode().String(), fmt.Sprintf("%x", dataH.Sum(nil)), nil
	}

	sort.Strings(fileStatusPathList)

	for _, fileStatusPath := range fileStatusPathList {
		fileStatus := r.fileStatusList[fileStatusPath]
		fileStatusFilepath := filepath.FromSlash(fileStatusPath)
		fileStatusFullFilepath := filepath.Join(r.repositoryFullFilepath, fileStatusFilepath)
		fileStatusAbsFilepath := filepath.Join(r.repositoryAbsFilepath, fileStatusFilepath)

		var modeAndFileDataShouldBeAdded bool
		var fileStatusToAdd git.StatusCode
		var extraToAdd string

		switch fileStatus.Staging {
		case git.Untracked:
			if fileStatus.Worktree == git.Untracked {
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Untracked
			} else {
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Unmodified:
			switch fileStatus.Worktree {
			case git.Modified:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Added:
			switch fileStatus.Worktree {
			case git.Unmodified, git.Modified:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Added
			case git.Deleted:
				continue
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Renamed:
			switch fileStatus.Worktree {
			case git.Unmodified:
				extraToAdd = fileStatus.Extra
				fileStatusToAdd = git.Renamed
			case git.Modified:
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Copied:
			switch fileStatus.Worktree {
			case git.Unmodified:
				extraToAdd = fileStatus.Extra
				fileStatusToAdd = git.Copied
			case git.Modified:
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Deleted:
			switch fileStatus.Worktree {
			case git.Untracked:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Untracked
			case git.Unmodified:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.UpdatedButUnmerged:
			exist, err := util.FileExists(fileStatusAbsFilepath)
			if err != nil {
				panic(err)
			}

			fileStatusToAdd = git.UpdatedButUnmerged
			if exist {
				modeAndFileDataShouldBeAdded = true
			}
		default:
			panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
		}

		fileStatusFullPath := filepath.ToSlash(fileStatusFullFilepath)

		var args []string
		args = append(args, fileStatusFullPath)

		if extraToAdd != "" {
			args = append(args, extraToAdd)
		}

		fileStatusName := fileStatusMapping[rune(fileStatusToAdd)]
		if fileStatusName != "" {
			args = append(args, fileStatusName)
		}

		if modeAndFileDataShouldBeAdded {
			mode, data, err := fileModeAndDataFunc(fileStatusAbsFilepath)
			if err != nil {
				return "", err
			}

			args = append(args, mode, data)
		}

		logboek.Debug.LogF("Args was added: %v\n", args)
		logboek.Debug.LogF("  worktree %s  staging %s  result %s\n", fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)], fileStatusMapping[rune(fileStatusToAdd)])
		h.Write([]byte(strings.Join(args, "🐜")))
	}

	sort.Slice(r.submoduleResults, func(i, j int) bool {
		return r.submoduleResults[i].repositoryFullFilepath < r.submoduleResults[j].repositoryFullFilepath
	})

	for _, sr := range r.submoduleResults {
		logboek.Debug.LogOptionalLn()
		logBlockMsg := fmt.Sprintf("submodule %s", sr.repositoryFullFilepath)
		if err := logboek.Debug.LogBlock(logBlockMsg, logboek.LevelLogBlockOptions{}, func() error {
			var srChecksumArgs []string

			srChecksumArgs = append(srChecksumArgs, sr.repositoryFullFilepath)

			if sr.isNotInitialized {
				srChecksumArgs = append(srChecksumArgs, "isNotInitialized")
				return nil
			} else {
				if sr.isNotClean {
					srChecksumArgs = append(srChecksumArgs, "isNotClean")
					srChecksumArgs = append(srChecksumArgs, sr.currentCommit)
				}

				srChecksum, err := sr.Checksum()
				if err != nil {
					return err
				}

				if srChecksum != "" {
					srChecksumArgs = append(srChecksumArgs, srChecksum)
				}
			}

			logboek.Debug.LogF("Args was added: %v\n", srChecksumArgs)
			h.Write([]byte(strings.Join(srChecksumArgs, "🐜")))

			return nil
		}); err != nil {
			return "", fmt.Errorf("submodule %s checksum failed: %s", sr.repositoryFullFilepath, err)
		}
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func (r *Result) IsEmpty() bool {
	return len(r.fileStatusList) == 0 && len(r.submoduleResults) == 0
}

func (sr *SubmoduleResult) isEmpty() bool {
	return sr.Result.IsEmpty() && !sr.isNotClean && !sr.isNotInitialized
}
<commit_msg>[dockerfile] Fix handling specific status of a local uncommitted file<commit_after>package status

import (
	"crypto/sha256"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/go-git/go-git/v5"

	"github.com/werf/logboek"

	"github.com/werf/werf/pkg/path_matcher"
	"github.com/werf/werf/pkg/util"
)

type Result struct {
	repository             *git.Repository
	repositoryAbsFilepath  string
	repositoryFullFilepath string
	fileStatusList         git.Status
	submoduleResults       []*SubmoduleResult
}

type SubmoduleResult struct {
	*Result
	isNotInitialized bool
	isNotClean       bool
	currentCommit    string
}

func (r *Result) Status(pathMatcher path_matcher.PathMatcher) (*Result, error) {
	res := &Result{
		repository:             r.repository,
		repositoryAbsFilepath:  r.repositoryAbsFilepath,
		repositoryFullFilepath: r.repositoryFullFilepath,
		fileStatusList:         git.Status{},
		submoduleResults:       []*SubmoduleResult{},
	}

	for fileStatusPath, fileStatus := range r.fileStatusList {
		fileStatusFilepath := filepath.FromSlash(fileStatusPath)
		fileStatusFullFilepath := filepath.Join(r.repositoryFullFilepath, fileStatusFilepath)

		if pathMatcher.MatchPath(fileStatusFullFilepath) {
			res.fileStatusList[fileStatusPath] = fileStatus

			if debugProcess() {
				logboek.Debug.LogF(
					"File was added:         %s (worktree: %s, staging: %s)\n",
					fileStatusFullFilepath,
					fileStatusMapping[rune(fileStatus.Worktree)],
					fileStatusMapping[rune(fileStatus.Staging)],
				)
			}
		}
	}

	for _, submoduleResult := range r.submoduleResults {
		isMatched, shouldGoThrough := pathMatcher.ProcessDirOrSubmodulePath(submoduleResult.repositoryFullFilepath)
		if isMatched || shouldGoThrough {
			if debugProcess() {
				logboek.Debug.LogF("Submodule was checking: %s\n", submoduleResult.repositoryFullFilepath)
			}

			if submoduleResult.isNotInitialized {
				res.submoduleResults = append(res.submoduleResults, submoduleResult)

				if debugProcess() {
					logboek.Debug.LogFWithCustomStyle(
						logboek.StyleByName(logboek.FailStyleName),
						"Submodule is not initialized: path %s will be added to checksum\n",
						submoduleResult.repositoryFullFilepath,
					)
				}
				continue
			}

			if submoduleResult.isNotClean {
				if debugProcess() {
					logboek.Debug.LogFWithCustomStyle(
						logboek.StyleByName(logboek.FailStyleName),
						"Submodule is not clean: current commit %s will be added to checksum\n",
						submoduleResult.currentCommit,
					)
				}
			}

			newResult, err := submoduleResult.Status(pathMatcher)
			if err != nil {
				return nil, err
			}

			newSubmoduleResult := &SubmoduleResult{
				Result:           newResult,
				isNotInitialized: false,
				isNotClean:       submoduleResult.isNotClean,
				currentCommit:    submoduleResult.currentCommit,
			}

			if !newSubmoduleResult.isEmpty() {
				res.submoduleResults = append(res.submoduleResults, newSubmoduleResult)
			}
		}
	}

	return res, nil
}

func (r *Result) Checksum() (string, error) {
	if r.IsEmpty() {
		return "", nil
	}

	h := sha256.New()

	var fileStatusPathList []string
	for fileStatusPath := range r.fileStatusList {
		fileStatusPathList = append(fileStatusPathList, fileStatusPath)
	}

	fileModeAndDataFunc := func(fileStatusAbsFilepath string) (string, string, error) {
		stat, err := os.Lstat(fileStatusAbsFilepath)
		if err != nil {
			return "", "", fmt.Errorf("os stat %s failed: %s", fileStatusAbsFilepath, err)
		}

		dataH := sha256.New()
		if stat.Mode()&os.ModeSymlink != 0 {
			linkTo, err := os.Readlink(fileStatusAbsFilepath)
			if err != nil {
				return "", "", fmt.Errorf("os read link %s failed: %s", fileStatusAbsFilepath, err)
			}

			dataH.Write([]byte(linkTo))
		} else {
			data, err := ioutil.ReadFile(fileStatusAbsFilepath)
			if err != nil {
				return "", "", fmt.Errorf("os read file %s failed: %s", fileStatusAbsFilepath, err)
			}

			dataH.Write(data)
		}

		return stat.Mode().String(), fmt.Sprintf("%x", dataH.Sum(nil)), nil
	}

	sort.Strings(fileStatusPathList)

	for _, fileStatusPath := range fileStatusPathList {
		fileStatus := r.fileStatusList[fileStatusPath]
		fileStatusFilepath := filepath.FromSlash(fileStatusPath)
		fileStatusFullFilepath := filepath.Join(r.repositoryFullFilepath, fileStatusFilepath)
		fileStatusAbsFilepath := filepath.Join(r.repositoryAbsFilepath, fileStatusFilepath)

		var modeAndFileDataShouldBeAdded bool
		var fileStatusToAdd git.StatusCode
		var extraToAdd string

		switch fileStatus.Staging {
		case git.Untracked:
			if fileStatus.Worktree == git.Untracked {
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Untracked
			} else {
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Unmodified:
			switch fileStatus.Worktree {
			case git.Modified:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Added, git.Modified:
			switch fileStatus.Worktree {
			case git.Unmodified, git.Modified:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Added
			case git.Deleted:
				continue
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Renamed:
			switch fileStatus.Worktree {
			case git.Unmodified:
				extraToAdd = fileStatus.Extra
				fileStatusToAdd = git.Renamed
			case git.Modified:
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Copied:
			switch fileStatus.Worktree {
			case git.Unmodified:
				extraToAdd = fileStatus.Extra
				fileStatusToAdd = git.Copied
			case git.Modified:
				fileStatusToAdd = git.Modified
			case git.Deleted:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.Deleted:
			switch fileStatus.Worktree {
			case git.Untracked:
				modeAndFileDataShouldBeAdded = true
				fileStatusToAdd = git.Untracked
			case git.Unmodified:
				fileStatusToAdd = git.Deleted
			default:
				panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
			}
		case git.UpdatedButUnmerged:
			exist, err := util.FileExists(fileStatusAbsFilepath)
			if err != nil {
				panic(err)
			}

			fileStatusToAdd = git.UpdatedButUnmerged
			if exist {
				modeAndFileDataShouldBeAdded = true
			}
		default:
			panic(fmt.Sprintf("unexpected condition (path %s worktree %s, staging %s)", fileStatusAbsFilepath, fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)]))
		}

		fileStatusFullPath := filepath.ToSlash(fileStatusFullFilepath)

		var args []string
		args = append(args, fileStatusFullPath)

		if extraToAdd != "" {
			args = append(args, extraToAdd)
		}

		fileStatusName := fileStatusMapping[rune(fileStatusToAdd)]
		if fileStatusName != "" {
			args = append(args, fileStatusName)
		}

		if modeAndFileDataShouldBeAdded {
			mode, data, err := fileModeAndDataFunc(fileStatusAbsFilepath)
			if err != nil {
				return "", err
			}

			args = append(args, mode, data)
		}

		logboek.Debug.LogF("Args was added: %v\n", args)
		logboek.Debug.LogF("  worktree %s  staging %s  result %s\n", fileStatusMapping[rune(fileStatus.Worktree)], fileStatusMapping[rune(fileStatus.Staging)], fileStatusMapping[rune(fileStatusToAdd)])
		h.Write([]byte(strings.Join(args, "🐜")))
	}

	sort.Slice(r.submoduleResults, func(i, j int) bool {
		return r.submoduleResults[i].repositoryFullFilepath < r.submoduleResults[j].repositoryFullFilepath
	})

	for _, sr := range r.submoduleResults {
		logboek.Debug.LogOptionalLn()
		logBlockMsg := fmt.Sprintf("submodule %s", sr.repositoryFullFilepath)
		if err := logboek.Debug.LogBlock(logBlockMsg, logboek.LevelLogBlockOptions{}, func() error {
			var srChecksumArgs []string

			srChecksumArgs = append(srChecksumArgs, sr.repositoryFullFilepath)

			if sr.isNotInitialized {
				srChecksumArgs = append(srChecksumArgs, "isNotInitialized")
				return nil
			} else {
				if sr.isNotClean {
					srChecksumArgs = append(srChecksumArgs, "isNotClean")
					srChecksumArgs = append(srChecksumArgs, sr.currentCommit)
				}

				srChecksum, err := sr.Checksum()
				if err != nil {
					return err
				}

				if srChecksum != "" {
					srChecksumArgs = append(srChecksumArgs, srChecksum)
				}
			}

			logboek.Debug.LogF("Args was added: %v\n", srChecksumArgs)
			h.Write([]byte(strings.Join(srChecksumArgs, "🐜")))

			return nil
		}); err != nil {
			return "", fmt.Errorf("submodule %s checksum failed: %s", sr.repositoryFullFilepath, err)
		}
	}

	return fmt.Sprintf("%x", h.Sum(nil)), nil
}

func (r *Result) IsEmpty() bool {
	return len(r.fileStatusList) == 0 && len(r.submoduleResults) == 0
}

func (sr *SubmoduleResult) isEmpty() bool {
	return sr.Result.IsEmpty() && !sr.isNotClean && !sr.isNotInitialized
}
<|endoftext|>


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

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}

	// Force network usage, to verify the epoll (or whatever) fd
	// doesn't leak to the child,
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v; output %q", err, bs)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		// Now verify that there are no other open fds.
		var files []*os.File
		for wantfd := os.Stderr.Fd() + 2; wantfd <= 100; wantfd++ {
			f, err := os.Open(os.Args[0])
			if err != nil {
				fmt.Printf("error opening file with expected fd %d: %v", wantfd, err)
				fmt.Println(Command("lsof", "-p", fmt.Sprint(os.Getpid())).CombinedOutput())
				os.Exit(1)
			}
			if got := f.Fd(); got != wantfd {
				fmt.Printf("leaked parent file. fd = %d; want %d", got, wantfd)
				os.Exit(1)
			}
			files = append(files, f)
		}
		for _, f := range files {
			f.Close()
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<commit_msg>os/exec: put the print where it will help<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}

	// Force network usage, to verify the epoll (or whatever) fd
	// doesn't leak to the child,
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v; output %q", err, bs)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		// Now verify that there are no other open fds.
		var files []*os.File
		for wantfd := os.Stderr.Fd() + 2; wantfd <= 100; wantfd++ {
			f, err := os.Open(os.Args[0])
			if err != nil {
				fmt.Printf("error opening file with expected fd %d: %v", wantfd, err)
				os.Exit(1)
			}
			if got := f.Fd(); got != wantfd {
				fmt.Printf("leaked parent file. fd = %d; want %d", got, wantfd)
				fmt.Println(Command("lsof", "-p", fmt.Sprint(os.Getpid())).CombinedOutput())
				os.Exit(1)
			}
			files = append(files, f)
		}
		for _, f := range files {
			f.Close()
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<|endoftext|>


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

package runtime_test

import (
	"runtime"
	"testing"
)

var stop = make(chan bool, 1)

func perpetuumMobile() {
	select {
	case <-stop:
	default:
		go perpetuumMobile()
	}
}

func TestStopTheWorldDeadlock(t *testing.T) {
	if testing.Short() {
		t.Logf("skipping during short test")
		return
	}
	runtime.GOMAXPROCS(3)
	compl := make(chan int, 1)
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GC()
		}
		compl <- 0
	}()
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GOMAXPROCS(3)
		}
	}()
	go perpetuumMobile()
	<-compl
	stop <- true
}
<commit_msg>runtime: reset GOMAXPROCS during tests<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime_test

import (
	"runtime"
	"testing"
)

var stop = make(chan bool, 1)

func perpetuumMobile() {
	select {
	case <-stop:
	default:
		go perpetuumMobile()
	}
}

func TestStopTheWorldDeadlock(t *testing.T) {
	if testing.Short() {
		t.Logf("skipping during short test")
		return
	}
	maxprocs := runtime.GOMAXPROCS(3)
	compl := make(chan bool, 2)
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GC()
		}
		compl <- true
	}()
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GOMAXPROCS(3)
		}
		compl <- true
	}()
	go perpetuumMobile()
	<-compl
	<-compl
	stop <- true
	runtime.GOMAXPROCS(maxprocs)
}
<|endoftext|>


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

package testing

import (
	"flag"
	"fmt"
	"os"
	"time"
)

var matchBenchmarks = flag.String("benchmarks", "", "regular expression to select benchmarks to run")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type Benchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark Benchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least a second.
	for b.ns < 1e9 && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = 1e9 / int(b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		n = min(int(1.5*float(n)), 100*last)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	ns := b.nsPerOp()
	mb := ""
	if ns > 0 && b.bytes > 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", (float64(b.bytes)/1e6)/(float64(ns)/1e9))
	}
	fmt.Printf("%s\t%8d\t%10d ns/op%s\n", b.benchmark.Name, b.N, b.nsPerOp(), mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(benchmarks []Benchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	re, err := CompileRegexp(*matchBenchmarks)
	if err != "" {
		println("invalid regexp for -benchmarks:", err)
		os.Exit(1)
	}
	for _, Benchmark := range benchmarks {
		if !re.MatchString(Benchmark.Name) {
			continue
		}
		b := &B{benchmark: Benchmark}
		b.run()
	}
}
<commit_msg>testing/benchmark: paranoia - make sure n always grows<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"flag"
	"fmt"
	"os"
	"time"
)

var matchBenchmarks = flag.String("benchmarks", "", "regular expression to select benchmarks to run")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type Benchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark Benchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

func max(x, y int) int {
	if x < y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least a second.
	for b.ns < 1e9 && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = 1e9 / int(b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		// Be sure to run at least one more than last time.
		n = max(min(n+n/2, 100*last), last+1)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	ns := b.nsPerOp()
	mb := ""
	if ns > 0 && b.bytes > 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", (float64(b.bytes)/1e6)/(float64(ns)/1e9))
	}
	fmt.Printf("%s\t%8d\t%10d ns/op%s\n", b.benchmark.Name, b.N, b.nsPerOp(), mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(benchmarks []Benchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	re, err := CompileRegexp(*matchBenchmarks)
	if err != "" {
		println("invalid regexp for -benchmarks:", err)
		os.Exit(1)
	}
	for _, Benchmark := range benchmarks {
		if !re.MatchString(Benchmark.Name) {
			continue
		}
		b := &B{benchmark: Benchmark}
		b.run()
	}
}
<|endoftext|>