<commit_before>package controller

import (
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/gophercloud/gophercloud/openstack/identity/v3/tokens"
	api_v1 "k8s.io/api/core/v1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	kubernetes_informers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/scheme"
	typedcorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/record"

	v1 "github.com/sapcc/kubernikus/pkg/apis/kubernikus/v1"
	helmutil "github.com/sapcc/kubernikus/pkg/client/helm"
	kube "github.com/sapcc/kubernikus/pkg/client/kubernetes"
	"github.com/sapcc/kubernikus/pkg/client/kubernikus"
	"github.com/sapcc/kubernikus/pkg/client/openstack"
	"github.com/sapcc/kubernikus/pkg/controller/certs"
	"github.com/sapcc/kubernikus/pkg/controller/config"
	"github.com/sapcc/kubernikus/pkg/controller/deorbit"
	"github.com/sapcc/kubernikus/pkg/controller/flight"
	"github.com/sapcc/kubernikus/pkg/controller/hammertime"
	"github.com/sapcc/kubernikus/pkg/controller/launch"
	"github.com/sapcc/kubernikus/pkg/controller/migration"
	"github.com/sapcc/kubernikus/pkg/controller/nodeobservatory"
	"github.com/sapcc/kubernikus/pkg/controller/routegc"
	"github.com/sapcc/kubernikus/pkg/controller/servicing"
	kubernikus_informers "github.com/sapcc/kubernikus/pkg/generated/informers/externalversions"
	_ "github.com/sapcc/kubernikus/pkg/util/workqueue/prometheus"
	"github.com/sapcc/kubernikus/pkg/version"
)

type KubernikusOperatorOptions struct {
	KubeConfig string
	Context    string

	ChartDirectory string

	AuthURL           string
	AuthUsername      string
	AuthPassword      string
	AuthDomain        string
	AuthProject       string
	AuthProjectDomain string
	Region            string

	KubernikusDomain    string
	KubernikusProjectID string
	KubernikusNetworkID string
	Namespace           string
	Controllers         []string
	MetricPort          int
	LogLevel            int

	NodeUpdateHoldoff time.Duration
}

type KubernikusOperator struct {
	config.Config
	config.Factories
	config.Clients
	Logger log.Logger
}

const (
	DEFAULT_RECONCILIATION = 5 * time.Minute
)

func NewKubernikusOperator(options *KubernikusOperatorOptions, logger log.Logger) (*KubernikusOperator, error) {
	var err error

	imageRegistry, err := version.NewImageRegistry(path.Join(options.ChartDirectory, "images.yaml"), options.Region)
	if err != nil {
		return nil, fmt.Errorf("Unable to initialize image registry: %s", err)
	}

	o := &KubernikusOperator{
		Config: config.Config{
			Openstack: config.OpenstackConfig{
				AuthURL:           options.AuthURL,
				AuthUsername:      options.AuthUsername,
				AuthPassword:      options.AuthPassword,
				AuthProject:       options.AuthProjectDomain,
				AuthProjectDomain: options.AuthProjectDomain,
			},
			Helm: config.HelmConfig{
				ChartDirectory: options.ChartDirectory,
			},
			Kubernikus: config.KubernikusConfig{
				Domain:      options.KubernikusDomain,
				Namespace:   options.Namespace,
				ProjectID:   options.KubernikusProjectID,
				NetworkID:   options.KubernikusNetworkID,
				Controllers: make(map[string]config.Controller),
			},
			Images: *imageRegistry,
		},
		Logger: logger,
	}

	o.Clients.Kubernetes, err = kube.NewClient(options.KubeConfig, options.Context, logger)

	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernetes clients: %s", err)
	}

	o.Clients.Kubernikus, err = kubernikus.NewClient(options.KubeConfig, options.Context)
	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernikus clients: %s", err)
	}

	config, err := kube.NewConfig(options.KubeConfig, options.Context)
	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernetes config: %s", err)
	}
	o.Clients.Helm, err = helmutil.NewClient(o.Clients.Kubernetes, config, logger)
	if err != nil {
		return nil, fmt.Errorf("Failed to create helm 2 client: %s", err)
	}
	o.Clients.Helm3, err = helmutil.NewClient3(options.Namespace, options.KubeConfig, options.Context, logger)
	if err != nil {
		return nil, fmt.Errorf("Failed to create helm 3 client: %s", err)
	}

	apiextensionsclientset, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return nil, fmt.Errorf("Failed to create apiextenstionsclient: %s", err)
	}

	if err := kube.EnsureCRD(apiextensionsclientset, logger); err != nil {
		return nil, fmt.Errorf("Couldn't create CRD: %s", err)
	}

	adminAuthOptions := &tokens.AuthOptions{
		IdentityEndpoint: options.AuthURL,
		Username:         options.AuthUsername,
		Password:         options.AuthPassword,
		DomainName:       options.AuthDomain,
		AllowReauth:      true,
		Scope: tokens.Scope{
			ProjectName: options.AuthProject,
			DomainName:  options.AuthProjectDomain,
		},
	}

	o.Factories.Kubernikus = kubernikus_informers.NewFilteredSharedInformerFactory(o.Clients.Kubernikus, DEFAULT_RECONCILIATION, options.Namespace, nil)
	o.Factories.Kubernetes = kubernetes_informers.NewFilteredSharedInformerFactory(o.Clients.Kubernetes, DEFAULT_RECONCILIATION, options.Namespace, nil)

	klusters := o.Factories.Kubernikus.Kubernikus().V1().Klusters().Informer()

	if options.AuthURL != "" {
		o.Factories.Openstack = openstack.NewSharedOpenstackClientFactory(o.Clients.Kubernetes, klusters, adminAuthOptions, logger)
	} else {
		o.Factories.Openstack = openstack.NotAvailableFactory{}
	}

	o.Clients.Satellites = kube.NewSharedClientFactory(o.Clients.Kubernetes, klusters, logger)

	o.Factories.NodesObservatory = nodeobservatory.NewInformerFactory(o.Factories.Kubernikus.Kubernikus().V1().Klusters(), o.Clients.Satellites, logger)

	// Add kubernikus types to the default Kubernetes Scheme so events can be
	// logged for those types.
	v1.AddToScheme(scheme.Scheme)
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartEventWatcher(func(e *api_v1.Event) {
		logger.Log(
			"controller", "operator",
			"resource", "event",
			"msg", e.Message,
			"reason", e.Reason,
			"type", e.Type,
			"kind", e.InvolvedObject.Kind,
			"namespace", e.InvolvedObject.Namespace,
			"name", e.InvolvedObject.Name,
			"v", 2)
	})

	eventBroadcaster.StartRecordingToSink(&typedcorev1.EventSinkImpl{Interface: o.Clients.Kubernetes.CoreV1().Events(o.Config.Kubernikus.Namespace)})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, api_v1.EventSource{Component: "operator"})

	for _, k := range options.Controllers {
		switch k {
		case "groundctl":
			o.Config.Kubernikus.Controllers["groundctl"] = NewGroundController(10, o.Factories, o.Clients, recorder, o.Config, logger)
		case "launchctl":
			o.Config.Kubernikus.Controllers["launchctl"] = launch.NewController(10, o.Factories, o.Clients, recorder, o.Config.Images, logger)
		case "routegc":
			o.Config.Kubernikus.Controllers["routegc"] = routegc.New(300*time.Second, o.Factories, logger)
		case "deorbiter":
			o.Config.Kubernikus.Controllers["deorbiter"] = deorbit.NewController(10, o.Factories, o.Clients, recorder, logger)
		case "flight":
			o.Config.Kubernikus.Controllers["flight"] = flight.NewController(10, o.Factories, o.Clients, recorder, logger)
		case "migration":
			o.Config.Kubernikus.Controllers["migration"] = migration.NewController(10, o.Factories, o.Clients, recorder, logger)
		case "hammertime":
			o.Config.Kubernikus.Controllers["hammertime"] = hammertime.New(10*time.Second, 20*time.Second, o.Factories, o.Clients, recorder, logger)
		case "servicing":
			o.Config.Kubernikus.Controllers["servicing"] = servicing.NewController(10, o.Factories, o.Clients, recorder, options.NodeUpdateHoldoff, logger)
		case "certs":
			o.Config.Kubernikus.Controllers["certs"] = certs.New(12*time.Hour, o.Factories, o.Config, o.Clients, logger)
		}
	}

	return o, err
}

func (o *KubernikusOperator) Run(stopCh <-chan struct{}, wg *sync.WaitGroup) {
	o.Logger.Log(
		"msg", "starting kubernikus operator",
		"namespace", o.Config.Kubernikus.Namespace,
		"version", version.GitCommit)

	kube.WaitForServer(o.Clients.Kubernetes, stopCh, o.Logger)

	o.Factories.Kubernikus.Start(stopCh)
	o.Factories.Kubernetes.Start(stopCh)
	o.Factories.NodesObservatory.Start(stopCh)

	o.Factories.Kubernikus.WaitForCacheSync(stopCh)
	o.Factories.Kubernetes.WaitForCacheSync(stopCh)

	o.Logger.Log("msg", "Cache primed")
	if _, found := o.Config.Kubernikus.Controllers["migration"]; found {
		migration.UpdateMigrationStatus(o.Clients.Kubernikus, o.Factories.Kubernikus.Kubernikus().V1().Klusters().Lister())
	}
	o.Logger.Log("msg", "Starting controllers")
	for _, controller := range o.Config.Kubernikus.Controllers {
		go controller.Run(stopCh, wg)
	}
}
<commit_msg>Reduce concurrency of migration controller<commit_after>package controller

import (
	"fmt"
	"path"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/gophercloud/gophercloud/openstack/identity/v3/tokens"
	api_v1 "k8s.io/api/core/v1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	kubernetes_informers "k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/scheme"
	typedcorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/record"

	v1 "github.com/sapcc/kubernikus/pkg/apis/kubernikus/v1"
	helmutil "github.com/sapcc/kubernikus/pkg/client/helm"
	kube "github.com/sapcc/kubernikus/pkg/client/kubernetes"
	"github.com/sapcc/kubernikus/pkg/client/kubernikus"
	"github.com/sapcc/kubernikus/pkg/client/openstack"
	"github.com/sapcc/kubernikus/pkg/controller/certs"
	"github.com/sapcc/kubernikus/pkg/controller/config"
	"github.com/sapcc/kubernikus/pkg/controller/deorbit"
	"github.com/sapcc/kubernikus/pkg/controller/flight"
	"github.com/sapcc/kubernikus/pkg/controller/hammertime"
	"github.com/sapcc/kubernikus/pkg/controller/launch"
	"github.com/sapcc/kubernikus/pkg/controller/migration"
	"github.com/sapcc/kubernikus/pkg/controller/nodeobservatory"
	"github.com/sapcc/kubernikus/pkg/controller/routegc"
	"github.com/sapcc/kubernikus/pkg/controller/servicing"
	kubernikus_informers "github.com/sapcc/kubernikus/pkg/generated/informers/externalversions"
	_ "github.com/sapcc/kubernikus/pkg/util/workqueue/prometheus"
	"github.com/sapcc/kubernikus/pkg/version"
)

type KubernikusOperatorOptions struct {
	KubeConfig string
	Context    string

	ChartDirectory string

	AuthURL           string
	AuthUsername      string
	AuthPassword      string
	AuthDomain        string
	AuthProject       string
	AuthProjectDomain string
	Region            string

	KubernikusDomain    string
	KubernikusProjectID string
	KubernikusNetworkID string
	Namespace           string
	Controllers         []string
	MetricPort          int
	LogLevel            int

	NodeUpdateHoldoff time.Duration
}

