<commit_before>/*
Copyright 2021 The Vitess 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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
)

func TestOrderBy(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	defer func() {
		_, _ = exec(t, conn, `delete from t1`)
	}()

	// insert some data.
	checkedExec(t, conn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)

	// Gen4 only supported query.
	assertMatches(t, conn, `select col from t1 order by id`, `[[INT64(13)] [INT64(12)] [INT64(123)] [INT64(1234)]]`)

	// Gen4 unsupported query. v3 supported.
	assertMatches(t, conn, `select col from t1 order by 1`, `[[INT64(12)] [INT64(13)] [INT64(123)] [INT64(1234)]]`)

	// unsupported in v3 and Gen4.
	_, err = exec(t, conn, `select t1.* from t1 order by id`)
	require.Error(t, err)
}

func TestGroupBy(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	defer func() {
		_, _ = exec(t, conn, `delete from t1`)
		_, _ = exec(t, conn, `delete from t2`)
	}()

	// insert some data.
	checkedExec(t, conn, `insert into t1(id, col) values (1, 123),(2, 12),(3, 13),(4, 1234)`)
	checkedExec(t, conn, `insert into t2(id, tcol1, tcol2) values (1, 'A', 'A'),(2, 'B', 'C'),(3, 'A', 'C'),(4, 'C', 'A'),(5, 'A', 'A'),(6, 'B', 'C'),(7, 'B', 'A'),(8, 'C', 'B')`)

	// Gen4 only supported query.
	assertMatches(t, conn, `select tcol2, tcol1, count(id) from t2 group by tcol2, tcol1`,
		`[[VARCHAR("A") VARCHAR("A") INT64(2)] [VARCHAR("A") VARCHAR("B") INT64(1)] [VARCHAR("A") VARCHAR("C") INT64(1)] [VARCHAR("B") VARCHAR("C") INT64(1)] [VARCHAR("C") VARCHAR("A") INT64(1)] [VARCHAR("C") VARCHAR("B") INT64(2)]]`)

	assertMatches(t, conn, `select tcol1, tcol1 from t2 order by tcol1`,
		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")] [VARCHAR("C") VARCHAR("C")]]`)

	assertMatches(t, conn, `select tcol1, tcol1 from t1 join t2 on t1.id = t2.id order by tcol1`,
		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")]]`)

	assertMatches(t, conn, `select count(*) k, tcol1, tcol2, "abc" b from t2 group by tcol1, tcol2, b order by k, tcol2, tcol1`,
		`[[INT64(1) VARCHAR("B") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("C") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("C") VARCHAR("B") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("A") VARCHAR("C") VARCHAR("abc")] `+
			`[INT64(2) VARCHAR("A") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(2) VARCHAR("B") VARCHAR("C") VARCHAR("abc")]]`)
}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := checkedExec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}

func checkedExec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := exec(t, conn, query)
	require.NoError(t, err, "for query: "+query)
	return qr
}

func exec(t *testing.T, conn *mysql.Conn, query string) (*sqltypes.Result, error) {
	t.Helper()
	return conn.ExecuteFetch(query, 1000, true)
}
<commit_msg>gen4: added e2e test for count distinct function<commit_after>/*
Copyright 2021 The Vitess 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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
)

func TestOrderBy(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	defer func() {
		_, _ = exec(t, conn, `delete from t1`)
	}()

	// insert some data.
	checkedExec(t, conn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)

	// Gen4 only supported query.
	assertMatches(t, conn, `select col from t1 order by id`, `[[INT64(13)] [INT64(12)] [INT64(123)] [INT64(1234)]]`)

	// Gen4 unsupported query. v3 supported.
	assertMatches(t, conn, `select col from t1 order by 1`, `[[INT64(12)] [INT64(13)] [INT64(123)] [INT64(1234)]]`)

	// unsupported in v3 and Gen4.
	_, err = exec(t, conn, `select t1.* from t1 order by id`)
	require.Error(t, err)
}

func TestGroupBy(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	defer func() {
		_, _ = exec(t, conn, `delete from t1`)
		_, _ = exec(t, conn, `delete from t2`)
	}()

	// insert some data.
	checkedExec(t, conn, `insert into t1(id, col) values (1, 123),(2, 12),(3, 13),(4, 1234)`)
	checkedExec(t, conn, `insert into t2(id, tcol1, tcol2) values (1, 'A', 'A'),(2, 'B', 'C'),(3, 'A', 'C'),(4, 'C', 'A'),(5, 'A', 'A'),(6, 'B', 'C'),(7, 'B', 'A'),(8, 'C', 'B')`)

	// Gen4 only supported query.
	assertMatches(t, conn, `select tcol2, tcol1, count(id) from t2 group by tcol2, tcol1`,
		`[[VARCHAR("A") VARCHAR("A") INT64(2)] [VARCHAR("A") VARCHAR("B") INT64(1)] [VARCHAR("A") VARCHAR("C") INT64(1)] [VARCHAR("B") VARCHAR("C") INT64(1)] [VARCHAR("C") VARCHAR("A") INT64(1)] [VARCHAR("C") VARCHAR("B") INT64(2)]]`)

	assertMatches(t, conn, `select tcol1, tcol1 from t2 order by tcol1`,
		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")] [VARCHAR("C") VARCHAR("C")]]`)

	assertMatches(t, conn, `select tcol1, tcol1 from t1 join t2 on t1.id = t2.id order by tcol1`,
		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")]]`)

	assertMatches(t, conn, `select count(*) k, tcol1, tcol2, "abc" b from t2 group by tcol1, tcol2, b order by k, tcol2, tcol1`,
		`[[INT64(1) VARCHAR("B") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("C") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("C") VARCHAR("B") VARCHAR("abc")] `+
			`[INT64(1) VARCHAR("A") VARCHAR("C") VARCHAR("abc")] `+
			`[INT64(2) VARCHAR("A") VARCHAR("A") VARCHAR("abc")] `+
			`[INT64(2) VARCHAR("B") VARCHAR("C") VARCHAR("abc")]]`)
}

func TestAggregationFunc(t *testing.T) {
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	defer exec(t, conn, `delete from t2`)

	// insert some data.
	checkedExec(t, conn, `insert into t2(id, tcol1, tcol2) values (1, 'A', 'A'),(2, 'B', 'C'),(3, 'A', 'C'),(4, 'C', 'A'),(5, 'A', 'A'),(6, 'B', 'C'),(7, 'B', 'A'),(8, 'C', 'A')`)

	// on primary vindex
	assertMatches(t, conn, `select tcol1, count(distinct id) from t2 group by tcol1`,
		`[[VARCHAR("A") INT64(3)] [VARCHAR("B") INT64(3)] [VARCHAR("C") INT64(2)]]`)

	// on any column
	assertMatches(t, conn, `select tcol1, count(distinct tcol2) from t2 group by tcol1`,
		`[[VARCHAR("A") INT64(2)] [VARCHAR("B") INT64(2)] [VARCHAR("C") INT64(1)]]`)

}

func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
	t.Helper()
	qr := checkedExec(t, conn, query)
	got := fmt.Sprintf("%v", qr.Rows)
	diff := cmp.Diff(expected, got)
	if diff != "" {
		t.Errorf("Query: %s (-want +got):\n%s", query, diff)
	}
}

func checkedExec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
	t.Helper()
	qr, err := exec(t, conn, query)
	require.NoError(t, err, "for query: "+query)
	return qr
}

func exec(t *testing.T, conn *mysql.Conn, query string) (*sqltypes.Result, error) {
	t.Helper()
	return conn.ExecuteFetch(query, 1000, true)
}
<|endoftext|>


<commit_before>// +build !nopkcs11

// Package Key implements crypto.Signer for PKCS #11 private
// keys. Currently, only RSA keys are supported.
// See ftp://ftp.rsasecurity.com/pub/pkcs/pkcs-11/v2-30/pkcs-11v2-30b-d6.pdf for
// details of the Cryptoki PKCS#11 API.
package pkcs11key

import (
	"crypto"
	"crypto/rsa"
	"errors"
	"fmt"
	"io"
	"math/big"
	"sync"

	"github.com/miekg/pkcs11"
)

// from src/pkg/crypto/rsa/pkcs1v15.go
var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}

// Key is an implementation of the crypto.Signer interface using a key stored
// in a PKCS#11 hardware token.  This enables the use of PKCS#11 tokens with
// the Go x509 library's methods for signing certificates.
//
// Each Key represents one session. Its session handle is protected internally
// by a mutex, so at most one Sign operation can be active at a time. For best
// performance you may want to instantiate multiple Keys using pkcs11key.Pool.
// Each one will have its own session and can be used concurrently. Note that
// some smartcards like the Yubikey Neo do not support multiple simultaneous
// sessions and will error out on creation of the second Key object.
//
// Note: If you instantiate multiple Keys without using Pool, it is *highly*
// recommended that you create all your Key objects serially, on your main
// thread, checking for errors each time, and then farm them out for use by
// different goroutines. If you fail to do this, your application may attempt
// to login repeatedly with an incorrect PIN, locking the PKCS#11 token.
type Key struct {
	// The PKCS#11 library to use
	module *pkcs11.Ctx

	// The path to the PKCS#11 library
	modulePath string

	// The label of the token to be used (mandatory).
	// We will automatically search for this in the slot list.
	tokenLabel string

	// The PIN to be used to log in to the device
	pin string

	// The public key corresponding to the private key.
	publicKey rsa.PublicKey

	// The an ObjectHandle pointing to the private key on the HSM.
	privateKeyHandle pkcs11.ObjectHandle

	// A handle to the session used by this Key.
	session *pkcs11.SessionHandle
	sessionMu sync.Mutex
}

var modules = make(map[string]*pkcs11.Ctx);
var modulesMu sync.Mutex;

// initialize loads the given PKCS#11 module (shared library) if it is not
// already loaded. It's an error to load a PKCS#11 module multiple times, so we
// maintain a map of loaded modules. Note that there is no facility yet to
// unload a module ("finalize" in PKCS#11 parlance). In general, modules will
// be unloaded at the end of the process.  The only place where you are likely
// to need to explicitly unload a module is if you fork your process after a
// Key has already been created, and the child process also needs to use
// that module.
func initialize(modulePath string) (*pkcs11.Ctx, error) {
	modulesMu.Lock()
	defer modulesMu.Unlock()
	module, ok := modules[modulePath]
	if ok {
		return module, nil
	}

	module = pkcs11.New(modulePath)

	if module == nil {
		return nil, fmt.Errorf("unable to load PKCS#11 module")
	}

	err := module.Initialize()
	if err != nil {
		return nil, err
	}

	modules[modulePath] = module

	return module, nil
}

// New instantiates a new handle to a PKCS #11-backed key.
func New(modulePath, tokenLabel, pin, privateKeyLabel string) (ps *Key, err error) {
	module, err := initialize(modulePath)
	if err != nil {
		return
	}
	if module == nil {
		err = fmt.Errorf("nil module")
		return
	}

	// Initialize a partial key
	ps = &Key{
		module:          module,
		modulePath:      modulePath,
		tokenLabel:      tokenLabel,
		pin:             pin,
	}

	// Open a session
	ps.sessionMu.Lock()
	defer ps.sessionMu.Unlock()
	session, err := ps.openSession()
	if err != nil {
		return
	}
	ps.session = &session

	// Fetch the private key by its label
	privateKeyHandle, err := getPrivateKey(module, session, privateKeyLabel)
	if err != nil {
		ps.module.CloseSession(session)
		return
	}
	ps.privateKeyHandle = privateKeyHandle

	publicKey, err := getPublicKey(module, session, privateKeyHandle)
	if err != nil {
		ps.module.CloseSession(session)
		return
	}
	ps.publicKey = publicKey

	return
}

func getPrivateKey(module *pkcs11.Ctx, session pkcs11.SessionHandle, label string) (pkcs11.ObjectHandle, error) {
	var noHandle pkcs11.ObjectHandle
	template := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_CLASS, pkcs11.CKO_PRIVATE_KEY),
		pkcs11.NewAttribute(pkcs11.CKA_LABEL, label),
	}
	if err := module.FindObjectsInit(session, template); err != nil {
		return noHandle, err
	}
	objs, _, err := module.FindObjects(session, 2)
	if err != nil {
		return noHandle, err
	}
	if err = module.FindObjectsFinal(session); err != nil {
		return noHandle, err
	}

	if len(objs) == 0 {
		return noHandle, fmt.Errorf("private key not found")
	}
	return objs[0], nil
}

// Get the public key matching a private key
// TODO: Add support for non-RSA keys, switching on CKA_KEY_TYPE
func getPublicKey(module *pkcs11.Ctx, session pkcs11.SessionHandle, privateKeyHandle pkcs11.ObjectHandle) (rsa.PublicKey, error) {
	var noKey rsa.PublicKey
	template := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_MODULUS, nil),
		pkcs11.NewAttribute(pkcs11.CKA_PUBLIC_EXPONENT, nil),
	}
	attr, err := module.GetAttributeValue(session, privateKeyHandle, template)
	if err != nil {
		return noKey, err
	}

	n := big.NewInt(0)
	e := int(0)
	gotModulus, gotExponent := false, false
	for _, a := range attr {
		if a.Type == pkcs11.CKA_MODULUS {
			n.SetBytes(a.Value)
			gotModulus = true
		} else if a.Type == pkcs11.CKA_PUBLIC_EXPONENT {
			bigE := big.NewInt(0)
			bigE.SetBytes(a.Value)
			e = int(bigE.Int64())
			gotExponent = true
		}
	}
	if !gotModulus || !gotExponent {
		return noKey, errors.New("public key missing either modulus or exponent")
	}
	return rsa.PublicKey{
		N: n,
		E: e,
	}, nil
}

// Destroy tears down a Key by closing the session. It should be
// called before the key gets GC'ed, to avoid leaving dangling sessions.
func (ps *Key) Destroy() {
	if ps.session != nil {
		// NOTE: We do not want to call module.Logout here. module.Logout applies
		// application-wide. So if there are multiple sessions active, the other ones
		// would be logged out as well, causing CKR_OBJECT_HANDLE_INVALID next
		// time they try to sign something. It's also unnecessary to log out explicitly:
		// module.CloseSession will log out once the last session in the application is
		// closed.
		ps.sessionMu.Lock()
		ps.module.CloseSession(*ps.session)
		ps.session = nil
		ps.sessionMu.Unlock()
	}
}

func (ps *Key) openSession() (pkcs11.SessionHandle, error) {
	var noSession pkcs11.SessionHandle
	slots, err := ps.module.GetSlotList(true)
	if err != nil {
		return noSession, err
	}

	for _, slot := range slots {
		// Check that token label matches.
		tokenInfo, err := ps.module.GetTokenInfo(slot)
		if err != nil {
			return noSession, err
		}
		if tokenInfo.Label != ps.tokenLabel {
			continue
		}

		// Open session
		session, err := ps.module.OpenSession(slot, pkcs11.CKF_SERIAL_SESSION)
		if err != nil {
			return session, err
		}

		// Login
		// Note: Logged-in status is application-wide, not per session. But in
		// practice it appears to be okay to login to a token multiple times with the same
		// credentials.
		if err = ps.module.Login(session, pkcs11.CKU_USER, ps.pin); err != nil {
			ps.module.CloseSession(session)
			return session, err
		}

		return session, err
	}
	return noSession, fmt.Errorf("No slot found matching token label '%s'", ps.tokenLabel)
}

// Public returns the public key for the PKCS #11 key.
func (ps *Key) Public() crypto.PublicKey {
	return &ps.publicKey
}

// Sign performs a signature using the PKCS #11 key.
func (ps *Key) Sign(rand io.Reader, msg []byte, opts crypto.SignerOpts) (signature []byte, err error) {
	ps.sessionMu.Lock()
	defer ps.sessionMu.Unlock()
	if ps.session == nil {
		return nil, errors.New("Session was nil")
	}

	// Verify that the length of the hash is as expected
	hash := opts.HashFunc()
	hashLen := hash.Size()
	if len(msg) != hashLen {
		err = errors.New("input size does not match hash function output size")
		return
	}

	// Add DigestInfo prefix
	// TODO: Switch mechanisms based on CKA_KEY_TYPE
	mechanism := []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_RSA_PKCS, nil)}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		err = errors.New("unknown hash function")
		return
	}
	signatureInput := append(prefix, msg...)

	// Perform the sign operation
	err = ps.module.SignInit(*ps.session, mechanism, ps.privateKeyHandle)
	if err != nil {
		return nil, fmt.Errorf("SignInit problem: %s", err)
	}

	signature, err = ps.module.Sign(*ps.session, signatureInput)
	return
}
<commit_msg>Fix lint error.<commit_after>// +build !nopkcs11

// Package pkcs11key implements crypto.Signer for PKCS #11 private
// keys. Currently, only RSA keys are supported.
// See ftp://ftp.rsasecurity.com/pub/pkcs/pkcs-11/v2-30/pkcs-11v2-30b-d6.pdf for
// details of the Cryptoki PKCS#11 API.
package pkcs11key

import (
	"crypto"
	"crypto/rsa"
	"errors"
	"fmt"
	"io"
	"math/big"
	"sync"

	"github.com/miekg/pkcs11"
)

// from src/pkg/crypto/rsa/pkcs1v15.go
var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}

// Key is an implementation of the crypto.Signer interface using a key stored
// in a PKCS#11 hardware token.  This enables the use of PKCS#11 tokens with
// the Go x509 library's methods for signing certificates.
//
// Each Key represents one session. Its session handle is protected internally
// by a mutex, so at most one Sign operation can be active at a time. For best
// performance you may want to instantiate multiple Keys using pkcs11key.Pool.
// Each one will have its own session and can be used concurrently. Note that
// some smartcards like the Yubikey Neo do not support multiple simultaneous
// sessions and will error out on creation of the second Key object.
//
// Note: If you instantiate multiple Keys without using Pool, it is *highly*
// recommended that you create all your Key objects serially, on your main
// thread, checking for errors each time, and then farm them out for use by
// different goroutines. If you fail to do this, your application may attempt
// to login repeatedly with an incorrect PIN, locking the PKCS#11 token.
type Key struct {
	// The PKCS#11 library to use
	module *pkcs11.Ctx

	// The path to the PKCS#11 library
	modulePath string

	// The label of the token to be used (mandatory).
	// We will automatically search for this in the slot list.
	tokenLabel string

	// The PIN to be used to log in to the device
	pin string

	// The public key corresponding to the private key.
	publicKey rsa.PublicKey

	// The an ObjectHandle pointing to the private key on the HSM.
	privateKeyHandle pkcs11.ObjectHandle

	// A handle to the session used by this Key.
	session *pkcs11.SessionHandle
	sessionMu sync.Mutex
}

var modules = make(map[string]*pkcs11.Ctx);
var modulesMu sync.Mutex;

// initialize loads the given PKCS#11 module (shared library) if it is not
// already loaded. It's an error to load a PKCS#11 module multiple times, so we
// maintain a map of loaded modules. Note that there is no facility yet to
// unload a module ("finalize" in PKCS#11 parlance). In general, modules will
// be unloaded at the end of the process.  The only place where you are likely
// to need to explicitly unload a module is if you fork your process after a
// Key has already been created, and the child process also needs to use
// that module.
func initialize(modulePath string) (*pkcs11.Ctx, error) {
	modulesMu.Lock()
	defer modulesMu.Unlock()
	module, ok := modules[modulePath]
	if ok {
		return module, nil
	}

	module = pkcs11.New(modulePath)

	if module == nil {
		return nil, fmt.Errorf("unable to load PKCS#11 module")
	}

	err := module.Initialize()
	if err != nil {
		return nil, err
	}

	modules[modulePath] = module

	return module, nil
}

// New instantiates a new handle to a PKCS #11-backed key.
func New(modulePath, tokenLabel, pin, privateKeyLabel string) (ps *Key, err error) {
	module, err := initialize(modulePath)
	if err != nil {
		return
	}
	if module == nil {
		err = fmt.Errorf("nil module")
		return
	}

	// Initialize a partial key
	ps = &Key{
		module:          module,
		modulePath:      modulePath,
		tokenLabel:      tokenLabel,
		pin:             pin,
	}

	// Open a session
	ps.sessionMu.Lock()
	defer ps.sessionMu.Unlock()
	session, err := ps.openSession()
	if err != nil {
		return
	}
	ps.session = &session

	// Fetch the private key by its label
	privateKeyHandle, err := getPrivateKey(module, session, privateKeyLabel)
	if err != nil {
		ps.module.CloseSession(session)
		return
	}
	ps.privateKeyHandle = privateKeyHandle

	publicKey, err := getPublicKey(module, session, privateKeyHandle)
	if err != nil {
		ps.module.CloseSession(session)
		return
	}
	ps.publicKey = publicKey

	return
}

func getPrivateKey(module *pkcs11.Ctx, session pkcs11.SessionHandle, label string) (pkcs11.ObjectHandle, error) {
	var noHandle pkcs11.ObjectHandle
	template := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_CLASS, pkcs11.CKO_PRIVATE_KEY),
		pkcs11.NewAttribute(pkcs11.CKA_LABEL, label),
	}
	if err := module.FindObjectsInit(session, template); err != nil {
		return noHandle, err
	}
	objs, _, err := module.FindObjects(session, 2)
	if err != nil {
		return noHandle, err
	}
	if err = module.FindObjectsFinal(session); err != nil {
		return noHandle, err
	}

	if len(objs) == 0 {
		return noHandle, fmt.Errorf("private key not found")
	}
	return objs[0], nil
}

// Get the public key matching a private key
// TODO: Add support for non-RSA keys, switching on CKA_KEY_TYPE
func getPublicKey(module *pkcs11.Ctx, session pkcs11.SessionHandle, privateKeyHandle pkcs11.ObjectHandle) (rsa.PublicKey, error) {
	var noKey rsa.PublicKey
	template := []*pkcs11.Attribute{
		pkcs11.NewAttribute(pkcs11.CKA_MODULUS, nil),
		pkcs11.NewAttribute(pkcs11.CKA_PUBLIC_EXPONENT, nil),
	}
	attr, err := module.GetAttributeValue(session, privateKeyHandle, template)
	if err != nil {
		return noKey, err
	}

	n := big.NewInt(0)
	e := int(0)
	gotModulus, gotExponent := false, false
	for _, a := range attr {
		if a.Type == pkcs11.CKA_MODULUS {
			n.SetBytes(a.Value)
			gotModulus = true
		} else if a.Type == pkcs11.CKA_PUBLIC_EXPONENT {
			bigE := big.NewInt(0)
			bigE.SetBytes(a.Value)
			e = int(bigE.Int64())
			gotExponent = true
		}
	}
	if !gotModulus || !gotExponent {
		return noKey, errors.New("public key missing either modulus or exponent")
	}
	return rsa.PublicKey{
		N: n,
		E: e,
	}, nil
}

// Destroy tears down a Key by closing the session. It should be
// called before the key gets GC'ed, to avoid leaving dangling sessions.
func (ps *Key) Destroy() {
	if ps.session != nil {
		// NOTE: We do not want to call module.Logout here. module.Logout applies
		// application-wide. So if there are multiple sessions active, the other ones
		// would be logged out as well, causing CKR_OBJECT_HANDLE_INVALID next
		// time they try to sign something. It's also unnecessary to log out explicitly:
		// module.CloseSession will log out once the last session in the application is
		// closed.
		ps.sessionMu.Lock()
		ps.module.CloseSession(*ps.session)
		ps.session = nil
		ps.sessionMu.Unlock()
	}
}

func (ps *Key) openSession() (pkcs11.SessionHandle, error) {
	var noSession pkcs11.SessionHandle
	slots, err := ps.module.GetSlotList(true)
	if err != nil {
		return noSession, err
	}

	for _, slot := range slots {
		// Check that token label matches.
		tokenInfo, err := ps.module.GetTokenInfo(slot)
		if err != nil {
			return noSession, err
		}
		if tokenInfo.Label != ps.tokenLabel {
			continue
		}

		// Open session
		session, err := ps.module.OpenSession(slot, pkcs11.CKF_SERIAL_SESSION)
		if err != nil {
			return session, err
		}

		// Login
		// Note: Logged-in status is application-wide, not per session. But in
		// practice it appears to be okay to login to a token multiple times with the same
		// credentials.
		if err = ps.module.Login(session, pkcs11.CKU_USER, ps.pin); err != nil {
			ps.module.CloseSession(session)
			return session, err
		}

		return session, err
	}
	return noSession, fmt.Errorf("No slot found matching token label '%s'", ps.tokenLabel)
}

// Public returns the public key for the PKCS #11 key.
func (ps *Key) Public() crypto.PublicKey {
	return &ps.publicKey
}

// Sign performs a signature using the PKCS #11 key.
func (ps *Key) Sign(rand io.Reader, msg []byte, opts crypto.SignerOpts) (signature []byte, err error) {
	ps.sessionMu.Lock()
	defer ps.sessionMu.Unlock()
	if ps.session == nil {
		return nil, errors.New("Session was nil")
	}

	// Verify that the length of the hash is as expected
	hash := opts.HashFunc()
	hashLen := hash.Size()
	if len(msg) != hashLen {
		err = errors.New("input size does not match hash function output size")
		return
	}

	// Add DigestInfo prefix
	// TODO: Switch mechanisms based on CKA_KEY_TYPE
	mechanism := []*pkcs11.Mechanism{pkcs11.NewMechanism(pkcs11.CKM_RSA_PKCS, nil)}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		err = errors.New("unknown hash function")
		return
	}
	signatureInput := append(prefix, msg...)

	// Perform the sign operation
	err = ps.module.SignInit(*ps.session, mechanism, ps.privateKeyHandle)
	if err != nil {
		return nil, fmt.Errorf("SignInit problem: %s", err)
	}

	signature, err = ps.module.Sign(*ps.session, signatureInput)
	return
}
<|endoftext|>


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

package statuscheck

import (
	"context"
	"sync"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/health/statuscheck/model/status"
	"github.com/ligato/cn-infra/logging"
)

const (
	// Init state means that the initialization of the plugin is in progress.
	Init PluginState = "init"
	// OK state means that the plugin is healthy.
	OK PluginState = "ok"
	// Error state means that some error has occurred in the plugin.
	Error PluginState = "error"

	periodicWriteTimeout   time.Duration = time.Second * 10 // frequency of periodic writes of state data into ETCD
	periodicProbingTimeout time.Duration = time.Second * 5  // frequency of periodic plugin state probing
)

// Plugin struct holds all plugin-related data.
type Plugin struct {
	Deps

	access sync.Mutex // lock for the Plugin data

	agentStat   *status.AgentStatus             // overall agent status
	pluginStat  map[string]*status.PluginStatus // plugin's status
	pluginProbe map[string]PluginStateProbe     // registered status probes

	cancel context.CancelFunc // cancel can be used to cancel all goroutines and their jobs inside of the plugin
	wg     sync.WaitGroup     // wait group that allows to wait until all goroutines of the plugin have finished
}

// Deps is here to group injected dependencies of plugin
// to not mix with other plugin fields.
type Deps struct {
	Log        logging.PluginLogger       //inject
	PluginName core.PluginName            //inject
	Transport  datasync.KeyProtoValWriter //inject optional
}

// Init is the plugin entry point called by the Agent Core.
func (p *Plugin) Init() error {
	p.Log.SetLevel(logging.DebugLevel)

	// write initial status data into ETCD
	p.agentStat = &status.AgentStatus{
		BuildVersion: core.BuildVersion,
		BuildDate:    core.BuildDate,
		State:        status.OperationalState_INIT,
		StartTime:    time.Now().Unix(),
		LastChange:   time.Now().Unix(),
	}

	// init pluginStat map
	p.pluginStat = make(map[string]*status.PluginStatus)

	// init map with plugin state probes
	p.pluginProbe = make(map[string]PluginStateProbe)

	// prepare context for all go routines
	var ctx context.Context
	ctx, p.cancel = context.WithCancel(context.Background())

	// do periodic status probing for plugins that have provided the probe function
	go p.periodicProbing(ctx)

	// do periodic updates of the state data in ETCD
	go p.periodicUpdates(ctx)

	return nil
}

// AfterInit is called by the Agent Core after all plugins have been initialized.
func (p *Plugin) AfterInit() error {
	p.access.Lock()
	defer p.access.Unlock()

	p.publishAgentData()

	// transition to OK state if there are no plugins
	if len(p.pluginStat) == 0 {
		p.agentStat.State = status.OperationalState_OK
		p.agentStat.LastChange = time.Now().Unix()
		p.publishAgentData()
	}

	return nil
}

// Close is called by the Agent Core when it's time to clean up the plugin.
func (p *Plugin) Close() error {
	p.cancel()
	p.wg.Wait()

	return nil
}

// Register a plugin for status change reporting.
func (p *Plugin) Register(pluginName core.PluginName, probe PluginStateProbe) {
	p.access.Lock()
	defer p.access.Unlock()

	stat := &status.PluginStatus{
		State:      status.OperationalState_INIT,
		LastChange: time.Now().Unix(),
	}
	p.pluginStat[string(pluginName)] = stat

	if probe != nil {
		p.pluginProbe[string(pluginName)] = probe
	}

	// write initial status data into ETCD
	p.publishPluginData(pluginName, stat)

	p.Log.Infof("Plugin %v: status check probe registered", pluginName)
}

// ReportStateChange can be used to report a change in the status of a previously registered plugin.
func (p *Plugin) ReportStateChange(pluginName core.PluginName, state PluginState, lastError error) {
	p.access.Lock()
	defer p.access.Unlock()

	stat, ok := p.pluginStat[string(pluginName)]
	if !ok {
		p.Log.Errorf("Unregistered plugin %s is reporting the state, ignoring.", pluginName)
		return
	}

	// update the state only if it has really changed
	changed := true
	if stateToProto(state) == stat.State {
		if lastError == nil && stat.Error == "" {
			changed = false
		}
		if lastError != nil && lastError.Error() == stat.Error {
			changed = false
		}
	}
	if !changed {
		return
	}

	p.Log.WithFields(map[string]interface{}{"plugin": pluginName, "state": state, "lastErr": lastError}).Debug(
		"Agent plugin state update.")

	// update plugin state
	stat.State = stateToProto(state)
	stat.LastChange = time.Now().Unix()
	if lastError != nil {
		stat.Error = lastError.Error()
	} else {
		stat.Error = ""
	}
	p.publishPluginData(pluginName, stat)

	// update global state if needed
	changeGlobalState := true
	if state == OK {
		// by transition to OK state, check if all plugins are OK
		for _, s := range p.pluginStat {
			if s.State != status.OperationalState_OK {
				changeGlobalState = false
				break
			}
		}
	}
	if changeGlobalState {
		p.agentStat.State = stateToProto(state)
		p.agentStat.LastChange = time.Now().Unix()
		p.publishAgentData()
	}
}

// publishAgentData writes the current global agent state into ETCD.
func (p *Plugin) publishAgentData() error {
	p.agentStat.LastUpdate = time.Now().Unix()
	if p.Transport != nil {
		return p.Transport.Put(status.AgentStatusKey(), p.agentStat)
	}
	return nil
}

// publishPluginData writes the current plugin state into ETCD.
func (p *Plugin) publishPluginData(pluginName core.PluginName, pluginStat *status.PluginStatus) error {
	pluginStat.LastUpdate = time.Now().Unix()
	if p.Transport != nil {
		return p.Transport.Put(status.PluginStatusKey(string(pluginName)), pluginStat)
	}
	return nil
}

// publishAllData publishes global agent + all plugins state data into ETCD.
func (p *Plugin) publishAllData() {
	p.access.Lock()
	defer p.access.Unlock()

	p.publishAgentData()
	for name, s := range p.pluginStat {
		p.publishPluginData(core.PluginName(name), s)
	}
}

// periodicProbing does periodic status probing for all plugins that have registered probe functions.
func (p *Plugin) periodicProbing(ctx context.Context) {
	p.wg.Add(1)
	defer p.wg.Done()

	for {
		select {
		case <-time.After(periodicProbingTimeout):
			for pluginName, probe := range p.pluginProbe {
				state, lastErr := probe()
				p.ReportStateChange(core.PluginName(pluginName), state, lastErr)
				// just check in-between probes if the plugin is closing
				select {
				case <-ctx.Done():
					return
				default:
					continue
				}
			}

		case <-ctx.Done():
			return
		}
	}
}

// periodicUpdates does periodic writes of state data into ETCD.
func (p *Plugin) periodicUpdates(ctx context.Context) {
	p.wg.Add(1)
	defer p.wg.Done()

	for {
		select {
		case <-time.After(periodicWriteTimeout):
			p.publishAllData()

		case <-ctx.Done():
			return
		}
	}
}

// getAgentState return current global operational state of the agent.
func (p *Plugin) getAgentState() status.OperationalState {
	p.access.Lock()
	defer p.access.Unlock()
	return p.agentStat.State
}

// GetAgentStatus return current global operational state of the agent.
func (p *Plugin) GetAgentStatus() status.AgentStatus {
	p.access.Lock()
	defer p.access.Unlock()
	return *(p.agentStat)
}

// stateToProto converts agent state type into protobuf agent state type.
func stateToProto(state PluginState) status.OperationalState {
	switch state {
	case Init:
		return status.OperationalState_INIT
	case OK:
		return status.OperationalState_OK
	default:
		return status.OperationalState_ERROR
	}
}
<commit_msg>remove debug logs from statuscheck<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package statuscheck

import (
	"context"
	"sync"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync"
	"github.com/ligato/cn-infra/health/statuscheck/model/status"
	"github.com/ligato/cn-infra/logging"
)

const (
	// Init state means that the initialization of the plugin is in progress.
	Init PluginState = "init"
	// OK state means that the plugin is healthy.
	OK PluginState = "ok"
	// Error state means that some error has occurred in the plugin.
	Error PluginState = "error"

	periodicWriteTimeout   time.Duration = time.Second * 10 // frequency of periodic writes of state data into ETCD
	periodicProbingTimeout time.Duration = time.Second * 5  // frequency of periodic plugin state probing
)

// Plugin struct holds all plugin-related data.
type Plugin struct {
	Deps

	access sync.Mutex // lock for the Plugin data

	agentStat   *status.AgentStatus             // overall agent status
	pluginStat  map[string]*status.PluginStatus // plugin's status
	pluginProbe map[string]PluginStateProbe     // registered status probes

	cancel context.CancelFunc // cancel can be used to cancel all goroutines and their jobs inside of the plugin
	wg     sync.WaitGroup     // wait group that allows to wait until all goroutines of the plugin have finished
}

// Deps is here to group injected dependencies of plugin
// to not mix with other plugin fields.
type Deps struct {
	Log        logging.PluginLogger       //inject
	PluginName core.PluginName            //inject
	Transport  datasync.KeyProtoValWriter //inject optional
}

// Init is the plugin entry point called by the Agent Core.
func (p *Plugin) Init() error {
	// write initial status data into ETCD
	p.agentStat = &status.AgentStatus{
		BuildVersion: core.BuildVersion,
		BuildDate:    core.BuildDate,
		State:        status.OperationalState_INIT,
		StartTime:    time.Now().Unix(),
		LastChange:   time.Now().Unix(),
	}

	// init pluginStat map
	p.pluginStat = make(map[string]*status.PluginStatus)

	// init map with plugin state probes
	p.pluginProbe = make(map[string]PluginStateProbe)

	// prepare context for all go routines
	var ctx context.Context
	ctx, p.cancel = context.WithCancel(context.Background())

	// do periodic status probing for plugins that have provided the probe function
	go p.periodicProbing(ctx)

	// do periodic updates of the state data in ETCD
	go p.periodicUpdates(ctx)

	return nil
}

// AfterInit is called by the Agent Core after all plugins have been initialized.
func (p *Plugin) AfterInit() error {
	p.access.Lock()
	defer p.access.Unlock()

	p.publishAgentData()

	// transition to OK state if there are no plugins
	if len(p.pluginStat) == 0 {
		p.agentStat.State = status.OperationalState_OK
		p.agentStat.LastChange = time.Now().Unix()
		p.publishAgentData()
	}

	return nil
}

// Close is called by the Agent Core when it's time to clean up the plugin.
func (p *Plugin) Close() error {
	p.cancel()
	p.wg.Wait()

	return nil
}

// Register a plugin for status change reporting.
func (p *Plugin) Register(pluginName core.PluginName, probe PluginStateProbe) {
	p.access.Lock()
	defer p.access.Unlock()

	stat := &status.PluginStatus{
		State:      status.OperationalState_INIT,
		LastChange: time.Now().Unix(),
	}
	p.pluginStat[string(pluginName)] = stat

	if probe != nil {
		p.pluginProbe[string(pluginName)] = probe
	}

	// write initial status data into ETCD
	p.publishPluginData(pluginName, stat)

	p.Log.Infof("Plugin %v: status check probe registered", pluginName)
}

// ReportStateChange can be used to report a change in the status of a previously registered plugin.
func (p *Plugin) ReportStateChange(pluginName core.PluginName, state PluginState, lastError error) {
	p.access.Lock()
	defer p.access.Unlock()

	stat, ok := p.pluginStat[string(pluginName)]
	if !ok {
		p.Log.Errorf("Unregistered plugin %s is reporting the state, ignoring.", pluginName)
		return
	}

	// update the state only if it has really changed
	changed := true
	if stateToProto(state) == stat.State {
		if lastError == nil && stat.Error == "" {
			changed = false
		}
		if lastError != nil && lastError.Error() == stat.Error {
			changed = false
		}
	}
	if !changed {
		return
	}

	p.Log.WithFields(map[string]interface{}{"plugin": pluginName, "state": state, "lastErr": lastError}).Info(
		"Agent plugin state update.")

	// update plugin state
	stat.State = stateToProto(state)
	stat.LastChange = time.Now().Unix()
	if lastError != nil {
		stat.Error = lastError.Error()
	} else {
		stat.Error = ""
	}
	p.publishPluginData(pluginName, stat)

	// update global state if needed
	changeGlobalState := true
	if state == OK {
		// by transition to OK state, check if all plugins are OK
		for _, s := range p.pluginStat {
			if s.State != status.OperationalState_OK {
				changeGlobalState = false
				break
			}
		}
	}
	if changeGlobalState {
		p.agentStat.State = stateToProto(state)
		p.agentStat.LastChange = time.Now().Unix()
		p.publishAgentData()
	}
}

// publishAgentData writes the current global agent state into ETCD.
func (p *Plugin) publishAgentData() error {
	p.agentStat.LastUpdate = time.Now().Unix()
	if p.Transport != nil {
		return p.Transport.Put(status.AgentStatusKey(), p.agentStat)
	}
	return nil
}

// publishPluginData writes the current plugin state into ETCD.
func (p *Plugin) publishPluginData(pluginName core.PluginName, pluginStat *status.PluginStatus) error {
	pluginStat.LastUpdate = time.Now().Unix()
	if p.Transport != nil {
		return p.Transport.Put(status.PluginStatusKey(string(pluginName)), pluginStat)
	}
	return nil
}

// publishAllData publishes global agent + all plugins state data into ETCD.
func (p *Plugin) publishAllData() {
	p.access.Lock()
	defer p.access.Unlock()

	p.publishAgentData()
	for name, s := range p.pluginStat {
		p.publishPluginData(core.PluginName(name), s)
	}
}

// periodicProbing does periodic status probing for all plugins that have registered probe functions.
func (p *Plugin) periodicProbing(ctx context.Context) {
	p.wg.Add(1)
	defer p.wg.Done()

	for {
		select {
		case <-time.After(periodicProbingTimeout):
			for pluginName, probe := range p.pluginProbe {
				state, lastErr := probe()
				p.ReportStateChange(core.PluginName(pluginName), state, lastErr)
				// just check in-between probes if the plugin is closing
				select {
				case <-ctx.Done():
					return
				default:
					continue
				}
			}

		case <-ctx.Done():
			return
		}
	}
}

// periodicUpdates does periodic writes of state data into ETCD.
func (p *Plugin) periodicUpdates(ctx context.Context) {
	p.wg.Add(1)
	defer p.wg.Done()

	for {
		select {
		case <-time.After(periodicWriteTimeout):
			p.publishAllData()

		case <-ctx.Done():
			return
		}
	}
}

// getAgentState return current global operational state of the agent.
func (p *Plugin) getAgentState() status.OperationalState {
	p.access.Lock()
	defer p.access.Unlock()
	return p.agentStat.State
}

// GetAgentStatus return current global operational state of the agent.
func (p *Plugin) GetAgentStatus() status.AgentStatus {
	p.access.Lock()
	defer p.access.Unlock()
	return *(p.agentStat)
}

// stateToProto converts agent state type into protobuf agent state type.
func stateToProto(state PluginState) status.OperationalState {
	switch state {
	case Init:
		return status.OperationalState_INIT
	case OK:
		return status.OperationalState_OK
	default:
		return status.OperationalState_ERROR
	}
}
<|endoftext|>


<commit_before>package vault

import (
	"context"
	"errors"
	"time"

	"github.com/hashicorp/vault/logical"
)

var (
	// ErrBarrierSealed is returned if an operation is performed on
	// a sealed barrier. No operation is expected to succeed before unsealing
	ErrBarrierSealed = errors.New("Vault is sealed")

	// ErrBarrierAlreadyInit is returned if the barrier is already
	// initialized. This prevents a re-initialization.
	ErrBarrierAlreadyInit = errors.New("Vault is already initialized")

	// ErrBarrierNotInit is returned if a non-initialized barrier
	// is attempted to be unsealed.
	ErrBarrierNotInit = errors.New("Vault is not initialized")

	// ErrBarrierInvalidKey is returned if the Unseal key is invalid
	ErrBarrierInvalidKey = errors.New("Unseal failed, invalid key")
)

const (
	// barrierInitPath is the path used to store our init sentinel file
	barrierInitPath = "barrier/init"

	// keyringPath is the location of the keyring data. This is encrypted
	// by the master key.
	keyringPath   = "core/keyring"
	keyringPrefix = "core/"

	// keyringUpgradePrefix is the path used to store keyring update entries.
	// When running in HA mode, the active instance will install the new key
	// and re-write the keyring. For standby instances, they need an upgrade
	// path from key N to N+1. They cannot just use the master key because
	// in the event of a rekey, that master key can no longer decrypt the keyring.
	// When key N+1 is installed, we create an entry at "prefix/N" which uses
	// encryption key N to provide the N+1 key. The standby instances scan
	// for this periodically and refresh their keyring. The upgrade keys
	// are deleted after a few minutes, but this provides enough time for the
	// standby instances to upgrade without causing any disruption.
	keyringUpgradePrefix = "core/upgrade/"

	// masterKeyPath is the location of the master key. This is encrypted
	// by the latest key in the keyring. This is only used by standby instances
	// to handle the case of a rekey. If the active instance does a rekey,
	// the standby instances can no longer reload the keyring since they
	// have the old master key. This key can be decrypted if you have the
	// keyring to discover the new master key. The new master key is then
	// used to reload the keyring itself.
	masterKeyPath = "core/master"
)

// SecurityBarrier is a critical component of Vault. It is used to wrap
// an untrusted physical backend and provide a single point of encryption,
// decryption and checksum verification. The goal is to ensure that any
// data written to the barrier is confidential and that integrity is preserved.
// As a real-world analogy, this is the steel and concrete wrapper around
// a Vault. The barrier should only be Unlockable given its key.
type SecurityBarrier interface {
	// Initialized checks if the barrier has been initialized
	// and has a master key set.
	Initialized(ctx context.Context) (bool, error)

	// Initialize works only if the barrier has not been initialized
	// and makes use of the given master key.
	Initialize(context.Context, []byte) error

	// GenerateKey is used to generate a new key
	GenerateKey() ([]byte, error)

	// KeyLength is used to sanity check a key
	KeyLength() (int, int)

	// Sealed checks if the barrier has been unlocked yet. The Barrier
	// is not expected to be able to perform any CRUD until it is unsealed.
	Sealed() (bool, error)

	// Unseal is used to provide the master key which permits the barrier
	// to be unsealed. If the key is not correct, the barrier remains sealed.
	Unseal(ctx context.Context, key []byte) error

	// VerifyMaster is used to check if the given key matches the master key
	VerifyMaster(key []byte) error

	// SetMasterKey is used to directly set a new master key. This is used in
	// repliated scenarios due to the chicken and egg problem of reloading the
	// keyring from disk before we have the master key to decrypt it.
	SetMasterKey(key []byte) error

	// ReloadKeyring is used to re-read the underlying keyring.
	// This is used for HA deployments to ensure the latest keyring
	// is present in the leader.
	ReloadKeyring(ctx context.Context) error

	// ReloadMasterKey is used to re-read the underlying masterkey.
	// This is used for HA deployments to ensure the latest master key
	// is available for keyring reloading.
	ReloadMasterKey(ctx context.Context) error

	// Seal is used to re-seal the barrier. This requires the barrier to
	// be unsealed again to perform any further operations.
	Seal() error

	// Rotate is used to create a new encryption key. All future writes
	// should use the new key, while old values should still be decryptable.
	Rotate(ctx context.Context) (uint32, error)

	// CreateUpgrade creates an upgrade path key to the given term from the previous term
	CreateUpgrade(ctx context.Context, term uint32) error

	// DestroyUpgrade destroys the upgrade path key to the given term
	DestroyUpgrade(ctx context.Context, term uint32) error

	// CheckUpgrade looks for an upgrade to the current term and installs it
	CheckUpgrade(ctx context.Context) (bool, uint32, error)

	// ActiveKeyInfo is used to inform details about the active key
	ActiveKeyInfo() (*KeyInfo, error)

	// Rekey is used to change the master key used to protect the keyring
	Rekey(context.Context, []byte) error

	// For replication we must send over the keyring, so this must be available
	Keyring() (*Keyring, error)

	// SecurityBarrier must provide the storage APIs
	BarrierStorage

	// SecurityBarrier must provide the encryption APIs
	BarrierEncryptor
}

// BarrierStorage is the storage only interface required for a Barrier.
type BarrierStorage interface {
	// Put is used to insert or update an entry
	Put(ctx context.Context, entry *Entry) error

	// Get is used to fetch an entry
	Get(ctx context.Context, key string) (*Entry, error)

	// Delete is used to permanently delete an entry
	Delete(ctx context.Context, key string) error

	// List is used ot list all the keys under a given
	// prefix, up to the next prefix.
	List(ctx context.Context, prefix string) ([]string, error)
}

// BarrierEncryptor is the in memory only interface that does not actually
// use the underlying barrier. It is used for lower level modules like the
// Write-Ahead-Log and Merkle index to allow them to use the barrier.
type BarrierEncryptor interface {
	Encrypt(key string, plaintext []byte) ([]byte, error)
	Decrypt(key string, ciphertext []byte) ([]byte, error)
}

// Entry is used to represent data stored by the security barrier
type Entry struct {
	Key      string
	Value    []byte
	SealWrap bool
}

// Logical turns the Entry into a logical storage entry.
func (e *Entry) Logical() *logical.StorageEntry {
	return &logical.StorageEntry{
		Key:      e.Key,
		Value:    e.Value,
		SealWrap: e.SealWrap,
	}
}

// KeyInfo is used to convey information about the encryption key
type KeyInfo struct {
	Term        int
	InstallTime time.Time
}
<commit_msg>Add context to barrier encryptor interface<commit_after>package vault

import (
	"context"
	"errors"
	"time"

	"github.com/hashicorp/vault/logical"
)

var (
	// ErrBarrierSealed is returned if an operation is performed on
	// a sealed barrier. No operation is expected to succeed before unsealing
	ErrBarrierSealed = errors.New("Vault is sealed")

	// ErrBarrierAlreadyInit is returned if the barrier is already
	// initialized. This prevents a re-initialization.
	ErrBarrierAlreadyInit = errors.New("Vault is already initialized")

	// ErrBarrierNotInit is returned if a non-initialized barrier
	// is attempted to be unsealed.
	ErrBarrierNotInit = errors.New("Vault is not initialized")

	// ErrBarrierInvalidKey is returned if the Unseal key is invalid
	ErrBarrierInvalidKey = errors.New("Unseal failed, invalid key")
)

const (
	// barrierInitPath is the path used to store our init sentinel file
	barrierInitPath = "barrier/init"

	// keyringPath is the location of the keyring data. This is encrypted
	// by the master key.
	keyringPath   = "core/keyring"
	keyringPrefix = "core/"

	// keyringUpgradePrefix is the path used to store keyring update entries.
	// When running in HA mode, the active instance will install the new key
	// and re-write the keyring. For standby instances, they need an upgrade
	// path from key N to N+1. They cannot just use the master key because
	// in the event of a rekey, that master key can no longer decrypt the keyring.
	// When key N+1 is installed, we create an entry at "prefix/N" which uses
	// encryption key N to provide the N+1 key. The standby instances scan
	// for this periodically and refresh their keyring. The upgrade keys
	// are deleted after a few minutes, but this provides enough time for the
	// standby instances to upgrade without causing any disruption.
	keyringUpgradePrefix = "core/upgrade/"

	// masterKeyPath is the location of the master key. This is encrypted
	// by the latest key in the keyring. This is only used by standby instances
	// to handle the case of a rekey. If the active instance does a rekey,
	// the standby instances can no longer reload the keyring since they
	// have the old master key. This key can be decrypted if you have the
	// keyring to discover the new master key. The new master key is then
	// used to reload the keyring itself.
	masterKeyPath = "core/master"
)

// SecurityBarrier is a critical component of Vault. It is used to wrap
// an untrusted physical backend and provide a single point of encryption,
// decryption and checksum verification. The goal is to ensure that any
// data written to the barrier is confidential and that integrity is preserved.
// As a real-world analogy, this is the steel and concrete wrapper around
// a Vault. The barrier should only be Unlockable given its key.
type SecurityBarrier interface {
	// Initialized checks if the barrier has been initialized
	// and has a master key set.
	Initialized(ctx context.Context) (bool, error)

	// Initialize works only if the barrier has not been initialized
	// and makes use of the given master key.
	Initialize(context.Context, []byte) error

	// GenerateKey is used to generate a new key
	GenerateKey() ([]byte, error)

	// KeyLength is used to sanity check a key
	KeyLength() (int, int)

	// Sealed checks if the barrier has been unlocked yet. The Barrier
	// is not expected to be able to perform any CRUD until it is unsealed.
	Sealed() (bool, error)

	// Unseal is used to provide the master key which permits the barrier
	// to be unsealed. If the key is not correct, the barrier remains sealed.
	Unseal(ctx context.Context, key []byte) error

	// VerifyMaster is used to check if the given key matches the master key
	VerifyMaster(key []byte) error

	// SetMasterKey is used to directly set a new master key. This is used in
	// repliated scenarios due to the chicken and egg problem of reloading the
	// keyring from disk before we have the master key to decrypt it.
	SetMasterKey(key []byte) error

	// ReloadKeyring is used to re-read the underlying keyring.
	// This is used for HA deployments to ensure the latest keyring
	// is present in the leader.
	ReloadKeyring(ctx context.Context) error

	// ReloadMasterKey is used to re-read the underlying masterkey.
	// This is used for HA deployments to ensure the latest master key
	// is available for keyring reloading.
	ReloadMasterKey(ctx context.Context) error

	// Seal is used to re-seal the barrier. This requires the barrier to
	// be unsealed again to perform any further operations.
	Seal() error

	// Rotate is used to create a new encryption key. All future writes
	// should use the new key, while old values should still be decryptable.
	Rotate(ctx context.Context) (uint32, error)

	// CreateUpgrade creates an upgrade path key to the given term from the previous term
	CreateUpgrade(ctx context.Context, term uint32) error

	// DestroyUpgrade destroys the upgrade path key to the given term
	DestroyUpgrade(ctx context.Context, term uint32) error

	// CheckUpgrade looks for an upgrade to the current term and installs it
	CheckUpgrade(ctx context.Context) (bool, uint32, error)

	// ActiveKeyInfo is used to inform details about the active key
	ActiveKeyInfo() (*KeyInfo, error)

	// Rekey is used to change the master key used to protect the keyring
	Rekey(context.Context, []byte) error

	// For replication we must send over the keyring, so this must be available
	Keyring() (*Keyring, error)

	// SecurityBarrier must provide the storage APIs
	BarrierStorage

	// SecurityBarrier must provide the encryption APIs
	BarrierEncryptor
}

// BarrierStorage is the storage only interface required for a Barrier.
type BarrierStorage interface {
	// Put is used to insert or update an entry
	Put(ctx context.Context, entry *Entry) error

	// Get is used to fetch an entry
	Get(ctx context.Context, key string) (*Entry, error)

	// Delete is used to permanently delete an entry
	Delete(ctx context.Context, key string) error

	// List is used ot list all the keys under a given
	// prefix, up to the next prefix.
	List(ctx context.Context, prefix string) ([]string, error)
}

// BarrierEncryptor is the in memory only interface that does not actually
// use the underlying barrier. It is used for lower level modules like the
// Write-Ahead-Log and Merkle index to allow them to use the barrier.
type BarrierEncryptor interface {
	Encrypt(ctx context.Context, key string, plaintext []byte) ([]byte, error)
	Decrypt(ctx context.Context, key string, ciphertext []byte) ([]byte, error)
}

// Entry is used to represent data stored by the security barrier
type Entry struct {
	Key      string
	Value    []byte
	SealWrap bool
}

// Logical turns the Entry into a logical storage entry.
func (e *Entry) Logical() *logical.StorageEntry {
	return &logical.StorageEntry{
		Key:      e.Key,
		Value:    e.Value,
		SealWrap: e.SealWrap,
	}
}

// KeyInfo is used to convey information about the encryption key
type KeyInfo struct {
	Term        int
	InstallTime time.Time
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Marc René Arns. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

/*
	Package smfreader provides a reader of Standard MIDI Files (SMF).

	There is also a more comfortable handler package:

	  github.com/gomidi/midi/handler    (reading MIDI messages live or from SMF files)

	Usage

		import (
			"github.com/gomidi/midi/smf/smfreader"
			"github.com/gomidi/midi/midimessage/channel"    // (Channel Messages)

			// you may also want to use these
			// github.com/gomidi/midi/midimessage/meta       (Meta Messages)
			// github.com/gomidi/midi/midimessage/cc         (Control Change Messages)
			// github.com/gomidi/midi/midimessage/syscommon  (System Common Messages)
			// github.com/gomidi/midi/midimessage/sysex      (System Exclusive Messages)
		)

		var err1, err2 error

		readMIDI := func (rd smf.Reader) {

			var m midi.Message

			for {
				m, err1 = rd.Read()

				// to interrupt, the input.Read method must return io.EOF or any other error
				if err1 != nil {
					break
				}

				// deal with them based on a type switch
				switch msg := m.(type) {
				case channel.NoteOn:
					fmt.Printf(
					  "NoteOn at channel %v: key %v velocity: %v\n",
					  msg.Channel(), // MIDI channels 1-16 correspond to msg.Channel 0-15
					  msg.Key(),
					  msg.Velocity(),
					)
				case channel.NoteOff:
					...
				}
			}

		}

		err2 = smfreader.ReadFile("file.mid", readMIDI)

		// deal with err1 and err2

*/
package smfreader
<commit_msg>fix doc<commit_after>// Copyright (c) 2017 Marc René Arns. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

