<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package json_test

import (
	"fmt"
	"reflect"
	"strings"
	"testing"

	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/runtime/schema"
	"k8s.io/kubernetes/pkg/runtime/serializer/json"
	"k8s.io/kubernetes/pkg/util/diff"
)

type testDecodable struct {
	Other string
	Value int `json:"value"`
	gvk   schema.GroupVersionKind
}

func (d *testDecodable) GetObjectKind() schema.ObjectKind                { return d }
func (d *testDecodable) SetGroupVersionKind(gvk schema.GroupVersionKind) { d.gvk = gvk }
func (d *testDecodable) GroupVersionKind() schema.GroupVersionKind       { return d.gvk }

func TestDecode(t *testing.T) {
	testCases := []struct {
		creater runtime.ObjectCreater
		typer   runtime.ObjectTyper
		yaml    bool
		pretty  bool

		data       []byte
		defaultGVK *schema.GroupVersionKind
		into       runtime.Object

		errFn          func(error) bool
		expectedObject runtime.Object
		expectedGVK    *schema.GroupVersionKind
	}{
		{
			data: []byte("{}"),

			expectedGVK: &schema.GroupVersionKind{},
			errFn:       func(err error) bool { return strings.Contains(err.Error(), "Object 'Kind' is missing in") },
		},
		{
			data:       []byte("{}"),
			defaultGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:    &mockCreater{err: fmt.Errorf("fake error")},

			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			errFn:       func(err error) bool { return err.Error() == "fake error" },
		},
		{
			data:       []byte("{}"),
			defaultGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:    &mockCreater{err: fmt.Errorf("fake error")},

			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			errFn:       func(err error) bool { return err.Error() == "fake error" },
		},
		{
			data:           []byte("{}"),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
		},

		// version without group is not defaulted
		{
			data:           []byte(`{"apiVersion":"blah"}`),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "", Version: "blah"},
		},
		// group without version is defaulted
		{
			data:           []byte(`{"apiVersion":"other/"}`),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
		},

		// accept runtime.Unknown as into and bypass creator
		{
			data: []byte(`{}`),
			into: &runtime.Unknown{},

			expectedGVK: &schema.GroupVersionKind{},
			expectedObject: &runtime.Unknown{
				Raw:         []byte(`{}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},
		{
			data: []byte(`{"test":"object"}`),
			into: &runtime.Unknown{},

			expectedGVK: &schema.GroupVersionKind{},
			expectedObject: &runtime.Unknown{
				Raw:         []byte(`{"test":"object"}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},
		{
			data:        []byte(`{"test":"object"}`),
			into:        &runtime.Unknown{},
			defaultGVK:  &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.Unknown{
				TypeMeta:    runtime.TypeMeta{APIVersion: "other/blah", Kind: "Test"},
				Raw:         []byte(`{"test":"object"}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},

		// unregistered objects can be decoded into directly
		{
			data:        []byte(`{"kind":"Test","apiVersion":"other/blah","value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{err: runtime.NewNotRegisteredErr(schema.GroupVersionKind{}, nil)},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},
		// registered types get defaulted by the into object kind
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},
		// registered types get defaulted by the into object kind even without version, but return an error
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: ""}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: ""},
			errFn:       func(err error) bool { return strings.Contains(err.Error(), "Object 'apiVersion' is missing in") },
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},

		// runtime.VersionedObjects are decoded
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &runtime.VersionedObjects{Objects: []runtime.Object{}},
			creater:     &mockCreater{obj: &testDecodable{}},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			defaultGVK:  &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.VersionedObjects{
				Objects: []runtime.Object{
					&testDecodable{
						Other: "test",
						Value: 1,
					},
				},
			},
		},
		// runtime.VersionedObjects with an object are decoded into
		{
			data:        []byte(`{"Other":"test"}`),
			into:        &runtime.VersionedObjects{Objects: []runtime.Object{&testDecodable{Value: 2}}},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.VersionedObjects{
				Objects: []runtime.Object{
					&testDecodable{
						Other: "test",
						Value: 2,
					},
				},
			},
		},
	}

	for i, test := range testCases {
		var s runtime.Serializer
		if test.yaml {
			s = json.NewYAMLSerializer(json.DefaultMetaFactory, test.creater, test.typer)
		} else {
			s = json.NewSerializer(json.DefaultMetaFactory, test.creater, test.typer, test.pretty)
		}
		obj, gvk, err := s.Decode([]byte(test.data), test.defaultGVK, test.into)

		if !reflect.DeepEqual(test.expectedGVK, gvk) {
			t.Errorf("%d: unexpected GVK: %v", i, gvk)
		}

		switch {
		case err == nil && test.errFn != nil:
			t.Errorf("%d: failed: %v", i, err)
			continue
		case err != nil && test.errFn == nil:
			t.Errorf("%d: failed: %v", i, err)
			continue
		case err != nil:
			if !test.errFn(err) {
				t.Errorf("%d: failed: %v", i, err)
			}
			if obj != nil {
				t.Errorf("%d: should have returned nil object", i)
			}
			continue
		}

		if test.into != nil && test.into != obj {
			t.Errorf("%d: expected into to be returned: %v", i, obj)
			continue
		}

		if !reflect.DeepEqual(test.expectedObject, obj) {
			t.Errorf("%d: unexpected object:\n%s", i, diff.ObjectGoPrintSideBySide(test.expectedObject, obj))
		}
	}
}

type mockCreater struct {
	apiVersion string
	kind       string
	err        error
	obj        runtime.Object
}

func (c *mockCreater) New(kind schema.GroupVersionKind) (runtime.Object, error) {
	c.apiVersion, c.kind = kind.GroupVersion().String(), kind.Kind
	return c.obj, c.err
}

type mockTyper struct {
	gvk *schema.GroupVersionKind
	err error
}

func (t *mockTyper) ObjectKinds(obj runtime.Object) ([]schema.GroupVersionKind, bool, error) {
	if t.gvk == nil {
		return nil, false, t.err
	}
	return []schema.GroupVersionKind{*t.gvk}, false, t.err
}

func (t *mockTyper) Recognizes(_ schema.GroupVersionKind) bool {
	return false
}
<commit_msg>delete one testcase from json_test.go for duplicated data<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package json_test

import (
	"fmt"
	"reflect"
	"strings"
	"testing"

	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/runtime/schema"
	"k8s.io/kubernetes/pkg/runtime/serializer/json"
	"k8s.io/kubernetes/pkg/util/diff"
)

type testDecodable struct {
	Other string
	Value int `json:"value"`
	gvk   schema.GroupVersionKind
}

func (d *testDecodable) GetObjectKind() schema.ObjectKind                { return d }
func (d *testDecodable) SetGroupVersionKind(gvk schema.GroupVersionKind) { d.gvk = gvk }
func (d *testDecodable) GroupVersionKind() schema.GroupVersionKind       { return d.gvk }