type KubernikusOperator struct {
	config.Config
	config.Factories
	config.Clients
	Logger log.Logger
}

const (
	DEFAULT_RECONCILIATION = 5 * time.Minute
)

func NewKubernikusOperator(options *KubernikusOperatorOptions, logger log.Logger) (*KubernikusOperator, error) {
	var err error

	imageRegistry, err := version.NewImageRegistry(path.Join(options.ChartDirectory, "images.yaml"), options.Region)
	if err != nil {
		return nil, fmt.Errorf("Unable to initialize image registry: %s", err)
	}

	o := &KubernikusOperator{
		Config: config.Config{
			Openstack: config.OpenstackConfig{
				AuthURL:           options.AuthURL,
				AuthUsername:      options.AuthUsername,
				AuthPassword:      options.AuthPassword,
				AuthProject:       options.AuthProjectDomain,
				AuthProjectDomain: options.AuthProjectDomain,
			},
			Helm: config.HelmConfig{
				ChartDirectory: options.ChartDirectory,
			},
			Kubernikus: config.KubernikusConfig{
				Domain:      options.KubernikusDomain,
				Namespace:   options.Namespace,
				ProjectID:   options.KubernikusProjectID,
				NetworkID:   options.KubernikusNetworkID,
				Controllers: make(map[string]config.Controller),
			},
			Images: *imageRegistry,
		},
		Logger: logger,
	}

	o.Clients.Kubernetes, err = kube.NewClient(options.KubeConfig, options.Context, logger)

	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernetes clients: %s", err)
	}

	o.Clients.Kubernikus, err = kubernikus.NewClient(options.KubeConfig, options.Context)
	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernikus clients: %s", err)
	}

	config, err := kube.NewConfig(options.KubeConfig, options.Context)
	if err != nil {
		return nil, fmt.Errorf("Failed to create kubernetes config: %s", err)
	}
	o.Clients.Helm, err = helmutil.NewClient(o.Clients.Kubernetes, config, logger)
	if err != nil {
		return nil, fmt.Errorf("Failed to create helm 2 client: %s", err)
	}
	o.Clients.Helm3, err = helmutil.NewClient3(options.Namespace, options.KubeConfig, options.Context, logger)
	if err != nil {
		return nil, fmt.Errorf("Failed to create helm 3 client: %s", err)
	}

	apiextensionsclientset, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return nil, fmt.Errorf("Failed to create apiextenstionsclient: %s", err)
	}

	if err := kube.EnsureCRD(apiextensionsclientset, logger); err != nil {
		return nil, fmt.Errorf("Couldn't create CRD: %s", err)
	}

	adminAuthOptions := &tokens.AuthOptions{
		IdentityEndpoint: options.AuthURL,
		Username:         options.AuthUsername,
		Password:         options.AuthPassword,
		DomainName:       options.AuthDomain,
		AllowReauth:      true,
		Scope: tokens.Scope{
			ProjectName: options.AuthProject,
			DomainName:  options.AuthProjectDomain,
		},
	}

	o.Factories.Kubernikus = kubernikus_informers.NewFilteredSharedInformerFactory(o.Clients.Kubernikus, DEFAULT_RECONCILIATION, options.Namespace, nil)
	o.Factories.Kubernetes = kubernetes_informers.NewFilteredSharedInformerFactory(o.Clients.Kubernetes, DEFAULT_RECONCILIATION, options.Namespace, nil)

	klusters := o.Factories.Kubernikus.Kubernikus().V1().Klusters().Informer()

	if options.AuthURL != "" {
		o.Factories.Openstack = openstack.NewSharedOpenstackClientFactory(o.Clients.Kubernetes, klusters, adminAuthOptions, logger)
	} else {
		o.Factories.Openstack = openstack.NotAvailableFactory{}
	}

	o.Clients.Satellites = kube.NewSharedClientFactory(o.Clients.Kubernetes, klusters, logger)

	o.Factories.NodesObservatory = nodeobservatory.NewInformerFactory(o.Factories.Kubernikus.Kubernikus().V1().Klusters(), o.Clients.Satellites, logger)

	// Add kubernikus types to the default Kubernetes Scheme so events can be
	// logged for those types.
	v1.AddToScheme(scheme.Scheme)
	eventBroadcaster := record.NewBroadcaster()
	eventBroadcaster.StartEventWatcher(func(e *api_v1.Event) {
		logger.Log(
			"controller", "operator",
			"resource", "event",
			"msg", e.Message,
			"reason", e.Reason,
			"type", e.Type,
			"kind", e.InvolvedObject.Kind,
			"namespace", e.InvolvedObject.Namespace,
			"name", e.InvolvedObject.Name,
			"v", 2)
	})

	eventBroadcaster.StartRecordingToSink(&typedcorev1.EventSinkImpl{Interface: o.Clients.Kubernetes.CoreV1().Events(o.Config.Kubernikus.Namespace)})
	recorder := eventBroadcaster.NewRecorder(scheme.Scheme, api_v1.EventSource{Component: "operator"})

	for _, k := range options.Controllers {
		switch k {
		case "groundctl":
			o.Config.Kubernikus.Controllers["groundctl"] = NewGroundController(10, o.Factories, o.Clients, recorder, o.Config, logger)
		case "launchctl":
			o.Config.Kubernikus.Controllers["launchctl"] = launch.NewController(10, o.Factories, o.Clients, recorder, o.Config.Images, logger)
		case "routegc":
			o.Config.Kubernikus.Controllers["routegc"] = routegc.New(300*time.Second, o.Factories, logger)
		case "deorbiter":
			o.Config.Kubernikus.Controllers["deorbiter"] = deorbit.NewController(10, o.Factories, o.Clients, recorder, logger)
		case "flight":
			o.Config.Kubernikus.Controllers["flight"] = flight.NewController(10, o.Factories, o.Clients, recorder, logger)
		case "migration":
			o.Config.Kubernikus.Controllers["migration"] = migration.NewController(3, o.Factories, o.Clients, recorder, logger)
		case "hammertime":
			o.Config.Kubernikus.Controllers["hammertime"] = hammertime.New(10*time.Second, 20*time.Second, o.Factories, o.Clients, recorder, logger)
		case "servicing":
			o.Config.Kubernikus.Controllers["servicing"] = servicing.NewController(10, o.Factories, o.Clients, recorder, options.NodeUpdateHoldoff, logger)
		case "certs":
			o.Config.Kubernikus.Controllers["certs"] = certs.New(12*time.Hour, o.Factories, o.Config, o.Clients, logger)
		}
	}

	return o, err
}

func (o *KubernikusOperator) Run(stopCh <-chan struct{}, wg *sync.WaitGroup) {
	o.Logger.Log(
		"msg", "starting kubernikus operator",
		"namespace", o.Config.Kubernikus.Namespace,
		"version", version.GitCommit)

	kube.WaitForServer(o.Clients.Kubernetes, stopCh, o.Logger)

	o.Factories.Kubernikus.Start(stopCh)
	o.Factories.Kubernetes.Start(stopCh)
	o.Factories.NodesObservatory.Start(stopCh)

	o.Factories.Kubernikus.WaitForCacheSync(stopCh)
	o.Factories.Kubernetes.WaitForCacheSync(stopCh)

	o.Logger.Log("msg", "Cache primed")
	if _, found := o.Config.Kubernikus.Controllers["migration"]; found {
		migration.UpdateMigrationStatus(o.Clients.Kubernikus, o.Factories.Kubernikus.Kubernikus().V1().Klusters().Lister())
	}
	o.Logger.Log("msg", "Starting controllers")
	for _, controller := range o.Config.Kubernikus.Controllers {
		go controller.Run(stopCh, wg)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 database

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/google/exposure-notifications-verification-server/pkg/api"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

func TestSubject(t *testing.T) {
	testDay, err := time.Parse("2006-01-02", "2020-07-07")
	if err != nil {
		t.Fatalf("test setup error: %v", err)
	}

	cases := []struct {
		Name  string
		Sub   string
		Want  *Subject
		Error string
	}{
		{
			Name: "normal parse",
			Sub:  "confirmed.2020-07-07",
			Want: &Subject{
				TestType:    "confirmed",
				SymptomDate: &testDay,
			},
		},
		{
			Name: "no date",
			Sub:  "confirmed.",
			Want: &Subject{
				TestType:    "confirmed",
				SymptomDate: nil,
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			got, err := ParseSubject(tc.Sub)
			if err != nil {
				if tc.Error == "" {
					t.Fatalf("unexpected error: %v", err)
				} else if !strings.Contains(err.Error(), tc.Error) {
					t.Fatalf("wrong error, want: %v got: %v", tc.Error, err.Error())
				}
			} else if tc.Error != "" {
				t.Fatalf("wanted error: %v got: nil", tc.Error)
			}

			if diff := cmp.Diff(tc.Want, got, cmp.Options{cmpopts.EquateApproxTime(time.Second)}); diff != "" {
				t.Fatalf("mismatch (-want, +got):\n%s", diff)
			}
		})
	}
}

func TestIssueToken(t *testing.T) {
	t.Parallel()
	db := NewTestDatabase(t)

	codeAge := time.Hour
	symptomDate := time.Now().UTC().Truncate(24 * time.Hour)
	wrongSymptomDate := symptomDate.Add(-48 * time.Hour)

	acceptConfirmed := api.AcceptTypes{
		api.TestTypeConfirmed: struct{}{},
	}

	cases := []struct {
		Name         string
		Verification func() *VerificationCode
		Accept       api.AcceptTypes
		UseLongCode  bool
		Error        string
		Delay        time.Duration
		TokenAge     time.Duration
		Subject      *Subject
		ClaimError   string
	}{
		{
			Name: "normal_token_issue",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "12345678",
					LongCode:      "12345678ABC",
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:   acceptConfirmed,
			Error:    "",
			TokenAge: time.Hour,
		},
		{
			Name: "long_code_token_issue",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "22332244",
					LongCode:      "abcd1234efgh5678",
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(5 * time.Second),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:      acceptConfirmed,
			UseLongCode: true,
			Error:       "",
			TokenAge:    time.Hour,
		},
		{
			Name: "already_claimed",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000001",
					LongCode:      "00000001ABC",
					Claimed:       true,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:   acceptConfirmed,
			Error:    ErrVerificationCodeUsed.Error(),
			TokenAge: time.Hour,
		},
		{
			Name: "code_expired",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000002",
					LongCode:      "00000002ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(1 * time.Second),
					LongExpiresAt: time.Now().Add(1 * time.Second),
				}
			},
			Accept:   acceptConfirmed,
			Delay:    2 * time.Second,
			Error:    ErrVerificationCodeExpired.Error(),
			TokenAge: time.Hour,
		},
		{
			Name: "token_expired",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000003",
					LongCode:      "00000003ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			Delay:      time.Second,
			ClaimError: ErrTokenExpired.Error(),
			TokenAge:   time.Millisecond,
		},
		{
			Name: "wrong_test_type",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000005",
					LongCode:      "00000005ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			ClaimError: ErrTokenMetadataMismatch.Error(),
			TokenAge:   time.Hour,
			Subject:    &Subject{"negative", nil},
		},
		{
			Name: "wrong_test_date",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000007",
					LongCode:      "00000007ABC",
					Claimed:       false,
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			ClaimError: ErrTokenMetadataMismatch.Error(),
			TokenAge:   time.Hour,
			Subject:    &Subject{"confirmed", &wrongSymptomDate},
		},
		{
			Name: "unsupported_test_type",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000008",
					LongCode:      "00000008ABC",
					Claimed:       false,
					TestType:      "likely",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept: acceptConfirmed,
			Error:  ErrUnsupportedTestType.Error(),
		},
	}

	for _, tc := range cases {
		tc := tc

		t.Run(tc.Name, func(t *testing.T) {
			t.Parallel()

			realm := NewRealmWithDefaults(fmt.Sprintf("TestIssueToken/%s", tc.Name))
			if err := db.SaveRealm(realm, System); err != nil {
				t.Fatal(err)
			}

			// Create the verification. We do this here instead of inside the test
			// struct to mitigate as much time drift as possible. It also ensures we
			// get a new VerificationCode on each invocation.
			verification := tc.Verification()
			verification.RealmID = realm.ID

			// Extract the code before saving. After saving, the code on the struct
			// will be the HMAC.
			code := verification.Code
			if tc.UseLongCode {
				code = verification.LongCode
			}

			if err := db.SaveVerificationCode(verification, codeAge); err != nil {
				t.Fatalf("error creating verification code: %v", err)
			}

			if tc.Delay > 0 {
				time.Sleep(tc.Delay)
			}

			tok, err := db.VerifyCodeAndIssueToken(realm.ID, code, tc.Accept, tc.TokenAge)
			if err != nil {
				if tc.Error == "" {
					t.Fatalf("error issuing token: %v", err)
				} else if !strings.Contains(err.Error(), tc.Error) {
					t.Fatalf("wrong error, want: '%v', got '%v'", tc.Error, err.Error())
				}
			} else if tc.Error != "" {
				t.Fatalf("missing error, want: '%v', got: nil", tc.Error)
			}

			if tc.Error == "" {
				if tok.TestType != verification.TestType {
					t.Errorf("test type missmatch want: %v, got %v", verification.TestType, tok.TestType)
				}
				if tok.FormatSymptomDate() != verification.FormatSymptomDate() {
					t.Errorf("test date missmatch want: %v, got %v", verification.FormatSymptomDate(), tok.FormatSymptomDate())
				}

				got, err := db.FindTokenByID(tok.TokenID)
				if err != nil {
					t.Fatalf("error reading token from db: %v", err)
				}

				if diff := cmp.Diff(tok, got, approxTime); diff != "" {
					t.Fatalf("mismatch (-want, +got):\n%s", diff)
				}

				if tc.Delay > 0 {
					time.Sleep(tc.Delay)
				}

				subject := &Subject{TestType: verification.TestType, SymptomDate: verification.SymptomDate}
				if tc.Subject != nil {
					subject = tc.Subject
				}
				if err != nil {
					t.Fatalf("unable to parse subject: %v", err)
				}
				if err := db.ClaimToken(realm.ID, got.TokenID, subject); err != nil && tc.ClaimError == "" {
					t.Fatalf("unexpected error claiming token: %v", err)
				} else if tc.ClaimError != "" {
					if err == nil {
						t.Fatalf("wanted error '%v' got: nil", tc.ClaimError)
					} else if !strings.Contains(err.Error(), tc.ClaimError) {
						t.Fatalf("wrong error, want: '%v', got '%v'", tc.ClaimError, err)
					}
				}

				if tc.ClaimError == "" {
					got, err = db.FindTokenByID(tok.TokenID)
					if err != nil {
						t.Fatalf("error reading token from db: %v", err)
					}
					if !got.Used {
						t.Fatalf("claimed token is not marked as used")
					}
				}
			}
		})
	}
}