/*
	Package smfreader provides a reader of Standard MIDI Files (SMF).

	There is also a more comfortable handler package:

	  github.com/gomidi/midi/midihandler    (reading MIDI messages live or from SMF files)

	Usage

		import (
			"github.com/gomidi/midi/smf/smfreader"
			"github.com/gomidi/midi/midimessage/channel"    // (Channel Messages)

			// you may also want to use these
			// github.com/gomidi/midi/midimessage/meta       (Meta Messages)
			// github.com/gomidi/midi/midimessage/cc         (Control Change Messages)
			// github.com/gomidi/midi/midimessage/syscommon  (System Common Messages)
			// github.com/gomidi/midi/midimessage/sysex      (System Exclusive Messages)
		)

		var err1, err2 error

		readMIDI := func (rd smf.Reader) {

			var m midi.Message

			for {
				m, err1 = rd.Read()

				// to interrupt, the input.Read method must return io.EOF or any other error
				if err1 != nil {
					break
				}

				// deal with them based on a type switch
				switch msg := m.(type) {
				case channel.NoteOn:
					fmt.Printf(
					  "NoteOn at channel %v: key %v velocity: %v\n",
					  msg.Channel(), // MIDI channels 1-16 correspond to msg.Channel 0-15
					  msg.Key(),
					  msg.Velocity(),
					)
				case channel.NoteOff:
					...
				}
			}

		}

		err2 = smfreader.ReadFile("file.mid", readMIDI)

		// deal with err1 and err2

*/
package smfreader
<|endoftext|>