func TestDecode(t *testing.T) {
	testCases := []struct {
		creater runtime.ObjectCreater
		typer   runtime.ObjectTyper
		yaml    bool
		pretty  bool

		data       []byte
		defaultGVK *schema.GroupVersionKind
		into       runtime.Object

		errFn          func(error) bool
		expectedObject runtime.Object
		expectedGVK    *schema.GroupVersionKind
	}{
		{
			data: []byte("{}"),

			expectedGVK: &schema.GroupVersionKind{},
			errFn:       func(err error) bool { return strings.Contains(err.Error(), "Object 'Kind' is missing in") },
		},
		{
			data:       []byte("{}"),
			defaultGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:    &mockCreater{err: fmt.Errorf("fake error")},

			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			errFn:       func(err error) bool { return err.Error() == "fake error" },
		},
		{
			data:           []byte("{}"),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
		},

		// version without group is not defaulted
		{
			data:           []byte(`{"apiVersion":"blah"}`),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "", Version: "blah"},
		},
		// group without version is defaulted
		{
			data:           []byte(`{"apiVersion":"other/"}`),
			defaultGVK:     &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			creater:        &mockCreater{obj: &testDecodable{}},
			expectedObject: &testDecodable{},
			expectedGVK:    &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
		},

		// accept runtime.Unknown as into and bypass creator
		{
			data: []byte(`{}`),
			into: &runtime.Unknown{},

			expectedGVK: &schema.GroupVersionKind{},
			expectedObject: &runtime.Unknown{
				Raw:         []byte(`{}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},
		{
			data: []byte(`{"test":"object"}`),
			into: &runtime.Unknown{},

			expectedGVK: &schema.GroupVersionKind{},
			expectedObject: &runtime.Unknown{
				Raw:         []byte(`{"test":"object"}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},
		{
			data:        []byte(`{"test":"object"}`),
			into:        &runtime.Unknown{},
			defaultGVK:  &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.Unknown{
				TypeMeta:    runtime.TypeMeta{APIVersion: "other/blah", Kind: "Test"},
				Raw:         []byte(`{"test":"object"}`),
				ContentType: runtime.ContentTypeJSON,
			},
		},

		// unregistered objects can be decoded into directly
		{
			data:        []byte(`{"kind":"Test","apiVersion":"other/blah","value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{err: runtime.NewNotRegisteredErr(schema.GroupVersionKind{}, nil)},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},
		// registered types get defaulted by the into object kind
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},
		// registered types get defaulted by the into object kind even without version, but return an error
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &testDecodable{},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: ""}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: ""},
			errFn:       func(err error) bool { return strings.Contains(err.Error(), "Object 'apiVersion' is missing in") },
			expectedObject: &testDecodable{
				Other: "test",
				Value: 1,
			},
		},

		// runtime.VersionedObjects are decoded
		{
			data:        []byte(`{"value":1,"Other":"test"}`),
			into:        &runtime.VersionedObjects{Objects: []runtime.Object{}},
			creater:     &mockCreater{obj: &testDecodable{}},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			defaultGVK:  &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.VersionedObjects{
				Objects: []runtime.Object{
					&testDecodable{
						Other: "test",
						Value: 1,
					},
				},
			},
		},
		// runtime.VersionedObjects with an object are decoded into
		{
			data:        []byte(`{"Other":"test"}`),
			into:        &runtime.VersionedObjects{Objects: []runtime.Object{&testDecodable{Value: 2}}},
			typer:       &mockTyper{gvk: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"}},
			expectedGVK: &schema.GroupVersionKind{Kind: "Test", Group: "other", Version: "blah"},
			expectedObject: &runtime.VersionedObjects{
				Objects: []runtime.Object{
					&testDecodable{
						Other: "test",
						Value: 2,
					},
				},
			},
		},
	}

	for i, test := range testCases {
		var s runtime.Serializer
		if test.yaml {
			s = json.NewYAMLSerializer(json.DefaultMetaFactory, test.creater, test.typer)
		} else {
			s = json.NewSerializer(json.DefaultMetaFactory, test.creater, test.typer, test.pretty)
		}
		obj, gvk, err := s.Decode([]byte(test.data), test.defaultGVK, test.into)

		if !reflect.DeepEqual(test.expectedGVK, gvk) {
			t.Errorf("%d: unexpected GVK: %v", i, gvk)
		}

		switch {
		case err == nil && test.errFn != nil:
			t.Errorf("%d: failed: %v", i, err)
			continue
		case err != nil && test.errFn == nil:
			t.Errorf("%d: failed: %v", i, err)
			continue
		case err != nil:
			if !test.errFn(err) {
				t.Errorf("%d: failed: %v", i, err)
			}
			if obj != nil {
				t.Errorf("%d: should have returned nil object", i)
			}
			continue
		}

		if test.into != nil && test.into != obj {
			t.Errorf("%d: expected into to be returned: %v", i, obj)
			continue
		}

		if !reflect.DeepEqual(test.expectedObject, obj) {
			t.Errorf("%d: unexpected object:\n%s", i, diff.ObjectGoPrintSideBySide(test.expectedObject, obj))
		}
	}
}

type mockCreater struct {
	apiVersion string
	kind       string
	err        error
	obj        runtime.Object
}

func (c *mockCreater) New(kind schema.GroupVersionKind) (runtime.Object, error) {
	c.apiVersion, c.kind = kind.GroupVersion().String(), kind.Kind
	return c.obj, c.err
}

type mockTyper struct {
	gvk *schema.GroupVersionKind
	err error
}

func (t *mockTyper) ObjectKinds(obj runtime.Object) ([]schema.GroupVersionKind, bool, error) {
	if t.gvk == nil {
		return nil, false, t.err
	}
	return []schema.GroupVersionKind{*t.gvk}, false, t.err
}

func (t *mockTyper) Recognizes(_ schema.GroupVersionKind) bool {
	return false
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package ecs

import (
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/AlecAivazis/survey/v2/terminal"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"gopkg.in/ini.v1"

	"github.com/docker/compose-cli/context/store"
	"github.com/docker/compose-cli/errdefs"
	"github.com/docker/compose-cli/prompt"
)

type contextElements struct {
	AccessKey    string
	SecretKey    string
	SessionToken string
	Profile      string
	Region       string
	CredsFromEnv bool
}

func (c contextElements) HaveRequiredCredentials() bool {
	if c.AccessKey != "" && c.SecretKey != "" {
		return true
	}
	return false
}

type contextCreateAWSHelper struct {
	user prompt.UI
}

func newContextCreateHelper() contextCreateAWSHelper {
	return contextCreateAWSHelper{
		user: prompt.User{},
	}
}

func getEnvVars() contextElements {
	c := contextElements{}
	profile := os.Getenv("AWS_PROFILE")
	if profile != "" {
		c.Profile = profile
	}

	p := credentials.EnvProvider{}
	creds, err := p.Retrieve()
	if err != nil {
		return c
	}
	c.AccessKey = creds.AccessKeyID
	c.SecretKey = creds.SecretAccessKey
	c.SessionToken = creds.SessionToken
	return c
}

func (h contextCreateAWSHelper) createProfile(name string, c *contextElements) error {
	if c != nil {
		if c.AccessKey != "" && c.SecretKey != "" {
			return h.saveCredentials(name, c.AccessKey, c.SecretKey)
		}
		accessKey, secretKey, err := h.askCredentials()

		if err != nil {
			return err
		}
		c.AccessKey = accessKey
		c.SecretKey = secretKey
		return h.saveCredentials(name, c.AccessKey, c.SecretKey)
	}

	accessKey, secretKey, err := h.askCredentials()
	if err != nil {
		return err
	}
	if accessKey != "" && secretKey != "" {
		return h.saveCredentials(name, accessKey, secretKey)
	}
	return nil
}

func (h contextCreateAWSHelper) createContext(c *contextElements, description string) (interface{}, string) {
	if c.Profile == "default" {
		c.Profile = ""
	}
	description = strings.TrimSpace(
		fmt.Sprintf("%s (%s)", description, c.Region))
	if c.CredsFromEnv {
		return store.EcsContext{
			CredentialsFromEnv: c.CredsFromEnv,
			Profile:            c.Profile,
			Region:             c.Region,
		}, description
	}
	return store.EcsContext{
		Profile: c.Profile,
		Region:  c.Region,
	}, description
}

func (h contextCreateAWSHelper) createContextData(_ context.Context, opts ContextParams) (interface{}, string, error) {
	creds := contextElements{}

	options := []string{
		"Use AWS credentials set via environment variables",
		"Create a new profile with AWS credentials",
		"Select from existing local AWS profiles",
	}
	//if creds.HaveRequiredProps() {
	selected, err := h.user.Select("Would you like to create your context based on", options)
	if err != nil {
		if err == terminal.InterruptErr {
			return nil, "", errdefs.ErrCanceled
		}
		return nil, "", err
	}
	if creds.Region == "" {
		creds.Region = opts.Region
	}
	if creds.Profile == "" {
		creds.Profile = opts.Profile
	}

	switch selected {
	case 0:
		creds.CredsFromEnv = true
		// confirm region profile should target
		if creds.Region == "" {
			creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
			if err != nil {
				return nil, "", err
			}
		}
		/*if creds.Profile == "" {
			creds.Profile = opts.Name

		}
		fmt.Printf("Saving credentials under profile %s\n", creds.Profile)
		h.createProfile(creds.Profile, &creds)*/
	case 1:
		accessKey, secretKey, err := h.askCredentials()
		if err != nil {
			return nil, "", err
		}
		creds.AccessKey = accessKey
		creds.SecretKey = secretKey
		// we need a region set -- either read it from profile or prompt user

		// prompt for the region to use with this context
		creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
		if err != nil {
			return nil, "", err
		}
		// save as a profile
		if creds.Profile == "" {
			creds.Profile = opts.Name
		}
		fmt.Printf("Saving credentials under profile %s\n", creds.Profile)
		h.createProfile(creds.Profile, &creds)

	case 2:
		profilesList, err := h.getProfiles()
		if err != nil {
			return nil, "", err
		}
		// choose profile
		creds.Profile, err = h.chooseProfile(profilesList)
		if err != nil {
			return nil, "", err
		}
		if creds.Region == "" {
			creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
			if err != nil {
				return nil, "", err
			}
		}
	}

	//os.Exit(0)
	ecsCtx, descr := h.createContext(&creds, opts.Description)
	return ecsCtx, descr, nil
}

func (h contextCreateAWSHelper) saveCredentials(profile string, accessKeyID string, secretAccessKey string) error {
	p := credentials.SharedCredentialsProvider{Profile: profile}
	_, err := p.Retrieve()
	if err == nil {
		return fmt.Errorf("credentials already exist")
	}

	if err.(awserr.Error).Code() == "SharedCredsLoad" && err.(awserr.Error).Message() == "failed to load shared credentials file" {
		_, err := os.Create(p.Filename)
		if err != nil {
			return err
		}
	}
	credIni, err := ini.Load(p.Filename)
	if err != nil {
		return err
	}
	section, err := credIni.NewSection(profile)
	if err != nil {
		return err
	}
	_, err = section.NewKey("aws_access_key_id", accessKeyID)
	if err != nil {
		return err
	}
	_, err = section.NewKey("aws_secret_access_key", secretAccessKey)
	if err != nil {
		return err
	}
	return credIni.SaveTo(p.Filename)
}

func (h contextCreateAWSHelper) getProfiles() ([]string, error) {
	profiles := []string{}
	// parse both .aws/credentials and .aws/config for profiles
	configFiles := map[string]bool{
		defaults.SharedCredentialsFilename(): false,
		defaults.SharedConfigFilename():      true,
	}
	for f, prefix := range configFiles {
		sections, err := loadIniFile(f, prefix)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		for key := range sections {
			name := strings.ToLower(key)
			if !contains(profiles, name) {
				profiles = append(profiles, name)
			}
		}
	}
	return profiles, nil
}

func (h contextCreateAWSHelper) chooseProfile(profiles []string) (string, error) {
	options := []string{}
	options = append(options, profiles...)

	selected, err := h.user.Select("Select AWS Profile", options)
	if err != nil {
		if err == terminal.InterruptErr {
			return "", errdefs.ErrCanceled
		}
		return "", err
	}
	profile := options[selected]
	return profile, nil
}

func (h contextCreateAWSHelper) getRegionSuggestion(region string, profile string) (string, error) {
	if profile == "" {
		profile = "default"
	}
	// only load ~/.aws/config
	awsConfig := defaults.SharedConfigFilename()
	configIni, err := ini.Load(awsConfig)
	if err != nil {
		if !os.IsNotExist(err) {
			return "", err
		}
		configIni = ini.Empty()
	}

	var f func(string, string) string
	f = func(r string, p string) string {
		section, err := configIni.GetSection(p)
		if err == nil {
			reg, err := section.GetKey("region")
			if err == nil {
				r = reg.Value()
			}
		}
		if r == "" {
			switch p {
			case "":
				return "us-east-1"
			case "default":
				return f(r, "")
			}
			return f(r, "default")
		}
		return r
	}

	if profile != "default" {
		profile = fmt.Sprintf("profile %s", profile)
	}
	return f(region, profile), nil
}

func (h contextCreateAWSHelper) chooseRegion(region string, profile string) (string, error) {
	suggestion, err := h.getRegionSuggestion(region, profile)
	if err != nil {
		return "", err
	}
	// promp user for region
	region, err = h.user.Input("Region", suggestion)
	if err != nil {
		return "", err
	}
	if region == "" {
		return "", fmt.Errorf("region cannot be empty")
	}
	return region, nil
}

func (h contextCreateAWSHelper) askCredentials() (string, string, error) {
	/*confirm, err := h.user.Confirm("Enter AWS credentials", false)
	if err != nil {
		return "", "", err
	}
	if !confirm {
		return "", "", nil
	}*/

	accessKeyID, err := h.user.Input("AWS Access Key ID", "")
	if err != nil {
		return "", "", err
	}
	secretAccessKey, err := h.user.Password("Enter AWS Secret Access Key")
	if err != nil {
		return "", "", err
	}
	// validate access ID and password
	if len(accessKeyID) < 3 || len(secretAccessKey) < 3 {
		return "", "", fmt.Errorf("AWS Access/Secret Access Key must have more than 3 characters")
	}
	return accessKeyID, secretAccessKey, nil
}

func contains(values []string, value string) bool {
	for _, v := range values {
		if v == value {
			return true
		}
	}
	return false
}

func loadIniFile(path string, prefix bool) (map[string]ini.Section, error) {
	profiles := map[string]ini.Section{}
	credIni, err := ini.Load(path)
	if err != nil {
		return nil, err
	}
	for _, section := range credIni.Sections() {
		if prefix && strings.HasPrefix(section.Name(), "profile ") {
			profiles[section.Name()[len("profile "):]] = *section
		} else if !prefix || section.Name() == "default" {
			profiles[section.Name()] = *section
		}
	}
	return profiles, nil
}
<commit_msg>cleanup<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

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

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

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

package ecs

import (
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/AlecAivazis/survey/v2/terminal"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"gopkg.in/ini.v1"

	"github.com/docker/compose-cli/context/store"
	"github.com/docker/compose-cli/errdefs"
	"github.com/docker/compose-cli/prompt"
)

type contextElements struct {
	AccessKey    string
	SecretKey    string
	SessionToken string
	Profile      string
	Region       string
	CredsFromEnv bool
}

func (c contextElements) HaveRequiredCredentials() bool {
	if c.AccessKey != "" && c.SecretKey != "" {
		return true
	}
	return false
}

type contextCreateAWSHelper struct {
	user prompt.UI
}

func newContextCreateHelper() contextCreateAWSHelper {
	return contextCreateAWSHelper{
		user: prompt.User{},
	}
}

func getEnvVars() contextElements {
	c := contextElements{}
	profile := os.Getenv("AWS_PROFILE")
	if profile != "" {
		c.Profile = profile
	}

	p := credentials.EnvProvider{}
	creds, err := p.Retrieve()
	if err != nil {
		return c
	}
	c.AccessKey = creds.AccessKeyID
	c.SecretKey = creds.SecretAccessKey
	c.SessionToken = creds.SessionToken
	return c
}

func (h contextCreateAWSHelper) createProfile(name string, c *contextElements) error {
	if c != nil {
		if c.AccessKey != "" && c.SecretKey != "" {
			return h.saveCredentials(name, c.AccessKey, c.SecretKey)
		}
		accessKey, secretKey, err := h.askCredentials()

		if err != nil {
			return err
		}
		c.AccessKey = accessKey
		c.SecretKey = secretKey
		return h.saveCredentials(name, c.AccessKey, c.SecretKey)
	}

	accessKey, secretKey, err := h.askCredentials()
	if err != nil {
		return err
	}
	if accessKey != "" && secretKey != "" {
		return h.saveCredentials(name, accessKey, secretKey)
	}
	return nil
}

func (h contextCreateAWSHelper) createContext(c *contextElements, description string) (interface{}, string) {
	if c.Profile == "default" {
		c.Profile = ""
	}
	description = strings.TrimSpace(
		fmt.Sprintf("%s (%s)", description, c.Region))
	if c.CredsFromEnv {
		return store.EcsContext{
			CredentialsFromEnv: c.CredsFromEnv,
			Profile:            c.Profile,
			Region:             c.Region,
		}, description
	}
	return store.EcsContext{
		Profile: c.Profile,
		Region:  c.Region,
	}, description
}

func (h contextCreateAWSHelper) createContextData(_ context.Context, opts ContextParams) (interface{}, string, error) {
	creds := contextElements{}

	options := []string{
		"Use AWS credentials set via environment variables",
		"Create a new profile with AWS credentials",
		"Select from existing local AWS profiles",
	}
	//if creds.HaveRequiredProps() {
	selected, err := h.user.Select("Would you like to create your context based on", options)
	if err != nil {
		if err == terminal.InterruptErr {
			return nil, "", errdefs.ErrCanceled
		}
		return nil, "", err
	}
	if creds.Region == "" {
		creds.Region = opts.Region
	}
	if creds.Profile == "" {
		creds.Profile = opts.Profile
	}

	switch selected {
	case 0:
		creds.CredsFromEnv = true
		// confirm region profile should target
		if creds.Region == "" {
			creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
			if err != nil {
				return nil, "", err
			}
		}
	case 1:
		accessKey, secretKey, err := h.askCredentials()
		if err != nil {
			return nil, "", err
		}
		creds.AccessKey = accessKey
		creds.SecretKey = secretKey
		// we need a region set -- either read it from profile or prompt user

		// prompt for the region to use with this context
		creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
		if err != nil {
			return nil, "", err
		}
		// save as a profile
		if creds.Profile == "" {
			creds.Profile = opts.Name
		}
		fmt.Printf("Saving credentials under profile %s\n", creds.Profile)
		h.createProfile(creds.Profile, &creds)

	case 2:
		profilesList, err := h.getProfiles()
		if err != nil {
			return nil, "", err
		}
		// choose profile
		creds.Profile, err = h.chooseProfile(profilesList)
		if err != nil {
			return nil, "", err
		}
		if creds.Region == "" {
			creds.Region, err = h.chooseRegion(creds.Region, creds.Profile)
			if err != nil {
				return nil, "", err
			}
		}
	}

	ecsCtx, descr := h.createContext(&creds, opts.Description)
	return ecsCtx, descr, nil
}

func (h contextCreateAWSHelper) saveCredentials(profile string, accessKeyID string, secretAccessKey string) error {
	p := credentials.SharedCredentialsProvider{Profile: profile}
	_, err := p.Retrieve()
	if err == nil {
		return fmt.Errorf("credentials already exist")
	}

	if err.(awserr.Error).Code() == "SharedCredsLoad" && err.(awserr.Error).Message() == "failed to load shared credentials file" {
		_, err := os.Create(p.Filename)
		if err != nil {
			return err
		}
	}
	credIni, err := ini.Load(p.Filename)
	if err != nil {
		return err
	}
	section, err := credIni.NewSection(profile)
	if err != nil {
		return err
	}
	_, err = section.NewKey("aws_access_key_id", accessKeyID)
	if err != nil {
		return err
	}
	_, err = section.NewKey("aws_secret_access_key", secretAccessKey)
	if err != nil {
		return err
	}
	return credIni.SaveTo(p.Filename)
}

func (h contextCreateAWSHelper) getProfiles() ([]string, error) {
	profiles := []string{}
	// parse both .aws/credentials and .aws/config for profiles
	configFiles := map[string]bool{
		defaults.SharedCredentialsFilename(): false,
		defaults.SharedConfigFilename():      true,
	}
	for f, prefix := range configFiles {
		sections, err := loadIniFile(f, prefix)
		if err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		for key := range sections {
			name := strings.ToLower(key)
			if !contains(profiles, name) {
				profiles = append(profiles, name)
			}
		}
	}
	return profiles, nil
}

func (h contextCreateAWSHelper) chooseProfile(profiles []string) (string, error) {
	options := []string{}
	options = append(options, profiles...)

	selected, err := h.user.Select("Select AWS Profile", options)
	if err != nil {
		if err == terminal.InterruptErr {
			return "", errdefs.ErrCanceled
		}
		return "", err
	}
	profile := options[selected]
	return profile, nil
}

func (h contextCreateAWSHelper) getRegionSuggestion(region string, profile string) (string, error) {
	if profile == "" {
		profile = "default"
	}
	// only load ~/.aws/config
	awsConfig := defaults.SharedConfigFilename()
	configIni, err := ini.Load(awsConfig)
	if err != nil {
		if !os.IsNotExist(err) {
			return "", err
		}
		configIni = ini.Empty()
	}

	var f func(string, string) string
	f = func(r string, p string) string {
		section, err := configIni.GetSection(p)
		if err == nil {
			reg, err := section.GetKey("region")
			if err == nil {
				r = reg.Value()
			}
		}
		if r == "" {
			switch p {
			case "":
				return "us-east-1"
			case "default":
				return f(r, "")
			}
			return f(r, "default")
		}
		return r
	}

	if profile != "default" {
		profile = fmt.Sprintf("profile %s", profile)
	}
	return f(region, profile), nil
}

func (h contextCreateAWSHelper) chooseRegion(region string, profile string) (string, error) {
	suggestion, err := h.getRegionSuggestion(region, profile)
	if err != nil {
		return "", err
	}
	// promp user for region
	region, err = h.user.Input("Region", suggestion)
	if err != nil {
		return "", err
	}
	if region == "" {
		return "", fmt.Errorf("region cannot be empty")
	}
	return region, nil
}

func (h contextCreateAWSHelper) askCredentials() (string, string, error) {
	/*confirm, err := h.user.Confirm("Enter AWS credentials", false)
	if err != nil {
		return "", "", err
	}
	if !confirm {
		return "", "", nil
	}*/

	accessKeyID, err := h.user.Input("AWS Access Key ID", "")
	if err != nil {
		return "", "", err
	}
	secretAccessKey, err := h.user.Password("Enter AWS Secret Access Key")
	if err != nil {
		return "", "", err
	}
	// validate access ID and password
	if len(accessKeyID) < 3 || len(secretAccessKey) < 3 {
		return "", "", fmt.Errorf("AWS Access/Secret Access Key must have more than 3 characters")
	}
	return accessKeyID, secretAccessKey, nil
}

func contains(values []string, value string) bool {
	for _, v := range values {
		if v == value {
			return true
		}
	}
	return false
}

func loadIniFile(path string, prefix bool) (map[string]ini.Section, error) {
	profiles := map[string]ini.Section{}
	credIni, err := ini.Load(path)
	if err != nil {
		return nil, err
	}
	for _, section := range credIni.Sections() {
		if prefix && strings.HasPrefix(section.Name(), "profile ") {
			profiles[section.Name()[len("profile "):]] = *section
		} else if !prefix || section.Name() == "default" {
			profiles[section.Name()] = *section
		}
	}
	return profiles, nil
}
<|endoftext|>


<commit_before>package benchlist

import (
	"container/list"
	"math/rand"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/timer"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

// QueryBenchlist ...
type QueryBenchlist interface {
	// RegisterQuery registers a sent query and returns whether the query is subject to benchlist
	RegisterQuery(validatorID ids.ShortID, requestID uint32, msgType constants.MsgType) bool
	// RegisterResponse registers the response to a query message
	RegisterResponse(validatorID ids.ShortID, requstID uint32)
	// QueryFailed registers that a query did not receive a response within our synchrony bound
	QueryFailed(validatorID ids.ShortID, requestID uint32)
}

// If a peer consistently does not respond to queries, it will
// increase latencies on the network whenever that peer is polled.
// If we cannot terminate the poll early, then the poll will wait
// the full timeout before finalizing the poll and making progress.
// This can increase network latencies to an undesirable level.

// Therefore, a benchlist is used as a heurstic to immediately fail
// queries to nodes that are consistently not responding.

type queryBenchlist struct {
	vdrs validators.Set
	// Validator ID --> Request ID --> non-empty iff
	// there is an outstanding request to this validator
	// with the corresponding requestID
	pendingQueries map[[20]byte]map[uint32]pendingQuery
	// Map of consecutive query failures
	consecutiveFailures map[[20]byte]int

	// Maintain benchlist
	benchlistTimes map[[20]byte]time.Time
	benchlistOrder *list.List
	benchlistSet   ids.ShortSet

	threshold    int
	halfDuration time.Duration
	maxPortion   float64

	clock timer.Clock

	metrics *metrics
	ctx     *snow.Context

	lock sync.Mutex
}

type pendingQuery struct {
	registered time.Time
	msgType    constants.MsgType
}

// NewQueryBenchlist ...
func NewQueryBenchlist(validators validators.Set, ctx *snow.Context, threshold int, duration time.Duration, maxPortion float64, summaryEnabled bool) QueryBenchlist {
	metrics := &metrics{}
	metrics.Initialize(ctx, summaryEnabled)

	return &queryBenchlist{
		pendingQueries:      make(map[[20]byte]map[uint32]pendingQuery),
		consecutiveFailures: make(map[[20]byte]int),
		benchlistTimes:      make(map[[20]byte]time.Time),
		benchlistOrder:      list.New(),
		benchlistSet:        ids.ShortSet{},
		vdrs:                validators,
		threshold:           threshold,
		halfDuration:        duration / 2,
		maxPortion:          maxPortion,
		ctx:                 ctx,
		metrics:             metrics,
	}
}

// RegisterQuery attempts to register a query from [validatorID] and returns true
// if that request should be made (not subject to benchlisting)
func (b *queryBenchlist) RegisterQuery(validatorID ids.ShortID, requestID uint32, msgType constants.MsgType) bool {
	b.lock.Lock()
	defer b.lock.Unlock()

	key := validatorID.Key()
	if benched := b.benched(validatorID); benched {
		return false
	}
	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		validatorRequests = make(map[uint32]pendingQuery)
		b.pendingQueries[key] = validatorRequests
	}
	validatorRequests[requestID] = pendingQuery{registered: b.clock.Time(), msgType: msgType}

	return true
}

// RegisterResponse removes the query from pending
func (b *queryBenchlist) RegisterResponse(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	// Reset consecutive failures on success
	delete(b.consecutiveFailures, validatorID.Key())
}

// QueryFailed notes a failure and benchlists [validatorID] if necessary
func (b *queryBenchlist) QueryFailed(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	key := validatorID.Key()
	// Add a failure and benches [validatorID] if it has
	// passed the threshold
	b.consecutiveFailures[key]++
	if b.consecutiveFailures[key] >= b.threshold {
		b.bench(validatorID)
	}
}

func (b *queryBenchlist) bench(validatorID ids.ShortID) {
	if b.benchlistSet.Contains(validatorID) {
		return
	}

	key := validatorID.Key()

	// Add to benchlist times with randomized delay
	randomizedDuration := time.Duration(rand.Float64()*float64(b.halfDuration)) + b.halfDuration // #nosec G404
	b.benchlistTimes[key] = b.clock.Time().Add(randomizedDuration)
	b.benchlistOrder.PushBack(validatorID)
	b.benchlistSet.Add(validatorID)
	delete(b.consecutiveFailures, key)
	b.ctx.Log.Debug("Benching validator %s for %v after %d consecutive failed queries", validatorID, randomizedDuration, b.threshold)

	// Note: there could be a memory leak if a large number of
	// validators were added, sampled, benched, and never sampled
	// again. Due to the minimum staking amount and durations this
	// is not a realistic concern.
	b.cleanup()
}

// benched checks if [validatorID] is currently benched
// and calls cleanup if its benching period has elapsed
func (b *queryBenchlist) benched(validatorID ids.ShortID) bool {
	key := validatorID.Key()

	end, ok := b.benchlistTimes[key]
	if !ok {
		return false
	}

	if b.clock.Time().Before(end) {
		return true
	}

	// If a benched item has expired, cleanup the benchlist
	b.cleanup()
	return false
}

// cleanup ensures that we have not benched too much stake
// and removes anything from the benchlist whose time has expired
func (b *queryBenchlist) cleanup() {
	currentWeight, err := b.vdrs.SubsetWeight(b.benchlistSet)
	if err != nil {
		// Add log for this, should never happen
		b.ctx.Log.Error("Failed to calculate subset weight due to: %w. Resetting benchlist.", err)
		b.reset()
		return
	}

	benchLen := b.benchlistSet.Len()
	updatedWeight := currentWeight
	totalWeight := b.vdrs.Weight()
	maxBenchlistWeight := uint64(float64(totalWeight) * b.maxPortion)

	// Iterate over elements of the benchlist in order of expiration
	for e := b.benchlistOrder.Front(); e != nil; e = e.Next() {
		validatorID := e.Value.(ids.ShortID)
		key := validatorID.Key()
		end := b.benchlistTimes[key]
		// Remove elements with the next expiration until the next item has not
		// expired and the bench has less than the maximum weight
		// Note: this creates an edge case where benchlisting a validator
		// with a sufficient stake may clear the benchlist
		if b.clock.Time().Before(end) && currentWeight < maxBenchlistWeight {
			break
		}

		removeWeight, ok := b.vdrs.GetWeight(validatorID)
		if ok {
			newWeight, err := safemath.Sub64(currentWeight, removeWeight)
			if err != nil {
				b.ctx.Log.Error("Failed to calculate new subset weight due to: %w. Resetting benchlist.", err)
				b.reset()
				return
			}
			updatedWeight = newWeight
		}

		b.benchlistOrder.Remove(e)
		delete(b.benchlistTimes, key)
		b.benchlistSet.Remove(validatorID)
	}

	updatedBenchLen := b.benchlistSet.Len()
	b.ctx.Log.Debug("Benchlist weight: (%v/%v) -> (%v/%v). Benched Validators: %d -> %d",
		currentWeight,
		totalWeight,
		updatedWeight,
		totalWeight,
		benchLen,
		updatedBenchLen,
	)
	b.metrics.weightBenched.Set(float64(updatedWeight))
	b.metrics.numBenched.Set(float64(updatedBenchLen))
}

func (b *queryBenchlist) reset() {
	b.pendingQueries = make(map[[20]byte]map[uint32]pendingQuery)
	b.consecutiveFailures = make(map[[20]byte]int)
	b.benchlistTimes = make(map[[20]byte]time.Time)
	b.benchlistOrder.Init()
	b.benchlistSet.Clear()
	b.metrics.weightBenched.Set(0)
	b.metrics.numBenched.Set(0)
}

// removeQuery returns true if the query was present
func (b *queryBenchlist) removeQuery(validatorID ids.ShortID, requestID uint32) bool {
	key := validatorID.Key()

	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		return false
	}

	query, ok := validatorRequests[requestID]
	if ok {
		delete(validatorRequests, requestID)
		if len(validatorRequests) == 0 {
			delete(b.pendingQueries, key)
		}
		b.metrics.observe(validatorID, query.msgType, float64(b.clock.Time().Sub(query.registered)))
	}
	return ok
}
<commit_msg>maintain invariant that benchlist stays ordered despite randomized bench times<commit_after>package benchlist

import (
	"container/list"
	"math/rand"
	"sync"
	"time"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/timer"

	safemath "github.com/ava-labs/avalanchego/utils/math"
)

// QueryBenchlist ...
type QueryBenchlist interface {
	// RegisterQuery registers a sent query and returns whether the query is subject to benchlist
	RegisterQuery(validatorID ids.ShortID, requestID uint32, msgType constants.MsgType) bool
	// RegisterResponse registers the response to a query message
	RegisterResponse(validatorID ids.ShortID, requstID uint32)
	// QueryFailed registers that a query did not receive a response within our synchrony bound
	QueryFailed(validatorID ids.ShortID, requestID uint32)
}

// If a peer consistently does not respond to queries, it will
// increase latencies on the network whenever that peer is polled.
// If we cannot terminate the poll early, then the poll will wait
// the full timeout before finalizing the poll and making progress.
// This can increase network latencies to an undesirable level.

// Therefore, a benchlist is used as a heurstic to immediately fail
// queries to nodes that are consistently not responding.

type queryBenchlist struct {
	vdrs validators.Set
	// Validator ID --> Request ID --> non-empty iff
	// there is an outstanding request to this validator
	// with the corresponding requestID
	pendingQueries map[[20]byte]map[uint32]pendingQuery
	// Map of consecutive query failures
	consecutiveFailures map[[20]byte]int

	// Maintain benchlist
	benchlistTimes map[[20]byte]time.Time
	benchlistOrder *list.List
	benchlistSet   ids.ShortSet

	threshold  int
	duration   time.Duration
	maxPortion float64

	clock timer.Clock

	metrics *metrics
	ctx     *snow.Context

	lock sync.Mutex
}

type pendingQuery struct {
	registered time.Time
	msgType    constants.MsgType
}

// NewQueryBenchlist ...
func NewQueryBenchlist(validators validators.Set, ctx *snow.Context, threshold int, duration time.Duration, maxPortion float64, summaryEnabled bool) QueryBenchlist {
	metrics := &metrics{}
	metrics.Initialize(ctx, summaryEnabled)

	return &queryBenchlist{
		pendingQueries:      make(map[[20]byte]map[uint32]pendingQuery),
		consecutiveFailures: make(map[[20]byte]int),
		benchlistTimes:      make(map[[20]byte]time.Time),
		benchlistOrder:      list.New(),
		benchlistSet:        ids.ShortSet{},
		vdrs:                validators,
		threshold:           threshold,
		duration:            duration,
		maxPortion:          maxPortion,
		ctx:                 ctx,
		metrics:             metrics,
	}
}

// RegisterQuery attempts to register a query from [validatorID] and returns true
// if that request should be made (not subject to benchlisting)
func (b *queryBenchlist) RegisterQuery(validatorID ids.ShortID, requestID uint32, msgType constants.MsgType) bool {
	b.lock.Lock()
	defer b.lock.Unlock()

	key := validatorID.Key()
	if benched := b.benched(validatorID); benched {
		return false
	}
	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		validatorRequests = make(map[uint32]pendingQuery)
		b.pendingQueries[key] = validatorRequests
	}
	validatorRequests[requestID] = pendingQuery{registered: b.clock.Time(), msgType: msgType}

	return true
}

// RegisterResponse removes the query from pending
func (b *queryBenchlist) RegisterResponse(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	// Reset consecutive failures on success
	delete(b.consecutiveFailures, validatorID.Key())
}

// QueryFailed notes a failure and benchlists [validatorID] if necessary
func (b *queryBenchlist) QueryFailed(validatorID ids.ShortID, requestID uint32) {
	b.lock.Lock()
	defer b.lock.Unlock()

	if ok := b.removeQuery(validatorID, requestID); !ok {
		return
	}

	key := validatorID.Key()
	// Add a failure and benches [validatorID] if it has
	// passed the threshold
	b.consecutiveFailures[key]++
	if b.consecutiveFailures[key] >= b.threshold {
		b.bench(validatorID)
	}
}

func (b *queryBenchlist) bench(validatorID ids.ShortID) {
	if b.benchlistSet.Contains(validatorID) {
		return
	}

	key := validatorID.Key()

	// Goal:
	// Random end time in the range:
	// [max(lastEndTime,(currentTime + (duration/2)): currentTime + duration]
	// This maintains the invariant that validators in benchlistOrder are
	// ordered by the time that they should be unbenched
	currTime := b.clock.Time()
	minEndTime := currTime.Add(b.duration / 2)
	if elem := b.benchlistOrder.Back(); elem != nil {
		lastValidator := elem.Value.(ids.ShortID)
		lastEndTime := b.benchlistTimes[lastValidator.Key()]
		if lastEndTime.After(minEndTime) {
			minEndTime = lastEndTime
		}
	}
	maxEndTime := currTime.Add(b.duration)
	// Since maxEndTime is at least [duration] in the future and every element
	// added to benchlist was added in the past with an end time at most [duration]
	// in the future, this should never produce a negative duration.
	diff := maxEndTime.Sub(minEndTime)
	randomizedEndTime := minEndTime.Add(time.Duration(rand.Float64() * float64(diff)))

	// Add to benchlist times with randomized delay
	b.benchlistTimes[key] = randomizedEndTime
	b.benchlistOrder.PushBack(validatorID)
	b.benchlistSet.Add(validatorID)
	delete(b.consecutiveFailures, key)
	b.ctx.Log.Debug("Benching validator %s for %v after %d consecutive failed queries", validatorID, randomizedEndTime.Sub(currTime), b.threshold)

	// Note: there could be a memory leak if a large number of
	// validators were added, sampled, benched, and never sampled
	// again. Due to the minimum staking amount and durations this
	// is not a realistic concern.
	b.cleanup()
}

// benched checks if [validatorID] is currently benched
// and calls cleanup if its benching period has elapsed
func (b *queryBenchlist) benched(validatorID ids.ShortID) bool {
	key := validatorID.Key()

	end, ok := b.benchlistTimes[key]
	if !ok {
		return false
	}

	if b.clock.Time().Before(end) {
		return true
	}

	// If a benched item has expired, cleanup the benchlist
	b.cleanup()
	return false
}

// cleanup ensures that we have not benched too much stake
// and removes anything from the benchlist whose time has expired
func (b *queryBenchlist) cleanup() {
	currentWeight, err := b.vdrs.SubsetWeight(b.benchlistSet)
	if err != nil {
		// Add log for this, should never happen
		b.ctx.Log.Error("Failed to calculate subset weight due to: %w. Resetting benchlist.", err)
		b.reset()
		return
	}

	benchLen := b.benchlistSet.Len()
	updatedWeight := currentWeight
	totalWeight := b.vdrs.Weight()
	maxBenchlistWeight := uint64(float64(totalWeight) * b.maxPortion)

	// Iterate over elements of the benchlist in order of expiration
	for e := b.benchlistOrder.Front(); e != nil; e = e.Next() {
		validatorID := e.Value.(ids.ShortID)
		key := validatorID.Key()
		end := b.benchlistTimes[key]
		// Remove elements with the next expiration until the next item has not
		// expired and the bench has less than the maximum weight
		// Note: this creates an edge case where benchlisting a validator
		// with a sufficient stake may clear the benchlist
		if b.clock.Time().Before(end) && currentWeight < maxBenchlistWeight {
			break
		}

		removeWeight, ok := b.vdrs.GetWeight(validatorID)
		if ok {
			newWeight, err := safemath.Sub64(currentWeight, removeWeight)
			if err != nil {
				b.ctx.Log.Error("Failed to calculate new subset weight due to: %w. Resetting benchlist.", err)
				b.reset()
				return
			}
			updatedWeight = newWeight
		}

		b.benchlistOrder.Remove(e)
		delete(b.benchlistTimes, key)
		b.benchlistSet.Remove(validatorID)
	}

	updatedBenchLen := b.benchlistSet.Len()
	b.ctx.Log.Debug("Benchlist weight: (%v/%v) -> (%v/%v). Benched Validators: %d -> %d",
		currentWeight,
		totalWeight,
		updatedWeight,
		totalWeight,
		benchLen,
		updatedBenchLen,
	)
	b.metrics.weightBenched.Set(float64(updatedWeight))
	b.metrics.numBenched.Set(float64(updatedBenchLen))
}

func (b *queryBenchlist) reset() {
	b.pendingQueries = make(map[[20]byte]map[uint32]pendingQuery)
	b.consecutiveFailures = make(map[[20]byte]int)
	b.benchlistTimes = make(map[[20]byte]time.Time)
	b.benchlistOrder.Init()
	b.benchlistSet.Clear()
	b.metrics.weightBenched.Set(0)
	b.metrics.numBenched.Set(0)
}

// removeQuery returns true if the query was present
func (b *queryBenchlist) removeQuery(validatorID ids.ShortID, requestID uint32) bool {
	key := validatorID.Key()

	validatorRequests, ok := b.pendingQueries[key]
	if !ok {
		return false
	}

	query, ok := validatorRequests[requestID]
	if ok {
		delete(validatorRequests, requestID)
		if len(validatorRequests) == 0 {
			delete(b.pendingQueries, key)
		}
		b.metrics.observe(validatorID, query.msgType, float64(b.clock.Time().Sub(query.registered)))
	}
	return ok
}
<|endoftext|>


<commit_before>package editor

import (
	"fmt"
	"io"
)

func help(w io.Writer) error {
	for k, v := range quickref {
		_, err := w.Write([]byte(fmt.Sprintf("%s   %s\n", k, v)))
		if err != nil {
			return err
		}
	}
	return nil
}

// Reference: Vim's quickref.txt.
var quickref = map[string]string{
	"h": "left",
	"l": "right",
	"0": "to first character in the line",
	"^": "to first non-blank character in the line",
	"$": "to the last character in the line",
	"|": "to column N",
	"f": "to the Nth occurrence of {char} to the right",
	"F": "to the Nth occurrence of {char} to the left",

	"k": "go back history",
	"j": "go forward history",

	"-": "decrement the number at or after the cursor",
	"+": "increment the number at or after the cursor",

	"w": "N words forward",
	"W": "N blank-separated WORDs forward",
	"e": "forward to the end of the Nth word",
	"E": "forward to the end of the Nth blank-separated WORD",
	"b": "N words backward",
	"B": "N blank-separated WORDs backward",

	"/": "search forward",
	"?": "search backward",

	"n": "repeat last search",
	"N": "repeat last search, in opposite direction",

	"a":  "append text after the cursor",
	"A":  "append text at the end of the line",
	"i":  "insert text before the cursor",
	"I":  "insert text before the first non-blank in the line",
	"gI": "insert text in column 1",
}
<commit_msg>Add help for a few insert command<commit_after>package editor

import (
	"fmt"
	"io"
)

func help(w io.Writer) error {
	for k, v := range quickref {
		_, err := w.Write([]byte(fmt.Sprintf("%s   %s\n", k, v)))
		if err != nil {
			return err
		}
	}
	return nil
}

// Reference: Vim's quickref.txt.
var quickref = map[string]string{
	"h": "left",
	"l": "right",
	"0": "to first character in the line",
	"^": "to first non-blank character in the line",
	"$": "to the last character in the line",
	"|": "to column N",
	"f": "to the Nth occurrence of {char} to the right",
	"F": "to the Nth occurrence of {char} to the left",

	"k": "go back history",
	"j": "go forward history",

	"-": "decrement the number at or after the cursor",
	"+": "increment the number at or after the cursor",

	"w": "N words forward",
	"W": "N blank-separated WORDs forward",
	"e": "forward to the end of the Nth word",
	"E": "forward to the end of the Nth blank-separated WORD",
	"b": "N words backward",
	"B": "N blank-separated WORDs backward",

	"/": "search forward",
	"?": "search backward",

	"n": "repeat last search",
	"N": "repeat last search, in opposite direction",

	"a":  "append text after the cursor",
	"A":  "append text at the end of the line",
	"i":  "insert text before the cursor",
	"I":  "insert text before the first non-blank in the line",
	"gI": "insert text in column 1",

	// insert mode...

	"i_<Esc>":  "end Insert mode, back to Normal mode",
	"i_CTRL-C": "like <Esc>",
}
<|endoftext|>


<commit_before>package d800

const Model = ""<commit_msg>nikon/d800: Add model name<commit_after>package d800

const Model = "NIKON D800"
<|endoftext|>


<commit_before>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
	"math"
	"strings"
)

// WorkflowID is the type of unique workflow identifiers.
type WorkflowID int64

// Workflow represents the entire life cycle of a single document.
//
// A workflow begins with the creation of a document, and drives its
// life cycle through a sequence of responses to user actions or other
// system events.
//
// The engine in `flow` is visible primarily through workflows,
// documents and their behaviour.
//
// Currently, the topology of workflows is a graph, and is determined
// by the node definitions herein.
//
// N.B. It is highly recommended, but not necessary, that workflow
// names be defined in a system of hierarchical namespaces.
type Workflow struct {
	ID         WorkflowID `json:"ID,omitempty"`     // Globally-unique identifier of this workflow
	Name       string     `json:"Name,omitempty"`   // Globally-unique name of this workflow
	DocType    DocType    `json:"DocType"`          // Document type of which this workflow defines the life cycle
	BeginState DocState   `json:"BeginState"`       // Where this flow begins
	Active     bool       `json:"Active,omitempty"` // Is this workflow enabled?
}

// ApplyEvent takes an input user action or a system event, and
// applies its document action to the given document.  This results in
// a possibly new document state.  This method also prepares a message
// that is posted to applicable mailboxes.
func (w *Workflow) ApplyEvent(otx *sql.Tx, event *DocEvent, recipients []GroupID) (DocStateID, error) {
	if !w.Active {
		return 0, errors.New("this workflow is currently disabled")
	}
	if event == nil {
		return 0, errors.New("event should be non-nil")
	}
	if len(recipients) == 0 {
		return 0, errors.New("list of recipients should have length > 0")
	}
	if event.Status == EventStatusApplied {
		return 0, errors.New("event already applied; nothing to do")
	}
	if w.DocType.ID != event.DocType {
		return 0, fmt.Errorf("document type mismatch -- workflow's document type : %d, event's document type : %d", w.DocType.ID, event.DocType)
	}

	n, err := _nodes.GetByState(w.DocType.ID, event.State)
	if err != nil {
		return 0, err
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	nstate, err := n.applyEvent(tx, event, recipients)
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return nstate, nil
}

// Unexported type, only for convenience methods.
type _Workflows struct{}

var _workflows *_Workflows

func init() {
	_workflows = &_Workflows{}
}

// Workflows provides a resource-like interface to the workflows
// defined in this system.
func Workflows() *_Workflows {
	return _workflows
}

// New creates and initialises a workflow definition using the given
// name, the document type whose life cycle this workflow should
// manage, and the initial document state in which this workflow
// begins.
//
// N.B.  Workflow names must be globally-unique.
func (ws *_Workflows) New(otx *sql.Tx, name string, dtype DocTypeID, state DocStateID) (WorkflowID, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return 0, errors.New("name should not be empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_workflows(name, doctype_id, docstate_id, active)
	VALUES(?, ?, ?, 1)
	`
	res, err := tx.Exec(q, name, dtype, state)
	if err != nil {
		return 0, err
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return WorkflowID(id), nil
}

// List answers a subset of the workflows defined in the system,
// according to the given specification.
//
// Result set begins with ID >= `offset`, and has not more than
// `limit` elements.  A value of `0` for `offset` fetches from the
// beginning, while a value of `0` for `limit` fetches until the end.
func (ws *_Workflows) List(offset, limit int64) ([]*Workflow, error) {
	if offset < 0 || limit < 0 {
		return nil, errors.New("offset and limit must be non-negative integers")
	}
	if limit == 0 {
		limit = math.MaxInt64
	}

	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	ORDER BY wf.id
	LIMIT ? OFFSET ?
	`
	rows, err := db.Query(q, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ary := make([]*Workflow, 0, 10)
	for rows.Next() {
		var elem Workflow
		err = rows.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
			&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
		if err != nil {
			return nil, err
		}
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves the details of the requested workflow from the
// database.
//
// N.B.  This method retrieves the primary information of the
// workflow.  Information of the nodes comprising this workflow have
// to be fetched separately.
func (ws *_Workflows) Get(id WorkflowID) (*Workflow, error) {
	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	WHERE wf.id = ?
	`
	row := db.QueryRow(q, id)
	var elem Workflow
	err := row.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
		&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
	if err != nil {
		return nil, err
	}

	return &elem, nil
}

// Get retrieves the details of the requested workflow from the
// database.
//
// N.B.  This method retrieves the primary information of the
// workflow.  Information of the nodes comprising this workflow have
// to be fetched separately.
func (ws *_Workflows) GetByName(name string) (*Workflow, error) {
	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	WHERE wf.name = ?
	`
	row := db.QueryRow(q, name)
	var elem Workflow
	err := row.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
		&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
	if err != nil {
		return nil, err
	}

	return &elem, nil
}

// Rename assigns a new name to the given workflow.
func (ws *_Workflows) Rename(otx *sql.Tx, id WorkflowID, name string) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return errors.New("name should be non-empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	UPDATE wf_workflows SET name = ?
	WHERE id = ?
	`
	_, err := tx.Exec(q, name, id)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}

// SetActive sets the status of the workflow as either active or
// inactive, helping in workflow management and deprecation.
func (ws *_Workflows) SetActive(otx *sql.Tx, id WorkflowID, active bool) error {
	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	var flag int
	if active {
		flag = 1
	}
	q := `
	UPDATE wf_workflows SET active = ?
	WHERE id = ?
	`
	_, err := tx.Exec(q, flag, id)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}

// AddNode maps the given document state to the specified node.  This
// map is consulted by the workflow when performing a state transition
// of the system.
func (ws *_Workflows) AddNode(otx *sql.Tx, dtype DocTypeID, state DocStateID, wid WorkflowID,
	name string, ntype NodeType) (NodeID, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return 0, errors.New("name should not be empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_workflow_nodes(doctype_id, docstate_id, workflow_id, name, type)
	VALUES(?, ?, ?, ?, ?)
	`
	res, err := tx.Exec(q, dtype, state, wid, name, ntype)
	if err != nil {
		return 0, err
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return NodeID(id), nil
}

// RemoveNode unmaps the given document state to the specified node.
// This map is consulted by the workflow when performing a state
// transition of the system.
func (ws *_Workflows) RemoveNode(otx *sql.Tx, wid WorkflowID, nid NodeID) error {
	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	DELETE FROM wf_workflow_nodes
	WHERE workflow_id = ?
	AND id = ?
	`
	_, err := tx.Exec(q, wid, nid)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Correct type error when feeding node type to SQL when creating a node<commit_after>// (c) Copyright 2015-2017 JONNALAGADDA Srinivas
//
// 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 flow

import (
	"database/sql"
	"errors"
	"fmt"
	"math"
	"strings"
)

// WorkflowID is the type of unique workflow identifiers.
type WorkflowID int64

// Workflow represents the entire life cycle of a single document.
//
// A workflow begins with the creation of a document, and drives its
// life cycle through a sequence of responses to user actions or other
// system events.
//
// The engine in `flow` is visible primarily through workflows,
// documents and their behaviour.
//
// Currently, the topology of workflows is a graph, and is determined
// by the node definitions herein.
//
// N.B. It is highly recommended, but not necessary, that workflow
// names be defined in a system of hierarchical namespaces.
type Workflow struct {
	ID         WorkflowID `json:"ID,omitempty"`     // Globally-unique identifier of this workflow
	Name       string     `json:"Name,omitempty"`   // Globally-unique name of this workflow
	DocType    DocType    `json:"DocType"`          // Document type of which this workflow defines the life cycle
	BeginState DocState   `json:"BeginState"`       // Where this flow begins
	Active     bool       `json:"Active,omitempty"` // Is this workflow enabled?
}

// ApplyEvent takes an input user action or a system event, and
// applies its document action to the given document.  This results in
// a possibly new document state.  This method also prepares a message
// that is posted to applicable mailboxes.
func (w *Workflow) ApplyEvent(otx *sql.Tx, event *DocEvent, recipients []GroupID) (DocStateID, error) {
	if !w.Active {
		return 0, errors.New("this workflow is currently disabled")
	}
	if event == nil {
		return 0, errors.New("event should be non-nil")
	}
	if len(recipients) == 0 {
		return 0, errors.New("list of recipients should have length > 0")
	}
	if event.Status == EventStatusApplied {
		return 0, errors.New("event already applied; nothing to do")
	}
	if w.DocType.ID != event.DocType {
		return 0, fmt.Errorf("document type mismatch -- workflow's document type : %d, event's document type : %d", w.DocType.ID, event.DocType)
	}

	n, err := _nodes.GetByState(w.DocType.ID, event.State)
	if err != nil {
		return 0, err
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	nstate, err := n.applyEvent(tx, event, recipients)
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return nstate, nil
}

// Unexported type, only for convenience methods.
type _Workflows struct{}

var _workflows *_Workflows

func init() {
	_workflows = &_Workflows{}
}

// Workflows provides a resource-like interface to the workflows
// defined in this system.
func Workflows() *_Workflows {
	return _workflows
}

// New creates and initialises a workflow definition using the given
// name, the document type whose life cycle this workflow should
// manage, and the initial document state in which this workflow
// begins.
//
// N.B.  Workflow names must be globally-unique.
func (ws *_Workflows) New(otx *sql.Tx, name string, dtype DocTypeID, state DocStateID) (WorkflowID, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return 0, errors.New("name should not be empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_workflows(name, doctype_id, docstate_id, active)
	VALUES(?, ?, ?, 1)
	`
	res, err := tx.Exec(q, name, dtype, state)
	if err != nil {
		return 0, err
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return WorkflowID(id), nil
}

// List answers a subset of the workflows defined in the system,
// according to the given specification.
//
// Result set begins with ID >= `offset`, and has not more than
// `limit` elements.  A value of `0` for `offset` fetches from the
// beginning, while a value of `0` for `limit` fetches until the end.
func (ws *_Workflows) List(offset, limit int64) ([]*Workflow, error) {
	if offset < 0 || limit < 0 {
		return nil, errors.New("offset and limit must be non-negative integers")
	}
	if limit == 0 {
		limit = math.MaxInt64
	}

	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	ORDER BY wf.id
	LIMIT ? OFFSET ?
	`
	rows, err := db.Query(q, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ary := make([]*Workflow, 0, 10)
	for rows.Next() {
		var elem Workflow
		err = rows.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
			&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
		if err != nil {
			return nil, err
		}
		ary = append(ary, &elem)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}

	return ary, nil
}

// Get retrieves the details of the requested workflow from the
// database.
//
// N.B.  This method retrieves the primary information of the
// workflow.  Information of the nodes comprising this workflow have
// to be fetched separately.
func (ws *_Workflows) Get(id WorkflowID) (*Workflow, error) {
	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	WHERE wf.id = ?
	`
	row := db.QueryRow(q, id)
	var elem Workflow
	err := row.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
		&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
	if err != nil {
		return nil, err
	}

	return &elem, nil
}

// Get retrieves the details of the requested workflow from the
// database.
//
// N.B.  This method retrieves the primary information of the
// workflow.  Information of the nodes comprising this workflow have
// to be fetched separately.
func (ws *_Workflows) GetByName(name string) (*Workflow, error) {
	q := `
	SELECT wf.id, wf.name, dtm.id, dtm.name, dsm.id, dsm.name, wf.active
	FROM wf_workflows wf
	JOIN wf_doctypes_master dtm ON wf.doctype_id = dtm.id
	JOIN wf_docstates_master dsm ON wf.docstate_id = dsm.id
	WHERE wf.name = ?
	`
	row := db.QueryRow(q, name)
	var elem Workflow
	err := row.Scan(&elem.ID, &elem.Name, &elem.DocType.ID, &elem.DocType.Name,
		&elem.BeginState.ID, &elem.BeginState.Name, &elem.Active)
	if err != nil {
		return nil, err
	}

	return &elem, nil
}

// Rename assigns a new name to the given workflow.
func (ws *_Workflows) Rename(otx *sql.Tx, id WorkflowID, name string) error {
	name = strings.TrimSpace(name)
	if name == "" {
		return errors.New("name should be non-empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	UPDATE wf_workflows SET name = ?
	WHERE id = ?
	`
	_, err := tx.Exec(q, name, id)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}

// SetActive sets the status of the workflow as either active or
// inactive, helping in workflow management and deprecation.
func (ws *_Workflows) SetActive(otx *sql.Tx, id WorkflowID, active bool) error {
	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	var flag int
	if active {
		flag = 1
	}
	q := `
	UPDATE wf_workflows SET active = ?
	WHERE id = ?
	`
	_, err := tx.Exec(q, flag, id)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}

// AddNode maps the given document state to the specified node.  This
// map is consulted by the workflow when performing a state transition
// of the system.
func (ws *_Workflows) AddNode(otx *sql.Tx, dtype DocTypeID, state DocStateID, wid WorkflowID,
	name string, ntype NodeType) (NodeID, error) {
	name = strings.TrimSpace(name)
	if name == "" {
		return 0, errors.New("name should not be empty")
	}

	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return 0, err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	INSERT INTO wf_workflow_nodes(doctype_id, docstate_id, workflow_id, name, type)
	VALUES(?, ?, ?, ?, ?)
	`
	res, err := tx.Exec(q, dtype, state, wid, name, string(ntype))
	if err != nil {
		return 0, err
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return 0, err
		}
	}

	return NodeID(id), nil
}

// RemoveNode unmaps the given document state to the specified node.
// This map is consulted by the workflow when performing a state
// transition of the system.
func (ws *_Workflows) RemoveNode(otx *sql.Tx, wid WorkflowID, nid NodeID) error {
	var tx *sql.Tx
	if otx == nil {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		defer tx.Rollback()
	} else {
		tx = otx
	}

	q := `
	DELETE FROM wf_workflow_nodes
	WHERE workflow_id = ?
	AND id = ?
	`
	_, err := tx.Exec(q, wid, nid)
	if err != nil {
		return err
	}

	if otx == nil {
		err = tx.Commit()
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package alerting

import (
	"fmt"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/golang/glog"
	"github.com/influxdb/influxdb/client"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
	"skia.googlesource.com/buildbot.git/go/email"
	"skia.googlesource.com/buildbot.git/go/util"
)

const (
	NAG_MSG_TMPL = "This alert has been active for %s since the last update. Please verify that it is still valid and either fix the issue or dismiss/snooze the alert."
)

type queryable interface {
	Query(string, ...client.TimePrecision) ([]*client.Series, error)
}

func executeQuery(c queryable, q string) (float64, error) {
	results, err := c.Query(q)
	if err != nil {
		return 0.0, fmt.Errorf("Failed to query InfluxDB with query %q: %s", q, err)
	}
	if len(results) < 1 {
		return 0.0, fmt.Errorf("Query returned no data: %q", q)
	}
	points := results[0].Points
	if len(points) < 1 {
		return 0.0, fmt.Errorf("Query returned no points: %q", q)
	}
	if len(points) > 1 {
		return 0.0, fmt.Errorf("Query returned more than one point: %q", q)
	}
	valueColumn := 0
	for _, label := range results[0].Columns {
		if label == "time" || label == "sequence_number" {
			valueColumn++
		} else {
			break
		}
	}
	if len(results[0].Columns) != valueColumn+1 {
		return 0.0, fmt.Errorf("Query returned an incorrect set of columns: %q %v", q, results[0].Columns)
	}
	if len(results[0].Columns) != len(points[0]) {
		return 0.0, fmt.Errorf("Invalid data from InfluxDB: Point data does not match column spec.")
	}
	return points[0][valueColumn].(float64), nil
}

// Comment is an object representing a comment on an alert.
type Comment struct {
	Time    time.Time
	User    string
	Message string
}

// Alert is an object which represents an active alert.
type Alert struct {
	Id            int64
	lastTriggered time.Time
	snoozedUntil  time.Time
	dismissedAt   time.Time
	Comments      []*Comment
	Rule          *Rule
}

// addComment adds a comment about the alert.
func (a *Alert) addComment(c *Comment) {
	a.Comments = append(a.Comments, c)
	for _, action := range a.Rule.actions {
		go action.Followup(c.Message)
	}
}

// Snoozed indicates whether the Alert has been Snoozed.
func (a *Alert) Snoozed() bool {
	return a.snoozedUntil != time.Time{}
}

// Triggered gives the time when the alert was triggered.
func (a Alert) Triggered() time.Time {
	return a.lastTriggered
}

// SnoozedUntil gives the time until which the alert is snoozed.
func (a Alert) SnoozedUntil() time.Time {
	return a.snoozedUntil
}

// Rule is an object used for triggering Alerts.
type Rule struct {
	Id          string        `json:"id"`
	Name        string        `json:"name"`
	Query       string        `json:"query"`
	Condition   string        `json:"condition"`
	Message     string        `json:"message"`
	Nag         time.Duration `json:"nag"`
	client      queryable
	AutoDismiss bool `json:"autoDismiss"`
	actions     []Action
	activeAlert *Alert
}

// Fire causes the Alert to become Active() and not Snoozed(), and causes each
// action to be performed. Active Alerts do not perform new queries.
func (r *Rule) fire() {
	a := Alert{
		Id:            0,
		lastTriggered: time.Now().UTC(),
		Comments:      []*Comment{},
		Rule:          r,
	}
	r.activeAlert = &a
	for _, action := range r.actions {
		go action.Fire()
	}
}

func (r *Rule) tick() {
	if r.activeAlert != nil && r.activeAlert.Snoozed() {
		if r.activeAlert.snoozedUntil.Before(time.Now().UTC()) {
			r.activeAlert.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: "Dismissing; snooze period expired.",
			})
			r.activeAlert = nil
		}
	} else if r.AutoDismiss || r.activeAlert == nil {
		glog.Infof("Executing query [%s]", r.Query)
		d, err := executeQuery(r.client, r.Query)
		if err != nil {
			glog.Error(err)
			return
		}
		glog.Infof("Query [%s] returned %v", r.Query, d)
		doAlert, err := r.evaluate(d)
		if err != nil {
			glog.Error(err)
			return
		}
		if r.activeAlert != nil && r.AutoDismiss && !doAlert {
			r.activeAlert.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: "Auto-dismissing; condition is no longer true.",
			})
			r.activeAlert = nil
		} else if r.activeAlert == nil && doAlert {
			r.fire()
		}
	}
	r.maybeNag()
}

func (r *Rule) maybeNag() {
	a := r.activeAlert
	if a != nil && !a.Snoozed() && r.Nag != time.Duration(0) {
		lastMsgTime := a.Triggered()
		if len(a.Comments) > 0 {
			lastMsgTime = a.Comments[len(a.Comments)-1].Time
		}
		if time.Since(lastMsgTime) > r.Nag {
			a.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: fmt.Sprintf(NAG_MSG_TMPL, r.Nag.String()),
			})
		}
	}
}

func (r *Rule) evaluate(d float64) (bool, error) {
	pkg := types.NewPackage("evaluateme", "evaluateme")
	v := exact.MakeFloat64(d)
	pkg.Scope().Insert(types.NewConst(0, pkg, "x", types.Typ[types.Float64], v))
	typ, val, err := types.Eval(r.Condition, pkg, pkg.Scope())
	if err != nil {
		return false, fmt.Errorf("Failed to evaluate condition %q: %s", r.Condition, err)
	}
	if typ.String() != "untyped bool" {
		return false, fmt.Errorf("Rule \"%v\" does not return boolean type.", r.Condition)
	}
	return exact.BoolVal(val), nil
}

type parsedRule map[string]interface{}

func newRule(r parsedRule, client *client.Client, emailAuth *email.GMail, testing bool) (*Rule, error) {
	errString := "Alert rule missing field %q"
	name, ok := r["name"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "name")
	}
	query, ok := r["query"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "query")
	}
	condition, ok := r["condition"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "condition")
	}
	message, ok := r["message"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "message")
	}
	AutoDismiss, ok := r["auto-dismiss"].(bool)
	if !ok {
		return nil, fmt.Errorf(errString, "auto-dismiss")
	}
	actionsInterface, ok := r["actions"]
	if !ok {
		return nil, fmt.Errorf(errString, "actions")
	}
	nagDuration := time.Duration(0)
	nag, ok := r["nag"].(string)
	if ok {
		var err error
		nagDuration, err = time.ParseDuration(nag)
		if err != nil {
			return nil, fmt.Errorf("Invalid nag duration %q: %v", nag, err)
		}
	}
	id, err := util.GenerateID()
	if err != nil {
		return nil, err
	}
	rule := Rule{
		Id:          id,
		Name:        name,
		Query:       query,
		Condition:   condition,
		Message:     message,
		Nag:         nagDuration,
		client:      client,
		AutoDismiss: AutoDismiss,
		actions:     nil,
		activeAlert: nil,
	}
	if err := rule.parseActions(actionsInterface, emailAuth, testing); err != nil {
		return nil, err
	}
	// Verify that the condition can be evaluated.
	_, err = rule.evaluate(0.0)
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func parseAlertRules(cfgFile string) ([]parsedRule, error) {
	var cfg struct {
		Rule []parsedRule
	}
	_, err := toml.DecodeFile(cfgFile, &cfg)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse %s: %s", cfgFile, err)
	}
	return cfg.Rule, nil
}

func makeRules(cfgFile string, dbClient *client.Client, emailAuth *email.GMail, testing bool) ([]*Rule, error) {
	parsedRules, err := parseAlertRules(cfgFile)
	if err != nil {
		return nil, err
	}
	rules := []*Rule{}
	for _, r := range parsedRules {
		r, err := newRule(r, dbClient, emailAuth, testing)
		if err != nil {
			return nil, err
		}
		rules = append(rules, r)
	}
	return rules, nil
}
<commit_msg>AlertServer: display user name as part of comment emails<commit_after>package alerting

import (
	"fmt"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/golang/glog"
	"github.com/influxdb/influxdb/client"
	"golang.org/x/tools/go/exact"
	"golang.org/x/tools/go/types"
	"skia.googlesource.com/buildbot.git/go/email"
	"skia.googlesource.com/buildbot.git/go/util"
)

const (
	NAG_MSG_TMPL = "This alert has been active for %s since the last update. Please verify that it is still valid and either fix the issue or dismiss/snooze the alert."
)

type queryable interface {
	Query(string, ...client.TimePrecision) ([]*client.Series, error)
}

func executeQuery(c queryable, q string) (float64, error) {
	results, err := c.Query(q)
	if err != nil {
		return 0.0, fmt.Errorf("Failed to query InfluxDB with query %q: %s", q, err)
	}
	if len(results) < 1 {
		return 0.0, fmt.Errorf("Query returned no data: %q", q)
	}
	points := results[0].Points
	if len(points) < 1 {
		return 0.0, fmt.Errorf("Query returned no points: %q", q)
	}
	if len(points) > 1 {
		return 0.0, fmt.Errorf("Query returned more than one point: %q", q)
	}
	valueColumn := 0
	for _, label := range results[0].Columns {
		if label == "time" || label == "sequence_number" {
			valueColumn++
		} else {
			break
		}
	}
	if len(results[0].Columns) != valueColumn+1 {
		return 0.0, fmt.Errorf("Query returned an incorrect set of columns: %q %v", q, results[0].Columns)
	}
	if len(results[0].Columns) != len(points[0]) {
		return 0.0, fmt.Errorf("Invalid data from InfluxDB: Point data does not match column spec.")
	}
	return points[0][valueColumn].(float64), nil
}

// Comment is an object representing a comment on an alert.
type Comment struct {
	Time    time.Time
	User    string
	Message string
}

// Alert is an object which represents an active alert.
type Alert struct {
	Id            int64
	lastTriggered time.Time
	snoozedUntil  time.Time
	dismissedAt   time.Time
	Comments      []*Comment
	Rule          *Rule
}

// addComment adds a comment about the alert.
func (a *Alert) addComment(c *Comment) {
	a.Comments = append(a.Comments, c)
	for _, action := range a.Rule.actions {
		go action.Followup(fmt.Sprintf("%s: %s", c.User, c.Message))
	}
}

// Snoozed indicates whether the Alert has been Snoozed.
func (a *Alert) Snoozed() bool {
	return a.snoozedUntil != time.Time{}
}

// Triggered gives the time when the alert was triggered.
func (a Alert) Triggered() time.Time {
	return a.lastTriggered
}

// SnoozedUntil gives the time until which the alert is snoozed.
func (a Alert) SnoozedUntil() time.Time {
	return a.snoozedUntil
}

// Rule is an object used for triggering Alerts.
type Rule struct {
	Id          string        `json:"id"`
	Name        string        `json:"name"`
	Query       string        `json:"query"`
	Condition   string        `json:"condition"`
	Message     string        `json:"message"`
	Nag         time.Duration `json:"nag"`
	client      queryable
	AutoDismiss bool `json:"autoDismiss"`
	actions     []Action
	activeAlert *Alert
}

// Fire causes the Alert to become Active() and not Snoozed(), and causes each
// action to be performed. Active Alerts do not perform new queries.
func (r *Rule) fire() {
	a := Alert{
		Id:            0,
		lastTriggered: time.Now().UTC(),
		Comments:      []*Comment{},
		Rule:          r,
	}
	r.activeAlert = &a
	for _, action := range r.actions {
		go action.Fire()
	}
}

func (r *Rule) tick() {
	if r.activeAlert != nil && r.activeAlert.Snoozed() {
		if r.activeAlert.snoozedUntil.Before(time.Now().UTC()) {
			r.activeAlert.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: "Dismissing; snooze period expired.",
			})
			r.activeAlert = nil
		}
	} else if r.AutoDismiss || r.activeAlert == nil {
		glog.Infof("Executing query [%s]", r.Query)
		d, err := executeQuery(r.client, r.Query)
		if err != nil {
			glog.Error(err)
			return
		}
		glog.Infof("Query [%s] returned %v", r.Query, d)
		doAlert, err := r.evaluate(d)
		if err != nil {
			glog.Error(err)
			return
		}
		if r.activeAlert != nil && r.AutoDismiss && !doAlert {
			r.activeAlert.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: "Auto-dismissing; condition is no longer true.",
			})
			r.activeAlert = nil
		} else if r.activeAlert == nil && doAlert {
			r.fire()
		}
	}
	r.maybeNag()
}

func (r *Rule) maybeNag() {
	a := r.activeAlert
	if a != nil && !a.Snoozed() && r.Nag != time.Duration(0) {
		lastMsgTime := a.Triggered()
		if len(a.Comments) > 0 {
			lastMsgTime = a.Comments[len(a.Comments)-1].Time
		}
		if time.Since(lastMsgTime) > r.Nag {
			a.addComment(&Comment{
				Time:    time.Now().UTC(),
				User:    "AlertServer",
				Message: fmt.Sprintf(NAG_MSG_TMPL, r.Nag.String()),
			})
		}
	}
}

func (r *Rule) evaluate(d float64) (bool, error) {
	pkg := types.NewPackage("evaluateme", "evaluateme")
	v := exact.MakeFloat64(d)
	pkg.Scope().Insert(types.NewConst(0, pkg, "x", types.Typ[types.Float64], v))
	typ, val, err := types.Eval(r.Condition, pkg, pkg.Scope())
	if err != nil {
		return false, fmt.Errorf("Failed to evaluate condition %q: %s", r.Condition, err)
	}
	if typ.String() != "untyped bool" {
		return false, fmt.Errorf("Rule \"%v\" does not return boolean type.", r.Condition)
	}
	return exact.BoolVal(val), nil
}

type parsedRule map[string]interface{}

func newRule(r parsedRule, client *client.Client, emailAuth *email.GMail, testing bool) (*Rule, error) {
	errString := "Alert rule missing field %q"
	name, ok := r["name"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "name")
	}
	query, ok := r["query"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "query")
	}
	condition, ok := r["condition"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "condition")
	}
	message, ok := r["message"].(string)
	if !ok {
		return nil, fmt.Errorf(errString, "message")
	}
	AutoDismiss, ok := r["auto-dismiss"].(bool)
	if !ok {
		return nil, fmt.Errorf(errString, "auto-dismiss")
	}
	actionsInterface, ok := r["actions"]
	if !ok {
		return nil, fmt.Errorf(errString, "actions")
	}
	nagDuration := time.Duration(0)
	nag, ok := r["nag"].(string)
	if ok {
		var err error
		nagDuration, err = time.ParseDuration(nag)
		if err != nil {
			return nil, fmt.Errorf("Invalid nag duration %q: %v", nag, err)
		}
	}
	id, err := util.GenerateID()
	if err != nil {
		return nil, err
	}
	rule := Rule{
		Id:          id,
		Name:        name,
		Query:       query,
		Condition:   condition,
		Message:     message,
		Nag:         nagDuration,
		client:      client,
		AutoDismiss: AutoDismiss,
		actions:     nil,
		activeAlert: nil,
	}
	if err := rule.parseActions(actionsInterface, emailAuth, testing); err != nil {
		return nil, err
	}
	// Verify that the condition can be evaluated.
	_, err = rule.evaluate(0.0)
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func parseAlertRules(cfgFile string) ([]parsedRule, error) {
	var cfg struct {
		Rule []parsedRule
	}
	_, err := toml.DecodeFile(cfgFile, &cfg)
	if err != nil {
		return nil, fmt.Errorf("Failed to parse %s: %s", cfgFile, err)
	}
	return cfg.Rule, nil
}

func makeRules(cfgFile string, dbClient *client.Client, emailAuth *email.GMail, testing bool) ([]*Rule, error) {
	parsedRules, err := parseAlertRules(cfgFile)
	if err != nil {
		return nil, err
	}
	rules := []*Rule{}
	for _, r := range parsedRules {
		r, err := newRule(r, dbClient, emailAuth, testing)
		if err != nil {
			return nil, err
		}
		rules = append(rules, r)
	}
	return rules, nil
}
<|endoftext|>


<commit_before>package lora

import (
	"bytes"
	"crypto/aes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/jacobsa/crypto/cmac"
	"log"
)

type PHYPayload struct {
	MHDR       byte
	MACPayload []byte
	MIC        []byte
	FCtrl      byte
	FCnt       uint16
	FOpts      []byte
	FPort      byte
	DevAddr    uint32
}

func ParsePHYPayload(buf []byte) (*PHYPayload, error) {
	data := &PHYPayload{
		MHDR: buf[0],
	}

	majorVersion := data.MHDR & 0x3
	if majorVersion != 0 {
		return nil, errors.New(fmt.Sprintf("Major version %d not supported", majorVersion))
	}

	if len(buf) < 5 {
		return nil, errors.New("The data is should be at least 5 bytes")
	}

	data.MACPayload = buf[1 : len(buf)-4]
	data.MIC = buf[len(buf)-4 : len(buf)]

	if len(data.MACPayload) < 7 {
		return nil, errors.New("Payload should at least be 7 bytes")
	}

	binary.Read(bytes.NewReader(data.MACPayload[0:4]), binary.LittleEndian, &data.DevAddr)
	data.FCtrl = data.MACPayload[4]
	binary.Read(bytes.NewReader(data.MACPayload[5:7]), binary.LittleEndian, &data.FCnt)

	index := 7
	fOptsLen := int(data.FCtrl & 0xf)
	if fOptsLen > 0 {
		if len(data.MACPayload) < index+fOptsLen {
			return nil, errors.New("Payload does not contain indicated options length")
		}
		data.FOpts = data.MACPayload[index : index+fOptsLen]
		index += fOptsLen
	} else {
		data.FOpts = make([]byte, 0)
	}

	if len(data.MACPayload) > index+1 {
		data.FPort = data.MACPayload[index]
	}

	return data, nil
}

func (d *PHYPayload) DecryptPayload(key []byte) ([]byte, error) {
	if len(d.MACPayload) <= 8+len(d.FOpts) {
		return nil, errors.New("No data to decrypt")
	}
	data := d.MACPayload[8+len(d.FOpts):]

	// See LoRaWAN specification 1r0 4.3.3.1
	block, err := aes.NewCipher(key)
	if err != nil {
		log.Printf("Failed to create AES cipher: %s", err.Error())
		return nil, err
	}

	buf := bytes.NewBuffer(data)
	i := 1
	group := buf.Next(block.BlockSize())
	result := make([]byte, 0, len(data))
	for len(group) > 0 {
		a := new(bytes.Buffer)
		a.Write([]byte{0x1, 0x0, 0x0, 0x0, 0x0, 0x0})
		binary.Write(a, binary.LittleEndian, d.DevAddr)
		binary.Write(a, binary.LittleEndian, uint32(d.FCnt))
		a.WriteByte(0x0)
		a.WriteByte(byte(i))

		key := make([]byte, block.BlockSize())
		block.Encrypt(key, a.Bytes())

		for j := 0; j < len(group); j++ {
			result = append(result, group[j]^key[j])
		}

		i++
		group = buf.Next(block.BlockSize())
	}
	return result, nil
}

func (d *PHYPayload) TestIntegrity(key []byte) (bool, error) {
	// See LoRaWAN specification 1r0 4.4
	b0 := new(bytes.Buffer)
	b0.Write([]byte{0x49, 0x0, 0x0, 0x0, 0x0})
	b0.WriteByte(0x0)
	binary.Write(b0, binary.LittleEndian, d.DevAddr)
	binary.Write(b0, binary.LittleEndian, uint32(d.FCnt))
	b0.WriteByte(0x0)
	b0.WriteByte(byte(1 + len(d.MACPayload)))
	b0.WriteByte(d.MHDR)
	b0.Write(d.MACPayload)

	hash, err := cmac.New(key)
	if err != nil {
		log.Printf("Failed to initialize CMAC: %s", err.Error())
		return false, err
	}

	_, err = hash.Write(b0.Bytes())
	if err != nil {
		log.Printf("Failed to hash data: %s", err.Error())
		return false, err
	}

	calculatedMIC := hash.Sum([]byte{})[0:4]
	return bytes.Equal(calculatedMIC, d.MIC), nil
}
<commit_msg>Major version check and dump of data<commit_after>package lora

import (
	"bytes"
	"crypto/aes"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/jacobsa/crypto/cmac"
	"log"
)

type PHYPayload struct {
	MHDR       byte
	MACPayload []byte
	MIC        []byte
	FCtrl      byte
	FCnt       uint16
	FOpts      []byte
	FPort      byte
	DevAddr    uint32
}

func ParsePHYPayload(buf []byte) (*PHYPayload, error) {
	data := &PHYPayload{
		MHDR: buf[0],
	}

	mType := (data.MHDR & 0xe0) >> 5
	log.Printf("The message type is %d", mType)

	majorVersion := data.MHDR & 0x2
	if majorVersion != 0 {
		return nil, errors.New(fmt.Sprintf("Major version %d not supported\nData:%#v", majorVersion, buf[1:]))
	}

	if len(buf) < 5 {
		return nil, errors.New("The data is should be at least 5 bytes")
	}

	data.MACPayload = buf[1 : len(buf)-4]
	data.MIC = buf[len(buf)-4 : len(buf)]

	if len(data.MACPayload) < 7 {
		return nil, errors.New("Payload should at least be 7 bytes")
	}

	binary.Read(bytes.NewReader(data.MACPayload[0:4]), binary.LittleEndian, &data.DevAddr)
	data.FCtrl = data.MACPayload[4]
	binary.Read(bytes.NewReader(data.MACPayload[5:7]), binary.LittleEndian, &data.FCnt)

	index := 7
	fOptsLen := int(data.FCtrl & 0xf)
	if fOptsLen > 0 {
		if len(data.MACPayload) < index+fOptsLen {
			return nil, errors.New("Payload does not contain indicated options length")
		}
		data.FOpts = data.MACPayload[index : index+fOptsLen]
		index += fOptsLen
	} else {
		data.FOpts = make([]byte, 0)
	}

	if len(data.MACPayload) > index+1 {
		data.FPort = data.MACPayload[index]
	}

	return data, nil
}

func (d *PHYPayload) DecryptPayload(key []byte) ([]byte, error) {
	if len(d.MACPayload) <= 8+len(d.FOpts) {
		return nil, errors.New("No data to decrypt")
	}
	data := d.MACPayload[8+len(d.FOpts):]

	// See LoRaWAN specification 1r0 4.3.3.1
	block, err := aes.NewCipher(key)
	if err != nil {
		log.Printf("Failed to create AES cipher: %s", err.Error())
		return nil, err
	}

	buf := bytes.NewBuffer(data)
	i := 1
	group := buf.Next(block.BlockSize())
	result := make([]byte, 0, len(data))
	for len(group) > 0 {
		a := new(bytes.Buffer)
		a.Write([]byte{0x1, 0x0, 0x0, 0x0, 0x0, 0x0})
		binary.Write(a, binary.LittleEndian, d.DevAddr)
		binary.Write(a, binary.LittleEndian, uint32(d.FCnt))
		a.WriteByte(0x0)
		a.WriteByte(byte(i))

		key := make([]byte, block.BlockSize())
		block.Encrypt(key, a.Bytes())

		for j := 0; j < len(group); j++ {
			result = append(result, group[j]^key[j])
		}

		i++
		group = buf.Next(block.BlockSize())
	}
	return result, nil
}

func (d *PHYPayload) TestIntegrity(key []byte) (bool, error) {
	// See LoRaWAN specification 1r0 4.4
	b0 := new(bytes.Buffer)
	b0.Write([]byte{0x49, 0x0, 0x0, 0x0, 0x0})
	b0.WriteByte(0x0)
	binary.Write(b0, binary.LittleEndian, d.DevAddr)
	binary.Write(b0, binary.LittleEndian, uint32(d.FCnt))
	b0.WriteByte(0x0)
	b0.WriteByte(byte(1 + len(d.MACPayload)))
	b0.WriteByte(d.MHDR)
	b0.Write(d.MACPayload)

	hash, err := cmac.New(key)
	if err != nil {
		log.Printf("Failed to initialize CMAC: %s", err.Error())
		return false, err
	}

	_, err = hash.Write(b0.Bytes())
	if err != nil {
		log.Printf("Failed to hash data: %s", err.Error())
		return false, err
	}

	calculatedMIC := hash.Sum([]byte{})[0:4]
	return bytes.Equal(calculatedMIC, d.MIC), nil
}
<|endoftext|>