func TestPurgeTokens(t *testing.T) {
	t.Parallel()
	db := NewTestDatabase(t)

	now := time.Now()
	testData := []*Token{
		{TokenID: "111111", TestType: "negative", ExpiresAt: now.Add(time.Second)},
		{TokenID: "222222", TestType: "negative", ExpiresAt: now.Add(time.Second)},
		{TokenID: "333333", TestType: "negative", ExpiresAt: now.Add(time.Minute)},
	}
	for _, rec := range testData {
		if err := db.db.Save(rec).Error; err != nil {
			t.Fatalf("can't save test data: %v", err)
		}
	}

	// Need to let some time lapse since we can't back date records through normal channels.
	time.Sleep(2 * time.Second)

	if count, err := db.PurgeTokens(time.Millisecond * 500); err != nil {
		t.Fatalf("error doing purge: %v", err)
	} else if count != 2 {
		t.Fatalf("purge record count mismatch, want: 2, got: %v", count)
	}
}
<commit_msg>Fix race in TestIssueToken (#835)<commit_after>// Copyright 2020 Google LLC
//
// 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 database

import (
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/google/exposure-notifications-verification-server/pkg/api"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

func TestSubject(t *testing.T) {
	testDay, err := time.Parse("2006-01-02", "2020-07-07")
	if err != nil {
		t.Fatalf("test setup error: %v", err)
	}

	cases := []struct {
		Name  string
		Sub   string
		Want  *Subject
		Error string
	}{
		{
			Name: "normal parse",
			Sub:  "confirmed.2020-07-07",
			Want: &Subject{
				TestType:    "confirmed",
				SymptomDate: &testDay,
			},
		},
		{
			Name: "no date",
			Sub:  "confirmed.",
			Want: &Subject{
				TestType:    "confirmed",
				SymptomDate: nil,
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			got, err := ParseSubject(tc.Sub)
			if err != nil {
				if tc.Error == "" {
					t.Fatalf("unexpected error: %v", err)
				} else if !strings.Contains(err.Error(), tc.Error) {
					t.Fatalf("wrong error, want: %v got: %v", tc.Error, err.Error())
				}
			} else if tc.Error != "" {
				t.Fatalf("wanted error: %v got: nil", tc.Error)
			}

			if diff := cmp.Diff(tc.Want, got, cmp.Options{cmpopts.EquateApproxTime(time.Second)}); diff != "" {
				t.Fatalf("mismatch (-want, +got):\n%s", diff)
			}
		})
	}
}

func TestIssueToken(t *testing.T) {
	t.Parallel()

	codeAge := time.Hour
	symptomDate := time.Now().UTC().Truncate(24 * time.Hour)
	wrongSymptomDate := symptomDate.Add(-48 * time.Hour)

	acceptConfirmed := api.AcceptTypes{
		api.TestTypeConfirmed: struct{}{},
	}

	cases := []struct {
		Name         string
		Verification func() *VerificationCode
		Accept       api.AcceptTypes
		UseLongCode  bool
		Error        string
		Delay        time.Duration
		TokenAge     time.Duration
		Subject      *Subject
		ClaimError   string
	}{
		{
			Name: "normal_token_issue",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "12345678",
					LongCode:      "12345678ABC",
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:   acceptConfirmed,
			Error:    "",
			TokenAge: time.Hour,
		},
		{
			Name: "long_code_token_issue",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "22332244",
					LongCode:      "abcd1234efgh5678",
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(5 * time.Second),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:      acceptConfirmed,
			UseLongCode: true,
			Error:       "",
			TokenAge:    time.Hour,
		},
		{
			Name: "already_claimed",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000001",
					LongCode:      "00000001ABC",
					Claimed:       true,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:   acceptConfirmed,
			Error:    ErrVerificationCodeUsed.Error(),
			TokenAge: time.Hour,
		},
		{
			Name: "code_expired",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000002",
					LongCode:      "00000002ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(1 * time.Second),
					LongExpiresAt: time.Now().Add(1 * time.Second),
				}
			},
			Accept:   acceptConfirmed,
			Delay:    2 * time.Second,
			Error:    ErrVerificationCodeExpired.Error(),
			TokenAge: time.Hour,
		},
		{
			Name: "token_expired",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000003",
					LongCode:      "00000003ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			Delay:      time.Second,
			ClaimError: ErrTokenExpired.Error(),
			TokenAge:   time.Millisecond,
		},
		{
			Name: "wrong_test_type",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000005",
					LongCode:      "00000005ABC",
					Claimed:       false,
					TestType:      "confirmed",
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			ClaimError: ErrTokenMetadataMismatch.Error(),
			TokenAge:   time.Hour,
			Subject:    &Subject{"negative", nil},
		},
		{
			Name: "wrong_test_date",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000007",
					LongCode:      "00000007ABC",
					Claimed:       false,
					TestType:      "confirmed",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept:     acceptConfirmed,
			ClaimError: ErrTokenMetadataMismatch.Error(),
			TokenAge:   time.Hour,
			Subject:    &Subject{"confirmed", &wrongSymptomDate},
		},
		{
			Name: "unsupported_test_type",
			Verification: func() *VerificationCode {
				return &VerificationCode{
					Code:          "00000008",
					LongCode:      "00000008ABC",
					Claimed:       false,
					TestType:      "likely",
					SymptomDate:   &symptomDate,
					ExpiresAt:     time.Now().Add(time.Hour),
					LongExpiresAt: time.Now().Add(time.Hour),
				}
			},
			Accept: acceptConfirmed,
			Error:  ErrUnsupportedTestType.Error(),
		},
	}

	for _, tc := range cases {
		tc := tc

		t.Run(tc.Name, func(t *testing.T) {
			t.Parallel()
			db := NewTestDatabase(t)

			realm := NewRealmWithDefaults(fmt.Sprintf("TestIssueToken/%s", tc.Name))
			if err := db.SaveRealm(realm, System); err != nil {
				t.Fatal(err)
			}

			// Create the verification. We do this here instead of inside the test
			// struct to mitigate as much time drift as possible. It also ensures we
			// get a new VerificationCode on each invocation.
			verification := tc.Verification()
			verification.RealmID = realm.ID

			// Extract the code before saving. After saving, the code on the struct
			// will be the HMAC.
			code := verification.Code
			if tc.UseLongCode {
				code = verification.LongCode
			}

			if err := db.SaveVerificationCode(verification, codeAge); err != nil {
				t.Fatalf("error creating verification code: %v", err)
			}

			if tc.Delay > 0 {
				time.Sleep(tc.Delay)
			}

			tok, err := db.VerifyCodeAndIssueToken(realm.ID, code, tc.Accept, tc.TokenAge)
			if err != nil {
				if tc.Error == "" {
					t.Fatalf("error issuing token: %v", err)
				} else if !strings.Contains(err.Error(), tc.Error) {
					t.Fatalf("wrong error, want: '%v', got '%v'", tc.Error, err.Error())
				}
			} else if tc.Error != "" {
				t.Fatalf("missing error, want: '%v', got: nil", tc.Error)
			}

			if tc.Error == "" {
				if tok.TestType != verification.TestType {
					t.Errorf("test type missmatch want: %v, got %v", verification.TestType, tok.TestType)
				}
				if tok.FormatSymptomDate() != verification.FormatSymptomDate() {
					t.Errorf("test date missmatch want: %v, got %v", verification.FormatSymptomDate(), tok.FormatSymptomDate())
				}

				got, err := db.FindTokenByID(tok.TokenID)
				if err != nil {
					t.Fatalf("error reading token from db: %v", err)
				}

				if diff := cmp.Diff(tok, got, approxTime); diff != "" {
					t.Fatalf("mismatch (-want, +got):\n%s", diff)
				}

				if tc.Delay > 0 {
					time.Sleep(tc.Delay)
				}

				subject := &Subject{TestType: verification.TestType, SymptomDate: verification.SymptomDate}
				if tc.Subject != nil {
					subject = tc.Subject
				}
				if err != nil {
					t.Fatalf("unable to parse subject: %v", err)
				}
				if err := db.ClaimToken(realm.ID, got.TokenID, subject); err != nil && tc.ClaimError == "" {
					t.Fatalf("unexpected error claiming token: %v", err)
				} else if tc.ClaimError != "" {
					if err == nil {
						t.Fatalf("wanted error '%v' got: nil", tc.ClaimError)
					} else if !strings.Contains(err.Error(), tc.ClaimError) {
						t.Fatalf("wrong error, want: '%v', got '%v'", tc.ClaimError, err)
					}
				}

				if tc.ClaimError == "" {
					got, err = db.FindTokenByID(tok.TokenID)
					if err != nil {
						t.Fatalf("error reading token from db: %v", err)
					}
					if !got.Used {
						t.Fatalf("claimed token is not marked as used")
					}
				}
			}
		})
	}
}