<commit_before>package awstesting

import (
	"io"

	"github.com/aws/aws-sdk-go/private/util"
)

// ZeroReader is a io.Reader which will always write zeros to the byte slice provided.
type ZeroReader struct{}

// Read fills the provided byte slice with zeros returning the number of bytes written.
func (r *ZeroReader) Read(b []byte) (int, error) {
	for i := 0; i < len(b); i++ {
		b[i] = 0
	}
	return len(b), nil
}

// ReadCloser is a io.ReadCloser for unit testing.
// Designed to test for leaks and whether a handle has
// been closed
type ReadCloser struct {
	Size     int
	Closed   bool
	set      bool
	FillData func(bool, []byte, int, int)
}

// Read will call FillData and fill it with whatever data needed.
// Decrements the size until zero, then return io.EOF.
func (r *ReadCloser) Read(b []byte) (int, error) {
	if r.Closed {
		return 0, io.EOF
	}

	delta := len(b)
	if delta > r.Size {
		delta = r.Size
	}
	r.Size -= delta

	for i := 0; i < delta; i++ {
		b[i] = 'a'
	}

	if r.FillData != nil {
		r.FillData(r.set, b, r.Size, delta)
	}
	r.set = true

	if r.Size > 0 {
		return delta, nil
	}
	return delta, io.EOF
}