func TestPurgeTokens(t *testing.T) {
	t.Parallel()
	db := NewTestDatabase(t)

	now := time.Now()
	testData := []*Token{
		{TokenID: "111111", TestType: "negative", ExpiresAt: now.Add(time.Second)},
		{TokenID: "222222", TestType: "negative", ExpiresAt: now.Add(time.Second)},
		{TokenID: "333333", TestType: "negative", ExpiresAt: now.Add(time.Minute)},
	}
	for _, rec := range testData {
		if err := db.db.Save(rec).Error; err != nil {
			t.Fatalf("can't save test data: %v", err)
		}
	}

	// Need to let some time lapse since we can't back date records through normal channels.
	time.Sleep(2 * time.Second)

	if count, err := db.PurgeTokens(time.Millisecond * 500); err != nil {
		t.Fatalf("error doing purge: %v", err)
	} else if count != 2 {
		t.Fatalf("purge record count mismatch, want: 2, got: %v", count)
	}
}
<|endoftext|>


<commit_before><commit_msg>support mounting with the --no-kubernetes flag<commit_after><|endoftext|>


<commit_before>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package history

import (
	"github.com/nelsam/gxui"
	"github.com/nelsam/gxui/themes/basic"
	"github.com/nelsam/vidar/commander/bind"
	"github.com/nelsam/vidar/commander/input"
	"github.com/nelsam/vidar/plugin/command"
)

// Bindables returns the slice of bind.Bindable types that is implemented
// by this package.
func Bindables(_ command.Commander, _ gxui.Driver, theme *basic.Theme) []bind.Bindable {
	h := History{all: make(map[string]*branch)}
	h.resetCurrent("")
	onOpen := OnOpen{theme: theme}
	return []bind.Bindable{&h, &onOpen}
}

// History keeps track of change history for a file.
type History struct {
	current tree
	skip    linkedlist

	all map[string]*branch
}

// Name returns the name of h
func (h *History) Name() string {
	return "history"
}

// OpNames returns the name of bind.Op types that
// h needs to bind to.
func (h *History) OpNames() []string {
	return []string{"input-handler", "focus-location"}
}

// Init implements input.ChangeHook
func (h *History) Init(input.Editor, []rune) {}

// resetCurrent resets h.current.trunk to a previous history (if one
// exists for path) or a new empty branch.
func (h *History) resetCurrent(path string) {
	if n, ok := h.all[path]; ok {
		h.current.setTrunk(n)
		return
	}
	h.current.setTrunk(&branch{})
}

// addSkip takes an edit that has been returned by h (from either
// Rewind or FastForward) and adds it to h.skipP, so that h will
// ignore e when it is triggered by TextChanged.
func (h *History) addSkip(e input.Edit) {
	n := &node{edit: e}
	added := h.skip.casHead(nil, n)
	if !added {
		for curr := h.skip.head(); !added; curr = curr.next() {
			added = curr.casNext(nil, n)
		}
	}
}

// shouldSkip reports whether e is an edit that was created by h and
// should be skipped.
func (h *History) shouldSkip(e input.Edit) bool {
	skip := h.skip.head()
	if skip == nil {
		return false
	}
	return skip.edit.At == e.At &&
		string(skip.edit.Old) == string(e.Old) &&
		string(skip.edit.New) == string(e.New)
}

// TextChanged hooks into the input handler to trigger off of changes
// in the editor so that h can track the history of those changes.
func (h *History) TextChanged(_ input.Editor, e input.Edit) {
	if h.shouldSkip(e) {
		h.skip.setHead(h.skip.head().next())
		return
	}
	h.current.setTrunk(h.current.trunk().push(e))
}

// Rewind tells h to rewind its current state and return the
// input.Edit that needs to be applied in order to rewind the
// text to its previous state.
func (h *History) Rewind() input.Edit {
	curr := h.current.trunk()
	prev := curr.prev()
	if prev == nil {
		return input.Edit{At: -1}
	}
	e := curr.edit
	h.current.setTrunk(prev)
	undo := input.Edit{
		At:  e.At,
		Old: e.New,
		New: e.Old,
	}
	h.addSkip(undo)
	return undo
}

// Branches returns the number of branches available to fast
// forward to at the current branch.
func (h *History) Branches() uint {
	next := h.current.trunk().next(0)
	if next == nil {
		return 0
	}
	return uint(len(next.siblings()) + 1)
}

// FastForward moves h's state forward in the history, based
// on branch.  To fast forward the most recent undo, run
// h.FastForward(h.Branches() - 1).
func (h *History) FastForward(branch uint) input.Edit {
	ff := h.current.trunk().next(branch)
	if ff == nil {
		return input.Edit{At: -1}
	}
	h.current.setTrunk(ff)
	h.addSkip(ff.edit)
	return ff.edit
}

// Apply is unused on history - recording the changes is all
// that is needed.
func (h *History) Apply(input.Editor) error { return nil }

// FileChanged updates the current history when the focused
// file is changed.
func (h *History) FileChanged(oldPath, newPath string) {
	if oldPath != "" {
		h.all[oldPath] = h.current.trunk()
	}
	h.resetCurrent(newPath)
	h.skip.setHead(nil)
}
<commit_msg>command/history: synchronize access to the 'all' map, too<commit_after>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package history

import (
	"sync"

	"github.com/nelsam/gxui"
	"github.com/nelsam/gxui/themes/basic"
	"github.com/nelsam/vidar/commander/bind"
	"github.com/nelsam/vidar/commander/input"
	"github.com/nelsam/vidar/plugin/command"
)

// Bindables returns the slice of bind.Bindable types that is implemented
// by this package.
func Bindables(_ command.Commander, _ gxui.Driver, theme *basic.Theme) []bind.Bindable {
	h := History{all: make(map[string]*branch)}
	h.resetCurrent("")
	onOpen := OnOpen{theme: theme}
	return []bind.Bindable{&h, &onOpen}
}

// History keeps track of change history for a file.
type History struct {
	current tree
	skip    linkedlist

	// all stores history for all files - it is used when the open
	// file is changed, to store the history for the previously open
	// file.  We're just using a mutex to synchronize access to it
	// because it will only be accessed when the open file is changed.
	all   map[string]*branch
	allMu sync.Mutex
}

// Name returns the name of h
func (h *History) Name() string {
	return "history"
}

// OpNames returns the name of bind.Op types that
// h needs to bind to.
func (h *History) OpNames() []string {
	return []string{"input-handler", "focus-location"}
}

// Init implements input.ChangeHook
func (h *History) Init(input.Editor, []rune) {}

// resetCurrent resets h.current.trunk to a previous history (if one
// exists for path) or a new empty branch.
func (h *History) resetCurrent(path string) {
	if n, ok := h.all[path]; ok {
		h.current.setTrunk(n)
		return
	}
	h.current.setTrunk(&branch{})
}

// addSkip takes an edit that has been returned by h (from either
// Rewind or FastForward) and adds it to h.skipP, so that h will
// ignore e when it is triggered by TextChanged.
func (h *History) addSkip(e input.Edit) {
	n := &node{edit: e}
	added := h.skip.casHead(nil, n)
	if !added {
		for curr := h.skip.head(); !added; curr = curr.next() {
			added = curr.casNext(nil, n)
		}
	}
}

// shouldSkip reports whether e is an edit that was created by h and
// should be skipped.
func (h *History) shouldSkip(e input.Edit) bool {
	skip := h.skip.head()
	if skip == nil {
		return false
	}
	return skip.edit.At == e.At &&
		string(skip.edit.Old) == string(e.Old) &&
		string(skip.edit.New) == string(e.New)
}

// TextChanged hooks into the input handler to trigger off of changes
// in the editor so that h can track the history of those changes.
func (h *History) TextChanged(_ input.Editor, e input.Edit) {
	if h.shouldSkip(e) {
		h.skip.setHead(h.skip.head().next())
		return
	}
	h.current.setTrunk(h.current.trunk().push(e))
}

// Rewind tells h to rewind its current state and return the
// input.Edit that needs to be applied in order to rewind the
// text to its previous state.
func (h *History) Rewind() input.Edit {
	curr := h.current.trunk()
	prev := curr.prev()
	if prev == nil {
		return input.Edit{At: -1}
	}
	e := curr.edit
	h.current.setTrunk(prev)
	undo := input.Edit{
		At:  e.At,
		Old: e.New,
		New: e.Old,
	}
	h.addSkip(undo)
	return undo
}

// Branches returns the number of branches available to fast
// forward to at the current branch.
func (h *History) Branches() uint {
	next := h.current.trunk().next(0)
	if next == nil {
		return 0
	}
	return uint(len(next.siblings()) + 1)
}

// FastForward moves h's state forward in the history, based
// on branch.  To fast forward the most recent undo, run
// h.FastForward(h.Branches() - 1).
func (h *History) FastForward(branch uint) input.Edit {
	ff := h.current.trunk().next(branch)
	if ff == nil {
		return input.Edit{At: -1}
	}
	h.current.setTrunk(ff)
	h.addSkip(ff.edit)
	return ff.edit
}

// Apply is unused on history - recording the changes is all
// that is needed.
func (h *History) Apply(input.Editor) error { return nil }

// FileChanged updates the current history when the focused
// file is changed.
func (h *History) FileChanged(oldPath, newPath string) {
	h.allMu.Lock()
	defer h.allMu.Unlock()
	if oldPath != "" {
		h.all[oldPath] = h.current.trunk()
	}
	h.resetCurrent(newPath)
	h.skip.setHead(nil)
}
<|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 wait

import (
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/util"
)

func TestPoller(t *testing.T) {
	done := make(chan struct{})
	defer close(done)
	w := poller(time.Millisecond, 2*time.Millisecond)
	ch := w(done)
	count := 0
DRAIN:
	for {
		select {
		case _, open := <-ch:
			if !open {
				break DRAIN
			}
			count++
		case <-time.After(util.ForeverTestTimeout):
			t.Errorf("unexpected timeout after poll")
		}
	}
	if count > 3 {
		t.Errorf("expected up to three values, got %d", count)
	}
}

func fakeTicker(max int, used *int32) WaitFunc {
	return func(done <-chan struct{}) <-chan struct{} {
		ch := make(chan struct{})
		go func() {
			defer close(ch)
			for i := 0; i < max; i++ {
				select {
				case ch <- struct{}{}:
				case <-done:
					return
				}
				if used != nil {
					atomic.AddInt32(used, 1)
				}
			}
		}()
		return ch
	}
}

type fakePoller struct {
	max  int
	used int32 // accessed with atomics
}

func (fp *fakePoller) GetWaitFunc(interval, timeout time.Duration) WaitFunc {
	return fakeTicker(fp.max, &fp.used)
}

func TestPoll(t *testing.T) {
	invocations := 0
	f := ConditionFunc(func() (bool, error) {
		invocations++
		return true, nil
	})
	fp := fakePoller{max: 1}
	if err := pollInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err != nil {
		t.Fatalf("unexpected error %v", err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used := atomic.LoadInt32(&fp.used)
	if used != 1 {
		t.Errorf("Expected exactly one tick, got %d", used)
	}

	expectedError := errors.New("Expected error")
	f = ConditionFunc(func() (bool, error) {
		return false, expectedError
	})
	fp = fakePoller{max: 1}
	if err := pollInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err == nil || err != expectedError {
		t.Fatalf("Expected error %v, got none %v", expectedError, err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used = atomic.LoadInt32(&fp.used)
	if used != 1 {
		t.Errorf("Expected exactly one tick, got %d", used)
	}
}

func TestPollImmediate(t *testing.T) {
	invocations := 0
	f := ConditionFunc(func() (bool, error) {
		invocations++
		return true, nil
	})
	fp := fakePoller{max: 0}
	if err := pollImmediateInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err != nil {
		t.Fatalf("unexpected error %v", err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used := atomic.LoadInt32(&fp.used)
	if used != 0 {
		t.Errorf("Expected exactly zero ticks, got %d", used)
	}

	expectedError := errors.New("Expected error")
	f = ConditionFunc(func() (bool, error) {
		return false, expectedError
	})
	fp = fakePoller{max: 0}
	if err := pollImmediateInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err == nil || err != expectedError {
		t.Fatalf("Expected error %v, got none %v", expectedError, err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used = atomic.LoadInt32(&fp.used)
	if used != 0 {
		t.Errorf("Expected exactly zero ticks, got %d", used)
	}
}

func TestPollForever(t *testing.T) {
	ch := make(chan struct{})
	done := make(chan struct{}, 1)
	complete := make(chan struct{})
	go func() {
		f := ConditionFunc(func() (bool, error) {
			ch <- struct{}{}
			select {
			case <-done:
				return true, nil
			default:
			}
			return false, nil
		})

		if err := PollInfinite(time.Microsecond, f); err != nil {
			t.Fatalf("unexpected error %v", err)
		}

		close(ch)
		complete <- struct{}{}
	}()

	// ensure the condition is opened
	<-ch

	// ensure channel sends events
	for i := 0; i < 10; i++ {
		select {
		case _, open := <-ch:
			if !open {
				t.Fatalf("did not expect channel to be closed")
			}
		case <-time.After(util.ForeverTestTimeout):
			t.Fatalf("channel did not return at least once within the poll interval")
		}
	}

	// at most two poll notifications should be sent once we return from the condition
	done <- struct{}{}
	go func() {
		for i := 0; i < 2; i++ {
			_, open := <-ch
			if open {
				<-complete
				return
			}
		}
		t.Fatalf("expected closed channel after two iterations")
	}()
	<-complete
}

func TestWaitFor(t *testing.T) {
	var invocations int
	testCases := map[string]struct {
		F       ConditionFunc
		Ticks   int
		Invoked int
		Err     bool
	}{
		"invoked once": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return true, nil
			}),
			2,
			1,
			false,
		},
		"invoked and returns a timeout": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return false, nil
			}),
			2,
			3, // the contract of WaitFor() says the func is called once more at the end of the wait
			true,
		},
		"returns immediately on error": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return false, errors.New("test")
			}),
			2,
			1,
			true,
		},
	}
	for k, c := range testCases {
		invocations = 0
		ticker := fakeTicker(c.Ticks, nil)
		err := func() error {
			done := make(chan struct{})
			defer close(done)
			return WaitFor(ticker, c.F, done)
		}()
		switch {
		case c.Err && err == nil:
			t.Errorf("%s: Expected error, got nil", k)
			continue
		case !c.Err && err != nil:
			t.Errorf("%s: Expected no error, got: %#v", k, err)
			continue
		}
		if invocations != c.Invoked {
			t.Errorf("%s: Expected %d invocations, got %d", k, c.Invoked, invocations)
		}
	}
}

func TestWaitForWithDelay(t *testing.T) {
	done := make(chan struct{})
	defer close(done)
	WaitFor(poller(time.Millisecond, util.ForeverTestTimeout), func() (bool, error) {
		time.Sleep(10 * time.Millisecond)
		return true, nil
	}, done)
	// If polling goroutine doesn't see the done signal it will leak timers.
	select {
	case done <- struct{}{}:
	case <-time.After(util.ForeverTestTimeout):
		t.Errorf("expected an ack of the done signal.")
	}
}
<commit_msg>Fix wait_test flakes<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 wait

import (
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/util"
)

func TestPoller(t *testing.T) {
	done := make(chan struct{})
	defer close(done)
	w := poller(time.Millisecond, 2*time.Millisecond)
	ch := w(done)
	count := 0
DRAIN:
	for {
		select {
		case _, open := <-ch:
			if !open {
				break DRAIN
			}
			count++
		case <-time.After(util.ForeverTestTimeout):
			t.Errorf("unexpected timeout after poll")
		}
	}
	if count > 3 {
		t.Errorf("expected up to three values, got %d", count)
	}
}

func fakeTicker(max int, used *int32) WaitFunc {
	return func(done <-chan struct{}) <-chan struct{} {
		ch := make(chan struct{})
		go func() {
			defer close(ch)
			for i := 0; i < max; i++ {
				select {
				case ch <- struct{}{}:
				case <-done:
					return
				}
				if used != nil {
					atomic.AddInt32(used, 1)
				}
			}
		}()
		return ch
	}
}

type fakePoller struct {
	max  int
	used int32 // accessed with atomics
}

func (fp *fakePoller) GetWaitFunc(interval, timeout time.Duration) WaitFunc {
	return fakeTicker(fp.max, &fp.used)
}

func TestPoll(t *testing.T) {
	invocations := 0
	f := ConditionFunc(func() (bool, error) {
		invocations++
		return true, nil
	})
	fp := fakePoller{max: 1}
	if err := pollInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err != nil {
		t.Fatalf("unexpected error %v", err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used := atomic.LoadInt32(&fp.used)
	if used != 1 {
		t.Errorf("Expected exactly one tick, got %d", used)
	}

	expectedError := errors.New("Expected error")
	f = ConditionFunc(func() (bool, error) {
		return false, expectedError
	})
	fp = fakePoller{max: 1}
	if err := pollInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err == nil || err != expectedError {
		t.Fatalf("Expected error %v, got none %v", expectedError, err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used = atomic.LoadInt32(&fp.used)
	if used != 1 {
		t.Errorf("Expected exactly one tick, got %d", used)
	}
}

func TestPollImmediate(t *testing.T) {
	invocations := 0
	f := ConditionFunc(func() (bool, error) {
		invocations++
		return true, nil
	})
	fp := fakePoller{max: 0}
	if err := pollImmediateInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err != nil {
		t.Fatalf("unexpected error %v", err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used := atomic.LoadInt32(&fp.used)
	if used != 0 {
		t.Errorf("Expected exactly zero ticks, got %d", used)
	}

	expectedError := errors.New("Expected error")
	f = ConditionFunc(func() (bool, error) {
		return false, expectedError
	})
	fp = fakePoller{max: 0}
	if err := pollImmediateInternal(fp.GetWaitFunc(time.Microsecond, time.Microsecond), f); err == nil || err != expectedError {
		t.Fatalf("Expected error %v, got none %v", expectedError, err)
	}
	if invocations != 1 {
		t.Errorf("Expected exactly one invocation, got %d", invocations)
	}
	used = atomic.LoadInt32(&fp.used)
	if used != 0 {
		t.Errorf("Expected exactly zero ticks, got %d", used)
	}
}

func TestPollForever(t *testing.T) {
	ch := make(chan struct{})
	done := make(chan struct{}, 1)
	complete := make(chan struct{})
	go func() {
		f := ConditionFunc(func() (bool, error) {
			ch <- struct{}{}
			select {
			case <-done:
				return true, nil
			default:
			}
			return false, nil
		})

		if err := PollInfinite(time.Microsecond, f); err != nil {
			t.Fatalf("unexpected error %v", err)
		}

		close(ch)
		complete <- struct{}{}
	}()

	// ensure the condition is opened
	<-ch

	// ensure channel sends events
	for i := 0; i < 10; i++ {
		select {
		case _, open := <-ch:
			if !open {
				t.Fatalf("did not expect channel to be closed")
			}
		case <-time.After(util.ForeverTestTimeout):
			t.Fatalf("channel did not return at least once within the poll interval")
		}
	}

	// at most one poll notification should be sent once we return from the condition
	done <- struct{}{}
	go func() {
		for i := 0; i < 2; i++ {
			_, open := <-ch
			if !open {
				return
			}
		}
		t.Fatalf("expected closed channel after two iterations")
	}()
	<-complete
}

func TestWaitFor(t *testing.T) {
	var invocations int
	testCases := map[string]struct {
		F       ConditionFunc
		Ticks   int
		Invoked int
		Err     bool
	}{
		"invoked once": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return true, nil
			}),
			2,
			1,
			false,
		},
		"invoked and returns a timeout": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return false, nil
			}),
			2,
			3, // the contract of WaitFor() says the func is called once more at the end of the wait
			true,
		},
		"returns immediately on error": {
			ConditionFunc(func() (bool, error) {
				invocations++
				return false, errors.New("test")
			}),
			2,
			1,
			true,
		},
	}
	for k, c := range testCases {
		invocations = 0
		ticker := fakeTicker(c.Ticks, nil)
		err := func() error {
			done := make(chan struct{})
			defer close(done)
			return WaitFor(ticker, c.F, done)
		}()
		switch {
		case c.Err && err == nil:
			t.Errorf("%s: Expected error, got nil", k)
			continue
		case !c.Err && err != nil:
			t.Errorf("%s: Expected no error, got: %#v", k, err)
			continue
		}
		if invocations != c.Invoked {
			t.Errorf("%s: Expected %d invocations, got %d", k, c.Invoked, invocations)
		}
	}
}

func TestWaitForWithDelay(t *testing.T) {
	done := make(chan struct{})
	defer close(done)
	WaitFor(poller(time.Millisecond, util.ForeverTestTimeout), func() (bool, error) {
		time.Sleep(10 * time.Millisecond)
		return true, nil
	}, done)
	// If polling goroutine doesn't see the done signal it will leak timers.
	select {
	case done <- struct{}{}:
	case <-time.After(util.ForeverTestTimeout):
		t.Errorf("expected an ack of the done signal.")
	}
}
<|endoftext|>


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

// Package logger logs messages as appropriate.
package logger

import (
	"context"
	"fmt"
	"io"
	"os"
	"time"

	"cloud.google.com/go/logging"
)

var (
	// DeferredFatalFuncs is a slice of functions that will be called prior to os.Exit in Fatal.
	DeferredFatalFuncs []func()

	cloudLoggingClient *logging.Client
	cloudLogger        *logging.Logger
	debugEnabled       bool
	loggerName         string
	formatFunction     func(LogEntry) string

	writers []io.Writer
)

// LogOpts represents options for logging.
type LogOpts struct {
	Debug               bool
	ProjectName         string
	LoggerName          string
	DisableLocalLogging bool
	DisableCloudLogging bool
	// FormatFunction will produce the string representation of each log event.
	FormatFunction func(LogEntry) string
	// Additional writers that will be used during logging.
	Writers []io.Writer
}