// Close sets Closed to true and returns no error
func (r *ReadCloser) Close() error {
	r.Closed = true
	return nil
}

// SortedKeys returns a sorted slice of keys of a map.
func SortedKeys(m map[string]interface{}) []string {
	return util.SortedKeys(m)
}
<commit_msg>remove unused util func<commit_after>package awstesting

import "io"

// ZeroReader is a io.Reader which will always write zeros to the byte slice provided.
type ZeroReader struct{}

// Read fills the provided byte slice with zeros returning the number of bytes written.
func (r *ZeroReader) Read(b []byte) (int, error) {
	for i := 0; i < len(b); i++ {
		b[i] = 0
	}
	return len(b), nil
}

// ReadCloser is a io.ReadCloser for unit testing.
// Designed to test for leaks and whether a handle has
// been closed
type ReadCloser struct {
	Size     int
	Closed   bool
	set      bool
	FillData func(bool, []byte, int, int)
}

// Read will call FillData and fill it with whatever data needed.
// Decrements the size until zero, then return io.EOF.
func (r *ReadCloser) Read(b []byte) (int, error) {
	if r.Closed {
		return 0, io.EOF
	}

	delta := len(b)
	if delta > r.Size {
		delta = r.Size
	}
	r.Size -= delta

	for i := 0; i < delta; i++ {
		b[i] = 'a'
	}

	if r.FillData != nil {
		r.FillData(r.set, b, r.Size, delta)
	}
	r.set = true

	if r.Size > 0 {
		return delta, nil
	}
	return delta, io.EOF
}