// Init instantiates the logger.
func Init(ctx context.Context, opts LogOpts) error {
	if opts.LoggerName == "" {
		return fmt.Errorf("logger name must be set")
	}

	loggerName = opts.LoggerName
	debugEnabled = opts.Debug
	formatFunction = opts.FormatFunction
	writers = opts.Writers

	if !opts.DisableLocalLogging {
		if err := localSetup(loggerName); err != nil {
			return fmt.Errorf("logger Init localSetup error: %v", err)
		}
	}

	if !opts.DisableCloudLogging && opts.ProjectName != "" {
		var err error
		cloudLoggingClient, err = logging.NewClient(ctx, opts.ProjectName)
		if err != nil {
			Errorf("Continuing without cloud logging due to error in initialization: %v", err.Error())
			// Log but don't return this error, as it doesn't prevent continuing.
			return nil
		}

		// This automatically detects and associates with a GCE resource.
		cloudLogger = cloudLoggingClient.Logger(loggerName)

		go func() {
			for {
				time.Sleep(5 * time.Second)
				cloudLogger.Flush()
			}
		}()
	}

	return nil
}

// Close closes the logger.
func Close() {
	if cloudLoggingClient != nil {
		cloudLoggingClient.Close()
	}
	localClose()
}

// Log writes an entry to all outputs.
func Log(e LogEntry) {
	if e.Severity == Debug && !debugEnabled {
		return
	}
	if e.CallDepth == 0 {
		e.CallDepth = 2
	}
	e.LocalTimestamp = now()
	e.Source = caller(e.CallDepth)
	local(e)
	for _, w := range writers {
		w.Write(e.bytes())
	}

	var cloudSev logging.Severity
	if cloudLogger != nil {
		switch e.Severity {
		case Debug:
			cloudSev = logging.Debug
		case Info:
			cloudSev = logging.Info
		case Warning:
			cloudSev = logging.Warning
		case Error:
			cloudSev = logging.Error
		case Critical:
			cloudSev = logging.Critical
		default:
			cloudSev = logging.Default
		}
		cloudLogger.Log(logging.Entry{Severity: cloudSev, SourceLocation: e.Source, Payload: e, Labels: e.Labels})
	}
}

// Debugf logs debug information.
func Debugf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Debug})
}

// Infof logs general information.
func Infof(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Info})
}

// Warningf logs warning information.
func Warningf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Warning})
}

// Errorf logs error information.
func Errorf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Error})
}

// Fatalf logs critical error information and exits.
func Fatalf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Critical})

	for _, f := range DeferredFatalFuncs {
		f()
	}
	Close()
	os.Exit(1)
}
<commit_msg>Add SetDebugLogging to enable debug logging dynamically (#12)<commit_after>//  Copyright 2019 Google Inc. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// Package logger logs messages as appropriate.
package logger

import (
	"context"
	"fmt"
	"io"
	"os"
	"time"

	"cloud.google.com/go/logging"
)

var (
	// DeferredFatalFuncs is a slice of functions that will be called prior to os.Exit in Fatal.
	DeferredFatalFuncs []func()

	cloudLoggingClient *logging.Client
	cloudLogger        *logging.Logger
	debugEnabled       bool
	loggerName         string
	formatFunction     func(LogEntry) string

	writers []io.Writer
)

// LogOpts represents options for logging.
type LogOpts struct {
	Debug               bool
	ProjectName         string
	LoggerName          string
	DisableLocalLogging bool
	DisableCloudLogging bool
	// FormatFunction will produce the string representation of each log event.
	FormatFunction func(LogEntry) string
	// Additional writers that will be used during logging.
	Writers []io.Writer
}

// SetDebugLogging enables or disables debug level logging.
func SetDebugLogging(enabled bool) {
	debugEnabled = enabled
}

// Init instantiates the logger.
func Init(ctx context.Context, opts LogOpts) error {
	if opts.LoggerName == "" {
		return fmt.Errorf("logger name must be set")
	}

	loggerName = opts.LoggerName
	debugEnabled = opts.Debug
	formatFunction = opts.FormatFunction
	writers = opts.Writers

	if !opts.DisableLocalLogging {
		if err := localSetup(loggerName); err != nil {
			return fmt.Errorf("logger Init localSetup error: %v", err)
		}
	}

	if !opts.DisableCloudLogging && opts.ProjectName != "" {
		var err error
		cloudLoggingClient, err = logging.NewClient(ctx, opts.ProjectName)
		if err != nil {
			Errorf("Continuing without cloud logging due to error in initialization: %v", err.Error())
			// Log but don't return this error, as it doesn't prevent continuing.
			return nil
		}

		// This automatically detects and associates with a GCE resource.
		cloudLogger = cloudLoggingClient.Logger(loggerName)

		go func() {
			for {
				time.Sleep(5 * time.Second)
				cloudLogger.Flush()
			}
		}()
	}

	return nil
}

// Close closes the logger.
func Close() {
	if cloudLoggingClient != nil {
		cloudLoggingClient.Close()
	}
	localClose()
}

// Log writes an entry to all outputs.
func Log(e LogEntry) {
	if e.Severity == Debug && !debugEnabled {
		return
	}
	if e.CallDepth == 0 {
		e.CallDepth = 2
	}
	e.LocalTimestamp = now()
	e.Source = caller(e.CallDepth)
	local(e)
	for _, w := range writers {
		w.Write(e.bytes())
	}

	var cloudSev logging.Severity
	if cloudLogger != nil {
		switch e.Severity {
		case Debug:
			cloudSev = logging.Debug
		case Info:
			cloudSev = logging.Info
		case Warning:
			cloudSev = logging.Warning
		case Error:
			cloudSev = logging.Error
		case Critical:
			cloudSev = logging.Critical
		default:
			cloudSev = logging.Default
		}
		cloudLogger.Log(logging.Entry{Severity: cloudSev, SourceLocation: e.Source, Payload: e, Labels: e.Labels})
	}
}

// Debugf logs debug information.
func Debugf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Debug})
}

// Infof logs general information.
func Infof(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Info})
}

// Warningf logs warning information.
func Warningf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Warning})
}

// Errorf logs error information.
func Errorf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Error})
}

// Fatalf logs critical error information and exits.
func Fatalf(format string, v ...interface{}) {
	Log(LogEntry{Message: fmt.Sprintf(format, v...), Severity: Critical})

	for _, f := range DeferredFatalFuncs {
		f()
	}
	Close()
	os.Exit(1)
}
<|endoftext|>


<commit_before>package aetestutil

import (
	"net/http"
	"appengine"
	"appengine/aetest"
	"github.com/pcl/httptest/aeutil"
)

/**
 * The Context returned from this function is registered with the current ContextProvider,
 * so calls to aeutil.UncurryWithContext will resolve to this context.
 */
func CreateAndRegisterTestContext() (context aetest.Context, err error) {
	context, err = aetest.NewContext(nil)

	// TODO change this function to maintain and consult a map, so we can
	// make multiple requests in parallel / in a recursive fashion
	aeutil.ContextProvider = func(req *http.Request) appengine.Context {
		// Copy the request details from the user-created one to the
		// context's one, and edit it to add in the Appengine headers,
		// as seen in appengine/user/user_dev.go. This has to happen
		// before we make any calls to the user module, which works
		// out since this is called during UncurryWithContext.
		contextRequest := context.Request().(*http.Request)
		contextRequest.Body = req.Body
		contextRequest.ContentLength = req.ContentLength
		contextRequest.Form = req.Form
		contextRequest.Header = req.Header
		contextRequest.Host = req.Host
		contextRequest.Method = req.Method
		contextRequest.MultipartForm = req.MultipartForm
		contextRequest.PostForm = req.PostForm
		contextRequest.Proto = req.Proto
		contextRequest.ProtoMajor = req.ProtoMajor
		contextRequest.ProtoMinor = req.ProtoMinor
		contextRequest.RemoteAddr = req.RemoteAddr
		contextRequest.RequestURI = req.RequestURI
		contextRequest.TLS = req.TLS
		contextRequest.Trailer = req.Trailer
		contextRequest.TransferEncoding = req.TransferEncoding
		contextRequest.URL = req.URL

		return context
	}

	return
}
<commit_msg>set the username to something so that test calls are authenticated. Need to make this configurable.<commit_after>package aetestutil

import (
	"net/http"
	"appengine"
	"appengine/aetest"
	"github.com/pcl/httptest/aeutil"
)

/**
 * The Context returned from this function is registered with the current ContextProvider,
 * so calls to aeutil.UncurryWithContext will resolve to this context.
 */
func CreateAndRegisterTestContext() (context aetest.Context, err error) {
	context, err = aetest.NewContext(nil)

	// TODO change this function to maintain and consult a map, so we can
	// make multiple requests in parallel / in a recursive fashion
	aeutil.ContextProvider = func(req *http.Request) appengine.Context {
		// Copy the request details from the user-created one to the
		// context's one, and edit it to add in the Appengine headers,
		// as seen in appengine/user/user_dev.go. This has to happen
		// before we make any calls to the user module, which works
		// out since this is called during UncurryWithContext.
		contextRequest := context.Request().(*http.Request)
		contextRequest.Body = req.Body
		contextRequest.ContentLength = req.ContentLength
		contextRequest.Form = req.Form
		contextRequest.Header = req.Header
		contextRequest.Host = req.Host
		contextRequest.Method = req.Method
		contextRequest.MultipartForm = req.MultipartForm
		contextRequest.PostForm = req.PostForm
		contextRequest.Proto = req.Proto
		contextRequest.ProtoMajor = req.ProtoMajor
		contextRequest.ProtoMinor = req.ProtoMinor
		contextRequest.RemoteAddr = req.RemoteAddr
		contextRequest.RequestURI = req.RequestURI
		contextRequest.TLS = req.TLS
		contextRequest.Trailer = req.Trailer
		contextRequest.TransferEncoding = req.TransferEncoding
		contextRequest.URL = req.URL

		// ##### should have some way to register what header values to use here
		req.Header.Set("X-AppEngine-Internal-User-Email", "foo@example.com")
		return context
	}

	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 dbsync

import (
	"encoding/json"

	"github.com/golang/protobuf/proto"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/datasync/syncbase"
	"github.com/ligato/cn-infra/db/keyval"
)

// NewAdapter creates a new instance of Adapter.
func NewAdapter(name string, db keyval.BytesBroker, dbW keyval.BytesWatcher) *Adapter {
	return &Adapter{name, db, dbW, syncbase.NewWatcher()}
}

// Adapter is ETCD watcher/publisher in front of Agent Plugins
type Adapter struct {
	Name string
	db   keyval.BytesBroker
	dbW  keyval.BytesWatcher
	base *syncbase.Watcher
}

// WatchData using ETCD or any other default transport.
func (adapter *Adapter) Watch(resyncName string, changeChan chan datasync.ChangeEvent,
	resyncChan chan datasync.ResyncEvent, keyPrefixes ...string) (datasync.WatchDataRegistration, error) {

	reg, err := adapter.base.WatchData(resyncName, changeChan, resyncChan, keyPrefixes...)
	if err != nil {
		return nil, err
	}

	_, err = watchAndResyncBrokerKeys(resyncName, changeChan, resyncChan, adapter, keyPrefixes...)
	if err != nil {
		return nil, err
	}

	return reg, err
}

// Put using ETCD or any other default transport
func (adapter *Adapter) Put(key string, data proto.Message, options ...datasync.PutOption) error {
	if data == nil {
		_, err := adapter.db.Delete(key)
		return err
	}

	bytes, err := json.Marshal(data)
	if err != nil {
		return err
	}

	return adapter.db.Put(key, bytes)
}

// String return resyncName
func (adapter *Adapter) String() string {
	return adapter.Name
}
<commit_msg>ODPM-361 Fix golint<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 dbsync

import (
	"encoding/json"

	"github.com/golang/protobuf/proto"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/datasync/syncbase"
	"github.com/ligato/cn-infra/db/keyval"
)

// NewAdapter creates a new instance of Adapter.
func NewAdapter(name string, db keyval.BytesBroker, dbW keyval.BytesWatcher) *Adapter {
	return &Adapter{name, db, dbW, syncbase.NewWatcher()}
}

// Adapter is ETCD watcher/publisher in front of Agent Plugins
type Adapter struct {
	Name string
	db   keyval.BytesBroker
	dbW  keyval.BytesWatcher
	base *syncbase.Watcher
}

// Watch using ETCD or any other default transport.
func (adapter *Adapter) Watch(resyncName string, changeChan chan datasync.ChangeEvent,
	resyncChan chan datasync.ResyncEvent, keyPrefixes ...string) (datasync.WatchDataRegistration, error) {

	reg, err := adapter.base.WatchData(resyncName, changeChan, resyncChan, keyPrefixes...)
	if err != nil {
		return nil, err
	}

	_, err = watchAndResyncBrokerKeys(resyncName, changeChan, resyncChan, adapter, keyPrefixes...)
	if err != nil {
		return nil, err
	}

	return reg, err
}

// Put using ETCD or any other default transport
func (adapter *Adapter) Put(key string, data proto.Message, options ...datasync.PutOption) error {
	if data == nil {
		_, err := adapter.db.Delete(key)
		return err
	}

	bytes, err := json.Marshal(data)
	if err != nil {
		return err
	}

	return adapter.db.Put(key, bytes)
}

// String return resyncName
func (adapter *Adapter) String() string {
	return adapter.Name
}
<|endoftext|>


<commit_before>package warp

import (
	"fmt"
	"errors"
	"bytes"
	"os"
	"time"
	"log"
	"log/syslog"
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"sync"
	"encoding/binary"
	"encoding/json"
)

type Client struct {
	Config Config
	Conn *tls.Conn
	Logger *log.Logger
	Connected bool
	Input chan *Packet
	Output chan *Packet
}

func NewClient(cfg Config) *Client {

	var logger *log.Logger
	var err error

	switch {
	case cfg.LogTo == "stdout":
		logger = log.New(os.Stdout, "warp: ", 0)
	case cfg.LogTo == "stderr":
		logger = log.New(os.Stderr, "warp: ", 0)
	case cfg.LogTo == "syslog":
		logger, err = syslog.NewLogger(syslog.LOG_INFO | syslog.LOG_DAEMON, 0)
		if err != nil {
			fmt.Printf("cannot initialize logger: %v", err)
			os.Exit(1)
		}
	default:
		fmt.Printf("cannot initialize logger, bad configuration")
		os.Exit(1)
	}

	cert, err := tls.LoadX509KeyPair(cfg.Cert, cfg.PrivKey)
	if err != nil {
		logger.Fatalf("cannot load certificate pair: %v", err)
	}
	data, err := ioutil.ReadFile(cfg.CaCert)
	if err != nil {
		logger.Fatalf("cannot load cacertificate: %v", err)
	}
	capool := x509.NewCertPool()
	capool.AppendCertsFromPEM(data)

	tlscfg := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs: capool,
	}
	tlscfg.BuildNameToCertificate()

	input := make(chan *Packet, 100)
	output := make(chan *Packet, 100)
	envmap := make(map[string]string)
	env := MapEnvironment{Internal: envmap}

	client := &Client{Config: cfg, Conn: nil, Logger: logger, Connected: false, Input: input, Output: output}

	go BuildEnv(logger, envmap)

	go func() {
		for {
			if (client.Connected == false) {
				conn, err := tls.Dial("tcp", cfg.Server, tlscfg)
				if err != nil {
					logger.Printf("unabled to connect, will retry in 5 seconds: %v", err)
					time.Sleep(5 * time.Second)
					continue
				}
				client.Conn = conn
				client.Connected = true
				logger.Printf("connected")
			}
			p, err := ReadPacket(logger, client.Conn)
			if err != nil {
				logger.Printf("read error: %v", err)
				client.Connected = false
				client.Conn = nil
			} else {
				input <- p
			}
		}
	}()

	go func() {
		for {
			p := <- input
			client.HandleRequest(p, env)
		}
	}()

	go func() {
		for {
			p := <- output
			p.Host = cfg.Host
			if (client.Connected == false) {
				logger.Printf("not connected, dropping output packet")
				continue
			}
			err = WritePacket(logger, client.Conn, p)
			if err != nil {
				logger.Printf("failed to send packet")
			}
		}
	}()

	return client
}

func ReadPacket(logger *log.Logger, conn *tls.Conn) (*Packet, error) {

	inbuf := make([]byte, 1024)
	br, err := conn.Read(inbuf)
	if err != nil {
		return nil, err
	}
	if br == 0 {
		return nil, errors.New("disconnected")
	}
	if br < 4 {
		return nil, errors.New("short read")
	}

	br = br -4
	phead := inbuf[0:4]
	b := bytes.NewReader(phead)

	var pilen uint32
	err = binary.Read(b, binary.BigEndian, &pilen)
	plen := int(pilen) - 4
	if err != nil {
		conn.Close()
		return nil, err
	}

	pdata := inbuf[4:]
	for {
		if br >= plen {
			break
		}
		buf := make([]byte,(plen - br))
		nr, err := conn.Read(buf)
		if err != nil {
			conn.Close()
			return nil, err
		}
		if nr == 0 {
			return nil, err
		}
		pdata = append(pdata, buf...)
		br = br + nr
	}

	var p Packet
	err = json.Unmarshal(pdata[:plen], &p)

	if err != nil {
		return nil, err

	}
	return &p, nil
}

func WritePacket(logger *log.Logger, conn *tls.Conn, p *Packet) error {

	jsbuf, err := json.Marshal(&p)
	if err != nil {
		return err
	}

	logger.Printf("ready to send packet: %s", string(jsbuf))

	plen := len(jsbuf)
	pilen := uint32(plen)

	outlen := new(bytes.Buffer)
	binary.Write(outlen, binary.BigEndian, &pilen)

	jsbuf = append(outlen.Bytes(), jsbuf...)

	bw := 0
	for {
		if bw >= (plen + 4) {
			break
		}
		w, err := conn.Write(jsbuf)
		if err != nil {
			return err
		}
		bw = bw + w
	}
	return nil
}

func (client *Client) SendPacket(p Packet) {
	client.Output <- &p
}

func (client *Client) WaitForQuit() {
	var wg sync.WaitGroup
	wg.Add(1)
	wg.Wait()
}
<commit_msg>agent: @llambiel thinks the agent is too chatty<commit_after>package warp

import (
	"fmt"
	"errors"
	"bytes"
	"os"
	"time"
	"log"
	"log/syslog"
	"crypto/tls"
	"crypto/x509"
	"io/ioutil"
	"sync"
	"encoding/binary"
	"encoding/json"
)

type Client struct {
	Config Config
	Conn *tls.Conn
	Logger *log.Logger
	Connected bool
	Input chan *Packet
	Output chan *Packet
}

func NewClient(cfg Config) *Client {

	var logger *log.Logger
	var err error

	switch {
	case cfg.LogTo == "stdout":
		logger = log.New(os.Stdout, "warp: ", 0)
	case cfg.LogTo == "stderr":
		logger = log.New(os.Stderr, "warp: ", 0)
	case cfg.LogTo == "syslog":
		logger, err = syslog.NewLogger(syslog.LOG_INFO | syslog.LOG_DAEMON, 0)
		if err != nil {
			fmt.Printf("cannot initialize logger: %v", err)
			os.Exit(1)
		}
	default:
		fmt.Printf("cannot initialize logger, bad configuration")
		os.Exit(1)
	}

	cert, err := tls.LoadX509KeyPair(cfg.Cert, cfg.PrivKey)
	if err != nil {
		logger.Fatalf("cannot load certificate pair: %v", err)
	}
	data, err := ioutil.ReadFile(cfg.CaCert)
	if err != nil {
		logger.Fatalf("cannot load cacertificate: %v", err)
	}
	capool := x509.NewCertPool()
	capool.AppendCertsFromPEM(data)

	tlscfg := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs: capool,
	}
	tlscfg.BuildNameToCertificate()

	input := make(chan *Packet, 100)
	output := make(chan *Packet, 100)
	envmap := make(map[string]string)
	env := MapEnvironment{Internal: envmap}

	client := &Client{Config: cfg, Conn: nil, Logger: logger, Connected: false, Input: input, Output: output}

	go BuildEnv(logger, envmap)

	go func() {
		for {
			if (client.Connected == false) {
				conn, err := tls.Dial("tcp", cfg.Server, tlscfg)
				if err != nil {
					logger.Printf("unabled to connect, will retry in 5 seconds: %v", err)
					time.Sleep(5 * time.Second)
					continue
				}
				client.Conn = conn
				client.Connected = true
				logger.Printf("connected")
			}
			p, err := ReadPacket(logger, client.Conn)
			if err != nil {
				logger.Printf("read error: %v", err)
				client.Connected = false
				client.Conn = nil
			} else {
				input <- p
			}
		}
	}()

	go func() {
		for {
			p := <- input
			client.HandleRequest(p, env)
		}
	}()

	go func() {
		for {
			p := <- output
			p.Host = cfg.Host
			if (client.Connected == false) {
				logger.Printf("not connected, dropping output packet")
				continue
			}
			err = WritePacket(logger, client.Conn, p)
			if err != nil {
				logger.Printf("failed to send packet")
			}
		}
	}()

	return client
}

func ReadPacket(logger *log.Logger, conn *tls.Conn) (*Packet, error) {

	inbuf := make([]byte, 1024)
	br, err := conn.Read(inbuf)
	if err != nil {
		return nil, err
	}
	if br == 0 {
		return nil, errors.New("disconnected")
	}
	if br < 4 {
		return nil, errors.New("short read")
	}

	br = br -4
	phead := inbuf[0:4]
	b := bytes.NewReader(phead)

	var pilen uint32
	err = binary.Read(b, binary.BigEndian, &pilen)
	plen := int(pilen) - 4
	if err != nil {
		conn.Close()
		return nil, err
	}

	pdata := inbuf[4:]
	for {
		if br >= plen {
			break
		}
		buf := make([]byte,(plen - br))
		nr, err := conn.Read(buf)
		if err != nil {
			conn.Close()
			return nil, err
		}
		if nr == 0 {
			return nil, err
		}
		pdata = append(pdata, buf...)
		br = br + nr
	}

	var p Packet
	err = json.Unmarshal(pdata[:plen], &p)

	if err != nil {
		return nil, err

	}
	return &p, nil
}

func WritePacket(logger *log.Logger, conn *tls.Conn, p *Packet) error {

	jsbuf, err := json.Marshal(&p)
	if err != nil {
		return err
	}

	plen := len(jsbuf)
	pilen := uint32(plen)

	outlen := new(bytes.Buffer)
	binary.Write(outlen, binary.BigEndian, &pilen)

	jsbuf = append(outlen.Bytes(), jsbuf...)

	bw := 0
	for {
		if bw >= (plen + 4) {
			break
		}
		w, err := conn.Write(jsbuf)
		if err != nil {
			return err
		}
		bw = bw + w
	}
	return nil
}

func (client *Client) SendPacket(p Packet) {
	client.Output <- &p
}

func (client *Client) WaitForQuit() {
	var wg sync.WaitGroup
	wg.Add(1)
	wg.Wait()
}
<|endoftext|>


<commit_before>package car

import (
    "fmt"
    "path"
    "regexp"
    "strings"

    "github.com/jtyr/gbt/gbt/core/utils"
)