// Close sets Closed to true and returns no error
func (r *ReadCloser) Close() error {
	r.Closed = true
	return nil
}
<|endoftext|>


<commit_before>// Package command is used for parsing input parameters,
// and starting subcommands (aka tools).
package command

import (
	"errors"
	"flag"
	"strings"
)

// commandWordSep is a separator of words in a command name.
// For illustration, "generate stuff" is a valid name.
const commandWordSep = " "

// Context stores information about available subcommands (aka tools)
// and other related things that are necessary for their start.
type Context struct {
	list     []Handler // A list of registered subcommands (tools).
	defaultH *int      // Index of the command that will be started if no arguments received.
}

// NewContext gets a number of handlers as arguments, allocates
// a new Context and returns it.
func NewContext(handlers ...Handler) *Context {
	// Allocate a new context with handlers
	// as a list.
	c := &Context{
		list: handlers,
	}

	// Find default handler and add it to the context.
	for i := 0; i < len(handlers); i++ {
		if handlers[i].Default {
			c.defaultH = &i
			return c
		}
	}
	return c
}

// Handler is a type for representation of a subcommand (aka tool)
// such as "cli run" or "cli new".
type Handler struct {
	// Run is an entry function of the handler.
	// The args are the arguments after the command name.
	Run func(hs []Handler, i int, args []string) error

	// Default means the handler must be executed if no arguments are
	// received from user (in addition to when it is called explicitly).
	// Only first default handler is used, others will be ignored.
	Default bool

	Name  string // Name of the handler, e.g. "new" or "generate stuff".
	Usage string // Possible arguments of the command, e.g. "[input] [output]".
	Info  string // One line description of the command.
	Desc  string // Detailed description of what the command does.

	Flags flag.FlagSet // Set of flags specific to the command.
}