// ModelElement is an element of which the car model is constructed from.
type ModelElement struct {
    Bg string
    Fg string
    Fm string
    Text string
}

// Car is a type defining the model of the car.
type Car struct {
    Model map[string]ModelElement
    Display bool
    Sep string
    Wrap bool
    Params map[string]interface{}
}

// Shell type.
var Shell string = utils.GetEnv("GBT_SHELL", path.Base(utils.GetEnv("SHELL", "zsh")))

// List of named colors and their codes.
var colors = map[string]string {
    "black":          "0",
    "red":            "1",
    "green":          "2",
    "yellow":         "3",
    "blue":           "4",
    "magenta":        "5",
    "cyan":           "6",
    "light_gray":     "7",
    "dark_gray":      "8",
    "light_red":      "9",
    "light_green":   "10",
    "light_yellow":  "11",
    "light_blue":    "12",
    "light_magenta": "13",
    "light_cyan":    "14",
    "white":         "15",
}

// SetParamStr sets string value to a parameter.
func (c *Car) SetParamStr(name, value string) {
    if c.Params == nil {
        c.Params = make(map[string]interface{})
    }

    c.Params[name] = value
}

// GetModel returns the Model value.
func (c *Car) GetModel() map[string]ModelElement {
    return c.Model
}

// GetDisplay returns the Display value.
func (c *Car) GetDisplay() bool {
    return c.Display
}

// GetSep returns the Sep value.
func (c *Car) GetSep() string {
    return c.Sep
}

// GetWrap returns the Wrap value.
func (c *Car) GetWrap() bool {
    return c.Wrap
}

var reTemplating = regexp.MustCompile(`{{\s*(\w+)\s*}}`)

// Format initiates replacement of all templating elements.
func (c *Car) Format() string {
    if ! c.Display {
        return ""
    }

    text := fmt.Sprintf("%s%s", c.DecorateElement("root", "", "", "", ""), c.Model["root"].Text)

    for range make([]int, 10) {
        match := reTemplating.MatchString(text)

        if match {
            text = reTemplating.ReplaceAllStringFunc(text, c.replaceElement)
        } else {
            break
        }
    }

    return text
}

// Replaces the specific templating element.
func (c *Car) replaceElement(format string) string {
    match := reTemplating.FindStringSubmatch(format)[1]

    if _, ok := c.Model[match]; ! ok {
        return format
    }

    return fmt.Sprintf(
        "%s%s",
        c.DecorateElement(match, "", "", "", ""),
        c.DecorateElement("root", "", "", "", ""))
}

// DecorateElement decorates the element text with its colors and formatting.
func (c *Car) DecorateElement(element, bg, fg, fm, text string) string {
    fmEnd := ""

    if element != "" {
        e := c.Model[element]

        if element != "root" {
            text = e.Text
        } else {
            text = ""
        }

        bg = c.GetColor(e.Bg, false)
        fg = c.GetColor(e.Fg, true)
        fm = c.GetFormat(e.Fm, false)

        if fm != c.GetFormat("empty", false) {
            fmEnd = c.GetFormat(e.Fm, true)
        } else {
            fm = ""
        }
    }

    return fmt.Sprintf("%s%s%s%s%s", bg, fg, fm, text, fmEnd)
}

// Patterns to parse the color codes
var reColorNumber = regexp.MustCompile(`^\d{1,3}$`)
var reRgbTriplet = regexp.MustCompile(`^\d{1,3};\d{1,3};\d{1,3}$`)

// GetColor returns color sequence based on the color name or code.
func (c *Car) GetColor(name string, isFg bool) (ret string) {
    kind := 4
    seq := ""
    esc := "\x1b"

    if isFg {
        kind = 3
    }

    if Shell == "_bash" {
        esc = "\\e"
    }

    if name == "RESETALL" {
        seq = fmt.Sprintf("%s[0m", esc)
    } else if name == "default" {
        // Default
        seq = fmt.Sprintf("%s[%d9m", esc, kind)
    } else {
        if val, ok := colors[name]; ok {
            // Named color
            seq = fmt.Sprintf("%s[%d8;5;%sm", esc, kind, val)
        } else if match := reColorNumber.MatchString(name); match {
            // Color number
            seq = fmt.Sprintf("%s[%d8;5;%sm", esc, kind, name)
        } else if match := reRgbTriplet.MatchString(name); match {
            // RGB color
            seq = fmt.Sprintf("%s[%d8;2;%sm", esc, kind, name)
        } else {
            // If anything else, use default
            seq = fmt.Sprintf("%s[%d9m", esc, kind)
        }
    }

    ret = DecorateShell(seq)

    return
}

// GetFormat returns formatting sequence based on the format name.
func (c *Car) GetFormat(name string, end bool) (ret string) {
    seq := ""
    kind := 0
    esc := "\x1b"

    if end {
        kind = 2
    }

    if Shell == "_bash" {
        esc = "\\e"
    }

    if strings.Contains(name, "bold") {
        seq += fmt.Sprintf("%s[%d1m", esc, kind)
    }

    if strings.Contains(name, "underline") {
        seq += fmt.Sprintf("%s[%d4m", esc, kind)
    }

    if strings.Contains(name, "blink") {
        seq += fmt.Sprintf("%s[%d5m", esc, kind)
    }

    ret = DecorateShell(seq)

    return
}

// DecorateShell decorates the string with shell-specific closure.
func DecorateShell(seq string) (ret string) {
    if Shell == "zsh" {
        ret = fmt.Sprintf("%%{%s%%}", seq)
    } else if Shell == "_bash" {
        ret = fmt.Sprintf("\\[%s\\]", seq)
    } else {
        ret = fmt.Sprintf("\001%s\002", seq)
    }

    return
}
<commit_msg>Use string instead of int for the kind var; make shell decoration private<commit_after>package car

import (
    "fmt"
    "path"
    "regexp"
    "strings"

    "github.com/jtyr/gbt/gbt/core/utils"
)

// ModelElement is an element of which the car model is constructed from.
type ModelElement struct {
    Bg string
    Fg string
    Fm string
    Text string
}

// Car is a type defining the model of the car.
type Car struct {
    Model map[string]ModelElement
    Display bool
    Sep string
    Wrap bool
    Params map[string]interface{}
}

// Shell type.
var Shell string = utils.GetEnv("GBT_SHELL", path.Base(utils.GetEnv("SHELL", "zsh")))

// List of named colors and their codes.
var colors = map[string]string {
    "black":          "0",
    "red":            "1",
    "green":          "2",
    "yellow":         "3",
    "blue":           "4",
    "magenta":        "5",
    "cyan":           "6",
    "light_gray":     "7",
    "dark_gray":      "8",
    "light_red":      "9",
    "light_green":   "10",
    "light_yellow":  "11",
    "light_blue":    "12",
    "light_magenta": "13",
    "light_cyan":    "14",
    "white":         "15",
}

// SetParamStr sets string value to a parameter.
func (c *Car) SetParamStr(name, value string) {
    if c.Params == nil {
        c.Params = make(map[string]interface{})
    }

    c.Params[name] = value
}

// GetModel returns the Model value.
func (c *Car) GetModel() map[string]ModelElement {
    return c.Model
}

// GetDisplay returns the Display value.
func (c *Car) GetDisplay() bool {
    return c.Display
}

// GetSep returns the Sep value.
func (c *Car) GetSep() string {
    return c.Sep
}

// GetWrap returns the Wrap value.
func (c *Car) GetWrap() bool {
    return c.Wrap
}

var reTemplating = regexp.MustCompile(`{{\s*(\w+)\s*}}`)

// Format initiates replacement of all templating elements.
func (c *Car) Format() string {
    if ! c.Display {
        return ""
    }

    text := fmt.Sprintf("%s%s", c.DecorateElement("root", "", "", "", ""), c.Model["root"].Text)

    for range make([]int, 10) {
        match := reTemplating.MatchString(text)

        if match {
            text = reTemplating.ReplaceAllStringFunc(text, c.replaceElement)
        } else {
            break
        }
    }

    return text
}

// Replaces the specific templating element.
func (c *Car) replaceElement(format string) string {
    match := reTemplating.FindStringSubmatch(format)[1]

    if _, ok := c.Model[match]; ! ok {
        return format
    }

    return fmt.Sprintf(
        "%s%s",
        c.DecorateElement(match, "", "", "", ""),
        c.DecorateElement("root", "", "", "", ""))
}

// DecorateElement decorates the element text with its colors and formatting.
func (c *Car) DecorateElement(element, bg, fg, fm, text string) string {
    fmEnd := ""

    if element != "" {
        e := c.Model[element]

        if element != "root" {
            text = e.Text
        } else {
            text = ""
        }

        bg = c.GetColor(e.Bg, false)
        fg = c.GetColor(e.Fg, true)
        fm = c.GetFormat(e.Fm, false)

        if fm != c.GetFormat("empty", false) {
            fmEnd = c.GetFormat(e.Fm, true)
        } else {
            fm = ""
        }
    }

    return fmt.Sprintf("%s%s%s%s%s", bg, fg, fm, text, fmEnd)
}

// Patterns to parse the color codes
var reColorNumber = regexp.MustCompile(`^\d{1,3}$`)
var reRgbTriplet = regexp.MustCompile(`^\d{1,3};\d{1,3};\d{1,3}$`)

// GetColor returns color sequence based on the color name or code.
func (c *Car) GetColor(name string, isFg bool) (ret string) {
    kind := "4"
    seq := ""
    esc := "\x1b"

    if isFg {
        kind = "3"
    }

    if Shell == "_bash" {
        esc = "\\e"
    }

    if name == "RESETALL" {
        seq = fmt.Sprintf("%s[0m", esc)
    } else if name == "default" {
        // Default
        seq = fmt.Sprintf("%s[%s9m", esc, kind)
    } else {
        if val, ok := colors[name]; ok {
            // Named color
            seq = fmt.Sprintf("%s[%s8;5;%sm", esc, kind, val)
        } else if match := reColorNumber.MatchString(name); match {
            // Color number
            seq = fmt.Sprintf("%s[%s8;5;%sm", esc, kind, name)
        } else if match := reRgbTriplet.MatchString(name); match {
            // RGB color
            seq = fmt.Sprintf("%s[%s8;2;%sm", esc, kind, name)
        } else {
            // If anything else, use default
            seq = fmt.Sprintf("%s[%s9m", esc, kind)
        }
    }

    ret = decorateShell(seq)

    return
}

// GetFormat returns formatting sequence based on the format name.
func (c *Car) GetFormat(name string, end bool) (ret string) {
    seq := ""
    kind := ""
    esc := "\x1b"

    if end {
        kind = "2"
    }

    if Shell == "_bash" {
        esc = "\\e"
    }

    if strings.Contains(name, "bold") {
        seq += fmt.Sprintf("%s[%s1m", esc, kind)
    }

    if strings.Contains(name, "underline") {
        seq += fmt.Sprintf("%s[%s4m", esc, kind)
    }

    if strings.Contains(name, "blink") {
        seq += fmt.Sprintf("%s[%s5m", esc, kind)
    }

    ret = decorateShell(seq)

    return
}

// decorateShell decorates the string with shell-specific closure.
func decorateShell(seq string) (ret string) {
    if Shell == "zsh" {
        ret = fmt.Sprintf("%%{%s%%}", seq)
    } else if Shell == "_bash" {
        ret = fmt.Sprintf("\\[%s\\]", seq)
    } else {
        ret = fmt.Sprintf("\001%s\002", seq)
    }

    return
}
<|endoftext|>