// ErrIncorrectArgs is returned every time a user is trying
// to request a command that does not exist.
var ErrIncorrectArgs = errors.New("incorrect command requested")

// Run gets a list of arguments and either starts an entry function of the
// requested subcommand (aka tool) or returns an error.
func (c *Context) Run(args []string) error {
	// Start default handlers if no arguments are received.
	if len(args) == 0 {
		if c.defaultH != nil {
			return c.list[*c.defaultH].Run(c.list, *c.defaultH, args)
		}
		return nil
	}

	// Otherwise, iterating over all available handlers of subcommands (aka tools).
	for i := 0; i < len(c.list); i++ {
		// Check whether current handler belongs to the subcommand (tool)
		// that is requested by user.
		if lst, ok := c.list[i].requested(args); ok {
			// Parse flags if there are any.
			err := c.list[i].Flags.Parse(lst)
			if err != nil {
				return err
			}

			// Start the entry function of the handler.
			// Use h.Flags' non-flag values as arguments.
			return c.list[i].Run(c.list, i, c.list[i].Flags.Args())
		}
	}
	return ErrIncorrectArgs
}

// requested checks whether the handler is the one that is requested by user,
// i.e. handler's name or alias is a part of args.
// It returns arguments (not including the handler name) and true in case
// of success, and nil, false otherwise.
func (h Handler) requested(args []string) ([]string, bool) {
	// Calculate the number of words in the handler's name.
	// It is equal to the number of spaces plus one.
	num := strings.Count(h.Name, commandWordSep) + 1

	// If the number of arguments is less than the number of words
	// in handler's name that means this is not the command user wants.
	if len(args) < num {
		return nil, false
	}

	// Make sure the handler's name is equal to the one user requested.
	if h.Name != strings.Join(args[:num], commandWordSep) {
		return nil, false
	}

	return args[num:], true
}
<commit_msg>Return error when no arguments and default handlers<commit_after>// Package command is used for parsing input parameters,
// and starting subcommands (aka tools).
package command

import (
	"errors"
	"flag"
	"strings"
)

// commandWordSep is a separator of words in a command name.
// For illustration, "generate stuff" is a valid name.
const commandWordSep = " "

// Context stores information about available subcommands (aka tools)
// and other related things that are necessary for their start.
type Context struct {
	list     []Handler // A list of registered subcommands (tools).
	defaultH *int      // Index of the command that will be started if no arguments received.
}

// NewContext gets a number of handlers as arguments, allocates
// a new Context and returns it.
func NewContext(handlers ...Handler) *Context {
	// Allocate a new context with handlers
	// as a list.
	c := &Context{
		list: handlers,
	}

	// Find default handler and add it to the context.
	for i := 0; i < len(handlers); i++ {
		if handlers[i].Default {
			c.defaultH = &i
			return c
		}
	}
	return c
}

// Handler is a type for representation of a subcommand (aka tool)
// such as "cli run" or "cli new".
type Handler struct {
	// Run is an entry function of the handler.
	// The args are the arguments after the command name.
	Run func(hs []Handler, i int, args []string) error

	// Default means the handler must be executed if no arguments are
	// received from user (in addition to when it is called explicitly).
	// Only first default handler is used, others will be ignored.
	Default bool

	Name  string // Name of the handler, e.g. "new" or "generate stuff".
	Usage string // Possible arguments of the command, e.g. "[input] [output]".
	Info  string // One line description of the command.
	Desc  string // Detailed description of what the command does.

	Flags flag.FlagSet // Set of flags specific to the command.
}

// ErrIncorrectArgs is returned every time a user is trying
// to request a command that does not exist.
var ErrIncorrectArgs = errors.New("incorrect command requested")

// Run gets a list of arguments and either starts an entry function of the
// requested subcommand (aka tool) or returns an error.
func (c *Context) Run(args []string) error {
	// Start default handlers if no arguments are received.
	if len(args) == 0 {
		if c.defaultH != nil {
			return c.list[*c.defaultH].Run(c.list, *c.defaultH, args)
		}
		return ErrIncorrectArgs
	}

	// Otherwise, iterating over all available handlers of subcommands (aka tools).
	for i := 0; i < len(c.list); i++ {
		// Check whether current handler belongs to the subcommand (tool)
		// that is requested by user.
		if lst, ok := c.list[i].requested(args); ok {
			// Parse flags if there are any.
			err := c.list[i].Flags.Parse(lst)
			if err != nil {
				return err
			}

			// Start the entry function of the handler.
			// Use h.Flags' non-flag values as arguments.
			return c.list[i].Run(c.list, i, c.list[i].Flags.Args())
		}
	}
	return ErrIncorrectArgs
}

// requested checks whether the handler is the one that is requested by user,
// i.e. handler's name or alias is a part of args.
// It returns arguments (not including the handler name) and true in case
// of success, and nil, false otherwise.
func (h Handler) requested(args []string) ([]string, bool) {
	// Calculate the number of words in the handler's name.
	// It is equal to the number of spaces plus one.
	num := strings.Count(h.Name, commandWordSep) + 1

	// If the number of arguments is less than the number of words
	// in handler's name that means this is not the command user wants.
	if len(args) < num {
		return nil, false
	}

	// Make sure the handler's name is equal to the one user requested.
	if h.Name != strings.Join(args[:num], commandWordSep) {
		return nil, false
	}

	return args[num:], true
}
<|endoftext|>


<commit_before>type Locale interface {
	Code() string
}
type User interface {
	Locale() string
}
<commit_msg>Fixing framework package<commit_after>package framework

type Locale interface {
	Code() string
}
type User interface {
	Locale() string
}
<|endoftext|>


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

package grammar

import (
	"fmt"

	"github.com/google/badwolf/bql/lexer"
	"github.com/google/badwolf/bql/semantic"
)

// Element are the main components that define a derivation rule.
type Element struct {
	isSymbol  bool
	symbol    semantic.Symbol
	tokenType lexer.TokenType
}

// NewSymbol creates a new element from a symbol.
func NewSymbol(s semantic.Symbol) Element {
	return Element{
		isSymbol: true,
		symbol:   s,
	}
}

// NewTokenType creates a new element from a token.
func NewTokenType(t lexer.TokenType) Element {
	return Element{
		isSymbol:  false,
		tokenType: t,
	}
}

// Symbol returns the symbol box for the given element.
func (e Element) Symbol() semantic.Symbol {
	return e.symbol
}

// Token returns the value of the token box for the given element.
func (e Element) Token() lexer.TokenType {
	return e.tokenType
}

// Clause contains on clause of the derivation rule.
type Clause struct {
	Elements         []Element
	ProcessStart     semantic.ClauseHook
	ProcessEnd       semantic.ClauseHook
	ProcessedElement semantic.ElementHook
}

// Grammar contains the left factory LLk grammar to be parsed. All provided
// grammars *must* have the "START" symbol to initialte the parsing of input
// text.
type Grammar map[semantic.Symbol][]*Clause

// Parser implements a LLk recursive decend parser for left factorized grammars.
type Parser struct {
	grammar *Grammar
}

// NewParser creates a new recursive decend parser for a left factorized
// grammar.
func NewParser(grammar *Grammar) (*Parser, error) {
	// Check that the grammar is left factorized.
	for _, clauses := range *grammar {
		idx := 0
		for _, cls := range clauses {
			if len(cls.Elements) == 0 {
				if idx == 0 {
					idx++
					continue
				}
				return nil, fmt.Errorf("grammar.NewParser: invalid extra empty clause derivation %v", clauses)
			}
			if cls.Elements[0].isSymbol {
				return nil, fmt.Errorf("grammar.NewParser: not left factored grammar in %v", clauses)
			}
		}
	}
	return &Parser{
		grammar: grammar,
	}, nil
}

// Parse attempts to run the parser for the given input.
func (p *Parser) Parse(llk *LLk, st *semantic.Statement) error {
	b, err := p.consume(llk, st, "START")
	if err != nil {
		return err
	}
	if !b {
		return fmt.Errorf("Parser.Parse: inconsitent parser, no error found, and no tokens were consumed")
	}
	return nil
}

// consume attempts to consume all input tokens for the provided symbols given
// the parser grammar.
func (p *Parser) consume(llk *LLk, st *semantic.Statement, s semantic.Symbol) (bool, error) {
	for _, clause := range (*p.grammar)[s] {
		if len(clause.Elements) == 0 {
			return true, nil
		}
		elem := clause.Elements[0]
		if elem.isSymbol {
			return false, fmt.Errorf("Parser.consume: not left factored grammar in %v", clause)
		}
		if llk.CanAccept(elem.Token()) {
			return p.expect(llk, st, s, clause)
		}
	}
	return false, fmt.Errorf("Parser.consume: could not consume token %s in production %s", llk.Current(), s)
}

// expect given the input, symbol, and clause attemps to satisfy all elements.
func (p *Parser) expect(llk *LLk, st *semantic.Statement, s semantic.Symbol, cls *Clause) (bool, error) {
	if cls.ProcessStart != nil {
		if _, err := cls.ProcessStart(st, s); err != nil {
			return false, err
		}
	}
	for _, elem := range cls.Elements {
		tkn := llk.Current()
		if elem.isSymbol {
			if b, err := p.consume(llk, st, elem.Symbol()); !b || err != nil {
				return false, fmt.Errorf("Parser.parse: Failed to consume symbol %v, with error %v", elem.Symbol(), err)
			}
		} else {
			if !llk.Consume(elem.Token()) {
				return false, fmt.Errorf("Parser.parse: Failed to consume %s, got %s instead", elem.Token(), llk.Current().Type)
			}
		}
		if cls.ProcessedElement != nil {
			var ce semantic.ConsumedElement
			if elem.isSymbol {
				ce = semantic.NewConsumedSymbol(ce.Symbol())
			} else {
				ce = semantic.NewConsumedToken(tkn)
			}
			if _, err := cls.ProcessedElement(st, ce); err != nil {
				return false, err
			}
		}
	}
	if cls.ProcessEnd != nil {
		if _, err := cls.ProcessEnd(st, s); err != nil {
			return false, err
		}
	}
	return true, nil
}
<commit_msg>Fix parser semantic symbol creation for element hook<commit_after>// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package grammar

import (
	"fmt"

	"github.com/google/badwolf/bql/lexer"
	"github.com/google/badwolf/bql/semantic"
)

// Element are the main components that define a derivation rule.
type Element struct {
	isSymbol  bool
	symbol    semantic.Symbol
	tokenType lexer.TokenType
}

// NewSymbol creates a new element from a symbol.
func NewSymbol(s semantic.Symbol) Element {
	return Element{
		isSymbol: true,
		symbol:   s,
	}
}

// NewTokenType creates a new element from a token.
func NewTokenType(t lexer.TokenType) Element {
	return Element{
		isSymbol:  false,
		tokenType: t,
	}
}

// Symbol returns the symbol box for the given element.
func (e Element) Symbol() semantic.Symbol {
	return e.symbol
}

// Token returns the value of the token box for the given element.
func (e Element) Token() lexer.TokenType {
	return e.tokenType
}

// Clause contains on clause of the derivation rule.
type Clause struct {
	Elements         []Element
	ProcessStart     semantic.ClauseHook
	ProcessEnd       semantic.ClauseHook
	ProcessedElement semantic.ElementHook
}

// Grammar contains the left factory LLk grammar to be parsed. All provided
// grammars *must* have the "START" symbol to initialte the parsing of input
// text.
type Grammar map[semantic.Symbol][]*Clause

// Parser implements a LLk recursive decend parser for left factorized grammars.
type Parser struct {
	grammar *Grammar
}

// NewParser creates a new recursive decend parser for a left factorized
// grammar.
func NewParser(grammar *Grammar) (*Parser, error) {
	// Check that the grammar is left factorized.
	for _, clauses := range *grammar {
		idx := 0
		for _, cls := range clauses {
			if len(cls.Elements) == 0 {
				if idx == 0 {
					idx++
					continue
				}
				return nil, fmt.Errorf("grammar.NewParser: invalid extra empty clause derivation %v", clauses)
			}
			if cls.Elements[0].isSymbol {
				return nil, fmt.Errorf("grammar.NewParser: not left factored grammar in %v", clauses)
			}
		}
	}
	return &Parser{
		grammar: grammar,
	}, nil
}

// Parse attempts to run the parser for the given input.
func (p *Parser) Parse(llk *LLk, st *semantic.Statement) error {
	b, err := p.consume(llk, st, "START")
	if err != nil {
		return err
	}
	if !b {
		return fmt.Errorf("Parser.Parse: inconsitent parser, no error found, and no tokens were consumed")
	}
	return nil
}

// consume attempts to consume all input tokens for the provided symbols given
// the parser grammar.
func (p *Parser) consume(llk *LLk, st *semantic.Statement, s semantic.Symbol) (bool, error) {
	for _, clause := range (*p.grammar)[s] {
		if len(clause.Elements) == 0 {
			return true, nil
		}
		elem := clause.Elements[0]
		if elem.isSymbol {
			return false, fmt.Errorf("Parser.consume: not left factored grammar in %v", clause)
		}
		if llk.CanAccept(elem.Token()) {
			return p.expect(llk, st, s, clause)
		}
	}
	return false, fmt.Errorf("Parser.consume: could not consume token %s in production %s", llk.Current(), s)
}

// expect given the input, symbol, and clause attemps to satisfy all elements.
func (p *Parser) expect(llk *LLk, st *semantic.Statement, s semantic.Symbol, cls *Clause) (bool, error) {
	if cls.ProcessStart != nil {
		if _, err := cls.ProcessStart(st, s); err != nil {
			return false, err
		}
	}
	for _, elem := range cls.Elements {
		tkn := llk.Current()
		if elem.isSymbol {
			if b, err := p.consume(llk, st, elem.Symbol()); !b || err != nil {
				return false, fmt.Errorf("Parser.parse: Failed to consume symbol %v, with error %v", elem.Symbol(), err)
			}
		} else {
			if !llk.Consume(elem.Token()) {
				return false, fmt.Errorf("Parser.parse: Failed to consume %s, got %s instead", elem.Token(), llk.Current().Type)
			}
		}
		if cls.ProcessedElement != nil {
			var ce semantic.ConsumedElement
			if elem.isSymbol {
				ce = semantic.NewConsumedSymbol(elem.Symbol())
			} else {
				ce = semantic.NewConsumedToken(tkn)
			}
			if _, err := cls.ProcessedElement(st, ce); err != nil {
				return false, err
			}
		}
	}
	if cls.ProcessEnd != nil {
		if _, err := cls.ProcessEnd(st, s); err != nil {
			return false, err
		}
	}
	return true, nil
}
<|endoftext|>


<commit_before>// Copyright 2022 The LUCI 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 cancel contains code in charge of cancelling stale tryjobs.
//
// Cancellator responds to tasks scheduled when a new patch is uploaded,
// looking for and cancelling stale tryjobs.
package cancel

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/tryjob"
)

const CancelStaleTaskClass = "cancel-stale-tryjobs"

// Cancellator is patterned after Updater to support multiple tryjob backends.
type Cancellator struct {
	tqd *tq.Dispatcher

	// guards backends map.
	rwmutex  sync.RWMutex
	backends map[string]cancellatorBackend
}

func NewCancellator(tqd *tq.Dispatcher) *Cancellator {
	c := &Cancellator{
		tqd:      tqd,
		backends: make(map[string]cancellatorBackend),
	}
	c.tqd.RegisterTaskClass(tq.TaskClass{
		ID:        CancelStaleTaskClass,
		Prototype: &tryjob.CancelStaleTryjobsTask{},
		Queue:     "cancel-stale-tryjobs",
		Kind:      tq.Transactional,
		Handler: func(ctx context.Context, payload proto.Message) error {
			return common.TQifyError(ctx, c.handleTask(ctx, payload.(*tryjob.CancelStaleTryjobsTask)))
		},
	})
	return c
}

// RegisterBackend registers a backend.
//
// Panics if backend for the same kind is already registered.
func (c *Cancellator) RegisterBackend(b cancellatorBackend) {
	kind := b.Kind()
	if strings.ContainsRune(kind, '/') {
		panic(fmt.Errorf("backend %T of kind %q must not contain '/'", b, kind))
	}
	c.rwmutex.Lock()
	defer c.rwmutex.Unlock()
	if _, exists := c.backends[kind]; exists {
		panic(fmt.Errorf("backend %q is already registered", kind))
	}
	c.backends[kind] = b
}

func (c *Cancellator) handleTask(ctx context.Context, task *tryjob.CancelStaleTryjobsTask) error {
	if task.PreviousMinEquivPatchset >= task.CurrentMinEquivPatchset {
		panic(fmt.Errorf("patchset numbers expected to increase monotonically"))
	}
	atOrAfter := tryjob.MakeCLPatchset(common.CLID(task.GetClid()), task.GetPreviousMinEquivPatchset())
	before := tryjob.MakeCLPatchset(common.CLID(task.GetClid()), task.GetCurrentMinEquivPatchset())
	q := datastore.NewQuery(tryjob.TryjobKind).Gte("CLPatchsets", atOrAfter).Lt("CLPatchsets", before)
	// The task should be safe to re-run even if it makes partial progress.
	// If the cancellation succeeded but the entity was not successfully
	// updated, on the next run the cancellation should be a no-op and the
	// entity update will be tried again.
	// Any tryjobs that are successfully cancelled and their entities updated
	// will not appear in subsequent runs of the query, ensuring progress.
	return datastore.Run(ctx, q, func(tj *tryjob.Tryjob) error {
		if tj.IsEnded() {
			return nil
		}
		switch ended, err := c.allWatchingRunsEnded(ctx, tj); {
		case err != nil:
			return err
		// External ID is required to cancel tryjob.
		// It may be unset for a few reasons:
		//   - The task triggering the tryjob is stuck.
		//   - The transaction saving the external id got rolled back and is
		//     waiting to retry.
		//   - RM has a bug.
		case ended && tj.ExternalID != "":
			be, err := c.backendFor(tj)
			if err != nil {
				return err
			}
			err = be.CancelTryjob(ctx, tj)
			if err != nil {
				return err
			}
			return datastore.RunInTransaction(ctx, func(ctx context.Context) error {
				if err := datastore.Get(ctx, tj); err != nil {
					return err
				}
				if tj.IsEnded() {
					return nil
				}
				tj.Status = tryjob.Status_CANCELLED
				tj.EVersion++
				tj.EntityUpdateTime = datastore.RoundTime(clock.Now(ctx).UTC())
				return datastore.Put(ctx, tj)
			}, nil)
		}
		return nil
	})
}

// cancellatorBackend is implemented by tryjobs backends, e.g. buildbucket.
type cancellatorBackend interface {
	// Kind identifies the backend
	//
	// It's also the first part of the Tryjob's ExternalID, e.g. "buildbucket".
	// Must not contain a slash.
	Kind() string
	// CancelTryjob should cancel the tryjob given.
	//
	// MUST not modify the given Tryjob object.
	// If the tryjob was already cancelled, it should not return an error.
	CancelTryjob(ctx context.Context, tj *tryjob.Tryjob) error
}

// allWatchingRunsEnded checks if all of the runs that are watching the given
// tryjob have ended.
func (c *Cancellator) allWatchingRunsEnded(ctx context.Context, tryjob *tryjob.Tryjob) (bool, error) {
	runIds := tryjob.AllWatchingRuns()
	runs, errs := run.LoadRunsFromIDs(runIds...).Do(ctx)
	for i, r := range runs {
		switch {
		case errs[i] == datastore.ErrNoSuchEntity:
			return false, errors.Reason("Tryjob %s is associated with a non-existent Run %s", tryjob.ExternalID, runIds[i]).Err()
		case errs[i] != nil:
			return false, errors.Annotate(errs[i], "failed to load run %s", runIds[i]).Err()
		case !run.IsEnded(r.Status):
			return false, nil
		}
	}
	return true, nil
}

func (c *Cancellator) backendFor(t *tryjob.Tryjob) (cancellatorBackend, error) {
	kind, err := t.ExternalID.Kind()
	if err != nil {
		return nil, err
	}
	c.rwmutex.RLock()
	defer c.rwmutex.RUnlock()
	if b, exists := c.backends[kind]; exists {
		return b, nil
	}
	return nil, errors.Reason("%q backend is not supported", kind).Err()
}
<commit_msg>[cv] [stale-tj] Parallelize cancellation calls.<commit_after>// Copyright 2022 The LUCI 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 cancel contains code in charge of cancelling stale tryjobs.
//
// Cancellator responds to tasks scheduled when a new patch is uploaded,
// looking for and cancelling stale tryjobs.
package cancel

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"google.golang.org/protobuf/proto"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/sync/parallel"
	"go.chromium.org/luci/gae/service/datastore"
	"go.chromium.org/luci/server/tq"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/run"
	"go.chromium.org/luci/cv/internal/tryjob"
)

const CancelStaleTaskClass = "cancel-stale-tryjobs"

// Cancellator is patterned after Updater to support multiple tryjob backends.
type Cancellator struct {
	tqd *tq.Dispatcher

	// guards backends map.
	rwmutex  sync.RWMutex
	backends map[string]cancellatorBackend
}

func NewCancellator(tqd *tq.Dispatcher) *Cancellator {
	c := &Cancellator{
		tqd:      tqd,
		backends: make(map[string]cancellatorBackend),
	}
	c.tqd.RegisterTaskClass(tq.TaskClass{
		ID:        CancelStaleTaskClass,
		Prototype: &tryjob.CancelStaleTryjobsTask{},
		Queue:     "cancel-stale-tryjobs",
		Kind:      tq.Transactional,
		Handler: func(ctx context.Context, payload proto.Message) error {
			return common.TQifyError(ctx, c.handleTask(ctx, payload.(*tryjob.CancelStaleTryjobsTask)))
		},
	})
	return c
}

// RegisterBackend registers a backend.
//
// Panics if backend for the same kind is already registered.
func (c *Cancellator) RegisterBackend(b cancellatorBackend) {
	kind := b.Kind()
	if strings.ContainsRune(kind, '/') {
		panic(fmt.Errorf("backend %T of kind %q must not contain '/'", b, kind))
	}
	c.rwmutex.Lock()
	defer c.rwmutex.Unlock()
	if _, exists := c.backends[kind]; exists {
		panic(fmt.Errorf("backend %q is already registered", kind))
	}
	c.backends[kind] = b
}

func (c *Cancellator) handleTask(ctx context.Context, task *tryjob.CancelStaleTryjobsTask) error {
	if task.PreviousMinEquivPatchset >= task.CurrentMinEquivPatchset {
		panic(fmt.Errorf("patchset numbers expected to increase monotonically"))
	}
	atOrAfter := tryjob.MakeCLPatchset(common.CLID(task.GetClid()), task.GetPreviousMinEquivPatchset())
	before := tryjob.MakeCLPatchset(common.CLID(task.GetClid()), task.GetCurrentMinEquivPatchset())
	q := datastore.NewQuery(tryjob.TryjobKind).Gte("CLPatchsets", atOrAfter).Lt("CLPatchsets", before)

	toCancel := make([]*tryjob.Tryjob, 0)
	// The task should be safe to re-run even if it makes partial progress.
	// If the cancellation succeeded but the entity was not successfully
	// updated, on the next run the cancellation should be a no-op and the
	// entity update will be tried again.
	// Any tryjobs that are successfully cancelled and their entities updated
	// will not appear in subsequent runs of the query, ensuring progress.
	dsErr := datastore.Run(ctx, q, func(tj *tryjob.Tryjob) error {
		if tj.IsEnded() {
			return nil
		}
		switch ended, err := c.allWatchingRunsEnded(ctx, tj); {
		case err != nil:
			return err
		// External ID is required to cancel tryjob.
		// It may be unset for a few reasons:
		//   - The task triggering the tryjob is stuck.
		//   - The transaction saving the external id got rolled back and is
		//     waiting to retry.
		//   - RM has a bug.
		case ended && tj.ExternalID != "":
			toCancel = append(toCancel, tj)
		}
		return nil
	})
	switch {
	case len(toCancel) == 0:
		return dsErr
	case dsErr != nil:
		logging.Warningf(ctx, "the query to fetch stale tryjobs returned an error along with partial results, will continue to cancel the returned tryjobs first")
		fallthrough
	default:
		if err := c.cancelAll(ctx, toCancel); err != nil {
			return err
		}
		return dsErr
	}
}

func (c *Cancellator) cancelAll(ctx context.Context, tjs []*tryjob.Tryjob) error {
	if len(tjs) == 0 {
		return nil
	}
	errs := parallel.WorkPool(min(8, len(tjs)), func(work chan<- func() error) {
		for _, tj := range tjs {
			tj := tj
			work <- func() error {
				be, err := c.backendFor(tj)
				if err != nil {
					return err
				}
				// TODO(crbug/1308930): use Buildbucket's batch API to reduce
				// number of RPCs.
				err = be.CancelTryjob(ctx, tj)
				if err != nil {
					return err
				}
				return datastore.RunInTransaction(ctx, func(ctx context.Context) error {
					if err := datastore.Get(ctx, tj); err != nil {
						return err
					}
					if tj.IsEnded() {
						return nil
					}
					tj.Status = tryjob.Status_CANCELLED
					tj.EVersion++
					tj.EntityUpdateTime = datastore.RoundTime(clock.Now(ctx).UTC())
					return datastore.Put(ctx, tj)
				}, nil)
			}
		}
	})
	return common.MostSevereError(errs)
}

// cancellatorBackend is implemented by tryjobs backends, e.g. buildbucket.
type cancellatorBackend interface {
	// Kind identifies the backend
	//
	// It's also the first part of the Tryjob's ExternalID, e.g. "buildbucket".
	// Must not contain a slash.
	Kind() string
	// CancelTryjob should cancel the tryjob given.
	//
	// MUST not modify the given Tryjob object.
	// If the tryjob was already cancelled, it should not return an error.
	CancelTryjob(ctx context.Context, tj *tryjob.Tryjob) error
}

// allWatchingRunsEnded checks if all of the runs that are watching the given
// tryjob have ended.
func (c *Cancellator) allWatchingRunsEnded(ctx context.Context, tryjob *tryjob.Tryjob) (bool, error) {
	runIds := tryjob.AllWatchingRuns()
	runs, errs := run.LoadRunsFromIDs(runIds...).Do(ctx)
	for i, r := range runs {
		switch {
		case errs[i] == datastore.ErrNoSuchEntity:
			return false, errors.Reason("Tryjob %s is associated with a non-existent Run %s", tryjob.ExternalID, runIds[i]).Err()
		case errs[i] != nil:
			return false, errors.Annotate(errs[i], "failed to load run %s", runIds[i]).Err()
		case !run.IsEnded(r.Status):
			return false, nil
		}
	}
	return true, nil
}

func (c *Cancellator) backendFor(t *tryjob.Tryjob) (cancellatorBackend, error) {
	kind, err := t.ExternalID.Kind()
	if err != nil {
		return nil, err
	}
	c.rwmutex.RLock()
	defer c.rwmutex.RUnlock()
	if b, exists := c.backends[kind]; exists {
		return b, nil
	}
	return nil, errors.Reason("%q backend is not supported", kind).Err()
}

func min(x, y int) int {
	if x <= y {
		return x
	}
	return y
}
<|endoftext|>