<commit_before>/*
 * protector.go - functions for dealing with policies
 *
 * Copyright 2017 Google Inc.
 * Author: Joe Richey (joerichey@google.com)
 *
 * 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 actions

import (
	"fmt"
	"log"
	"reflect"

	"github.com/pkg/errors"

	"github.com/google/fscrypt/crypto"
	"github.com/google/fscrypt/filesystem"
	"github.com/google/fscrypt/metadata"
	"github.com/google/fscrypt/util"
)

// Errors relating to Policies
var (
	ErrMissingPolicyMetadata  = util.SystemError("missing policy metadata for encrypted directory")
	ErrPolicyMetadataMismatch = util.SystemError("inconsistent metadata between filesystem and directory")
	ErrDifferentFilesystem    = errors.New("policies may only protect files on the same filesystem")
	ErrOnlyProtector          = errors.New("cannot remove the only protector for a policy")
	ErrAlreadyProtected       = errors.New("policy already protected by protector")
	ErrNotProtected           = errors.New("policy not protected by protector")
)

// PurgeAllPolicies removes all policy keys on the filesystem from the kernel
// keyring. In order for this removal to have an effect, the filesystem should
// also be unmounted.
func PurgeAllPolicies(ctx *Context) error {
	if err := ctx.checkContext(); err != nil {
		return err
	}
	policies, err := ctx.Mount.ListPolicies()
	if err != nil {
		return err
	}

	for _, policyDescriptor := range policies {
		service := ctx.getService()
		err = crypto.RemovePolicyKey(policyDescriptor, service)

		switch errors.Cause(err) {
		case nil, crypto.ErrKeyringSearch:
			// We don't care if the key has already been removed
			break
		default:
			return err
		}
	}
	return nil
}

// Policy represents an unlocked policy, so it contains the PolicyData as well
// as the actual protector key. These unlocked Polices can then be applied to a
// directory, or have their key material inserted into the keyring (which will
// allow encrypted files to be accessed). As with the key struct, a Policy
// should be wiped after use.
type Policy struct {
	Context *Context
	data    *metadata.PolicyData
	key     *crypto.Key
	created bool
}

// CreatePolicy creates a Policy protected by given Protector and stores the
// appropriate data on the filesystem. On error, no data is changed on the
// filesystem.
func CreatePolicy(ctx *Context, protector *Protector) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}
	// Randomly create the underlying policy key (and wipe if we fail)
	key, err := crypto.NewRandomKey(metadata.PolicyKeyLen)
	if err != nil {
		return nil, err
	}

	policy := &Policy{
		Context: ctx,
		data: &metadata.PolicyData{
			Options:       ctx.Config.Options,
			KeyDescriptor: crypto.ComputeDescriptor(key),
		},
		key:     key,
		created: true,
	}

	if err = policy.AddProtector(protector); err != nil {
		policy.Lock()
		return nil, err
	}

	return policy, nil
}

// GetPolicy retrieves a locked policy with a specific descriptor. The Policy is
// still locked in this case, so it must be unlocked before using certain
// methods.
func GetPolicy(ctx *Context, descriptor string) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}
	data, err := ctx.Mount.GetPolicy(descriptor)
	if err != nil {
		return nil, err
	}
	log.Printf("got data for %s from %q", descriptor, ctx.Mount.Path)

	return &Policy{Context: ctx, data: data}, nil
}

// GetPolicyFromPath returns the locked policy descriptor for a file on the
// filesystem. The Policy is still locked in this case, so it must be unlocked
// before using certain methods. An error is returned if the metadata is
// inconsistent or the path is not encrypted.
func GetPolicyFromPath(ctx *Context, path string) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}

	// We double check that the options agree for both the data we get from
	// the path, and the data we get from the mountpoint.
	pathData, err := metadata.GetPolicy(path)
	if err != nil {
		return nil, err
	}
	descriptor := pathData.KeyDescriptor
	log.Printf("found policy %s for %q", descriptor, path)

	mountData, err := ctx.Mount.GetPolicy(descriptor)
	if err != nil {
		log.Printf("getting policy metadata: %v", err)
		return nil, errors.Wrap(ErrMissingPolicyMetadata, path)
	}
	log.Printf("found data for policy %s on %q", descriptor, ctx.Mount.Path)

	if !reflect.DeepEqual(pathData.Options, mountData.Options) {
		log.Printf("options from path: %+v", pathData.Options)
		log.Printf("options from mount: %+v", mountData.Options)
		return nil, errors.Wrapf(ErrPolicyMetadataMismatch, "policy %s", descriptor)
	}
	log.Print("data from filesystem and path agree")

	return &Policy{Context: ctx, data: mountData}, nil
}

// ProtectorOptions creates a slice of ProtectorOptions for the protectors
// protecting this policy.
func (policy *Policy) ProtectorOptions() []*ProtectorOption {
	options := make([]*ProtectorOption, len(policy.data.WrappedPolicyKeys))
	for i, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		options[i] = policy.Context.getProtectorOption(wrappedPolicyKey.ProtectorDescriptor)
	}
	return options
}

// ProtectorDescriptors creates a slice of the Protector descriptors for the
// protectors protecting this policy.
func (policy *Policy) ProtectorDescriptors() []string {
	descriptors := make([]string, len(policy.data.WrappedPolicyKeys))
	for i, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		descriptors[i] = wrappedPolicyKey.ProtectorDescriptor
	}
	return descriptors
}

// Descriptor returns the key descriptor for this policy.
func (policy *Policy) Descriptor() string {
	return policy.data.KeyDescriptor
}

// Destroy removes a policy from the filesystem. The internal key should still
// be wiped with Lock().
func (policy *Policy) Destroy() error {
	return policy.Context.Mount.RemovePolicy(policy.Descriptor())
}

// Revert destroys a policy if it was created, but does nothing if it was just
// queried from the filesystem.
func (policy *Policy) Revert() error {
	if !policy.created {
		return nil
	}
	return policy.Destroy()
}

func (policy *Policy) String() string {
	return fmt.Sprintf("Policy: %s\nMountpoint: %s\nOptions: %v\nProtectors:%+v",
		policy.Descriptor(), policy.Context.Mount, policy.data.Options,
		policy.ProtectorDescriptors())
}

// Unlock unwraps the Policy's internal key. As a Protector is needed to unlock
// the Policy, callbacks to select the Policy and get the key are needed. This
// method will retry the keyFn as necessary to get the correct key for the
// selected protector. Does nothing if policy is already unlocked.
func (policy *Policy) Unlock(optionFn OptionFunc, keyFn KeyFunc) error {
	if policy.key != nil {
		return nil
	}
	options := policy.ProtectorOptions()

	// The OptionFunc indicates which option and wrapped key we should use.
	idx, err := optionFn(policy.Descriptor(), options)
	if err != nil {
		return err
	}
	option := options[idx]
	if option.LoadError != nil {
		return option.LoadError
	}

	log.Printf("protector %s selected in callback", option.Descriptor())
	protectorKey, err := unwrapProtectorKey(option.ProtectorInfo, keyFn)
	if err != nil {
		return err
	}
	defer protectorKey.Wipe()

	log.Printf("unwrapping policy %s with protector", policy.Descriptor())
	wrappedPolicyKey := policy.data.WrappedPolicyKeys[idx].WrappedKey
	policy.key, err = crypto.Unwrap(protectorKey, wrappedPolicyKey)
	return err
}

// Lock wipes a Policy's internal Key. It should always be called after using a
// Policy. This is often done with a defer statement. There is no effect if
// called multiple times.
func (policy *Policy) Lock() error {
	err := policy.key.Wipe()
	policy.key = nil
	return err
}

// AddProtector updates the data that is wrapping the Policy Key so that the
// provided Protector is now protecting the specified Policy. If an error is
// returned, no data has been changed. If the policy and protector are on
// different filesystems, a link will be created between them. The policy and
// protector must both be unlocked.
func (policy *Policy) AddProtector(protector *Protector) error {
	if _, ok := policy.findWrappedKeyIndex(protector.Descriptor()); ok {
		return ErrAlreadyProtected
	}
	if policy.key == nil || protector.key == nil {
		return ErrLocked
	}

	// If the protector is on a different filesystem, we need to add a link
	// to it on the policy's filesystem.
	if policy.Context.Mount != protector.Context.Mount {
		log.Printf("policy on %s\n protector on %s\n", policy.Context.Mount, protector.Context.Mount)
		err := policy.Context.Mount.AddLinkedProtector(
			protector.Descriptor(), protector.Context.Mount)
		if err != nil {
			return err
		}
	} else {
		log.Printf("policy and protector both on %q", policy.Context.Mount)
	}

	// Create the wrapped policy key
	wrappedKey, err := crypto.Wrap(protector.key, policy.key)
	if err != nil {
		return err
	}

	// Append the wrapped key to the data
	policy.addKey(&metadata.WrappedPolicyKey{
		ProtectorDescriptor: protector.Descriptor(),
		WrappedKey:          wrappedKey,
	})

	if err := policy.commitData(); err != nil {
		// revert the addition on failure
		policy.removeKey(len(policy.data.WrappedPolicyKeys) - 1)
		return err
	}
	return nil
}

// RemoveProtector updates the data that is wrapping the Policy Key so that the
// provided Protector is no longer protecting the specified Policy. If an error
// is returned, no data has been changed. Note that no protector links are
// removed (in the case where the protector and policy are on different
// filesystems). The policy and protector can be locked or unlocked.
func (policy *Policy) RemoveProtector(protector *Protector) error {
	idx, ok := policy.findWrappedKeyIndex(protector.Descriptor())
	if !ok {
		return ErrNotProtected
	}

	if len(policy.data.WrappedPolicyKeys) == 1 {
		return ErrOnlyProtector
	}

	// Remove the wrapped key from the data
	toRemove := policy.removeKey(idx)

	if err := policy.commitData(); err != nil {
		// revert the removal on failure (order is irrelevant)
		policy.addKey(toRemove)
		return err
	}
	return nil
}

// Apply sets the Policy on a specified directory. Currently we impose the
// additional constraint that policies and the directories they are applied to
// must reside on the same filesystem.
func (policy *Policy) Apply(path string) error {
	if pathMount, err := filesystem.FindMount(path); err != nil {
		return err
	} else if pathMount != policy.Context.Mount {
		return ErrDifferentFilesystem
	}

	return metadata.SetPolicy(path, policy.data)
}

// IsProvisioned returns a boolean indicating if the policy has its key in the
// keyring, meaning files and directories using this policy are accessible.
func (policy *Policy) IsProvisioned() bool {
	_, _, err := crypto.FindPolicyKey(policy.Descriptor(), policy.Context.getService())
	return err == nil
}

// Provision inserts the Policy key into the kernel keyring. This allows reading
// and writing of files encrypted with this directory. Requires unlocked Policy.
func (policy *Policy) Provision() error {
	if policy.key == nil {
		return ErrLocked
	}
	return crypto.InsertPolicyKey(policy.key, policy.Descriptor(), policy.Context.getService())
}

// Deprovision removes the Policy key from the kernel keyring. This prevents
// reading and writing to the directory once the caches are cleared.
func (policy *Policy) Deprovision() error {
	return crypto.RemovePolicyKey(policy.Descriptor(), policy.Context.getService())
}

// commitData writes the Policy's current data to the filesystem.
func (policy *Policy) commitData() error {
	return policy.Context.Mount.AddPolicy(policy.data)
}

// findWrappedPolicyKey returns the index of the wrapped policy key
// corresponding to this policy and protector. The returned bool is false if no
// wrapped policy key corresponds to the specified protector, true otherwise.
func (policy *Policy) findWrappedKeyIndex(protectorDescriptor string) (int, bool) {
	for idx, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		if wrappedPolicyKey.ProtectorDescriptor == protectorDescriptor {
			return idx, true
		}
	}
	return 0, false
}

// addKey adds the wrapped policy key to end of the wrapped key data.
func (policy *Policy) addKey(toAdd *metadata.WrappedPolicyKey) {
	policy.data.WrappedPolicyKeys = append(policy.data.WrappedPolicyKeys, toAdd)
}

// remove removes the wrapped policy key at the specified index. This
// does not preserve the order of the wrapped policy key array. If no index is
// specified the last key is removed.
func (policy *Policy) removeKey(index int) *metadata.WrappedPolicyKey {
	lastIdx := len(policy.data.WrappedPolicyKeys) - 1
	toRemove := policy.data.WrappedPolicyKeys[index]

	// See https://github.com/golang/go/wiki/SliceTricks
	policy.data.WrappedPolicyKeys[index] = policy.data.WrappedPolicyKeys[lastIdx]
	policy.data.WrappedPolicyKeys[lastIdx] = nil
	policy.data.WrappedPolicyKeys = policy.data.WrappedPolicyKeys[:lastIdx]

	return toRemove
}
<commit_msg>actions: Policies now have Description method<commit_after>/*
 * protector.go - functions for dealing with policies
 *
 * Copyright 2017 Google Inc.
 * Author: Joe Richey (joerichey@google.com)
 *
 * 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 actions

import (
	"fmt"
	"log"
	"reflect"

	"github.com/pkg/errors"

	"github.com/google/fscrypt/crypto"
	"github.com/google/fscrypt/filesystem"
	"github.com/google/fscrypt/metadata"
	"github.com/google/fscrypt/util"
)

// Errors relating to Policies
var (
	ErrMissingPolicyMetadata  = util.SystemError("missing policy metadata for encrypted directory")
	ErrPolicyMetadataMismatch = util.SystemError("inconsistent metadata between filesystem and directory")
	ErrDifferentFilesystem    = errors.New("policies may only protect files on the same filesystem")
	ErrOnlyProtector          = errors.New("cannot remove the only protector for a policy")
	ErrAlreadyProtected       = errors.New("policy already protected by protector")
	ErrNotProtected           = errors.New("policy not protected by protector")
)

// PurgeAllPolicies removes all policy keys on the filesystem from the kernel
// keyring. In order for this removal to have an effect, the filesystem should
// also be unmounted.
func PurgeAllPolicies(ctx *Context) error {
	if err := ctx.checkContext(); err != nil {
		return err
	}
	policies, err := ctx.Mount.ListPolicies()
	if err != nil {
		return err
	}

	for _, policyDescriptor := range policies {
		service := ctx.getService()
		err = crypto.RemovePolicyKey(service + policyDescriptor)

		switch errors.Cause(err) {
		case nil, crypto.ErrKeyringSearch:
			// We don't care if the key has already been removed
			break
		default:
			return err
		}
	}
	return nil
}

// Policy represents an unlocked policy, so it contains the PolicyData as well
// as the actual protector key. These unlocked Polices can then be applied to a
// directory, or have their key material inserted into the keyring (which will
// allow encrypted files to be accessed). As with the key struct, a Policy
// should be wiped after use.
type Policy struct {
	Context *Context
	data    *metadata.PolicyData
	key     *crypto.Key
	created bool
}

// CreatePolicy creates a Policy protected by given Protector and stores the
// appropriate data on the filesystem. On error, no data is changed on the
// filesystem.
func CreatePolicy(ctx *Context, protector *Protector) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}
	// Randomly create the underlying policy key (and wipe if we fail)
	key, err := crypto.NewRandomKey(metadata.PolicyKeyLen)
	if err != nil {
		return nil, err
	}

	policy := &Policy{
		Context: ctx,
		data: &metadata.PolicyData{
			Options:       ctx.Config.Options,
			KeyDescriptor: crypto.ComputeDescriptor(key),
		},
		key:     key,
		created: true,
	}

	if err = policy.AddProtector(protector); err != nil {
		policy.Lock()
		return nil, err
	}

	return policy, nil
}

// GetPolicy retrieves a locked policy with a specific descriptor. The Policy is
// still locked in this case, so it must be unlocked before using certain
// methods.
func GetPolicy(ctx *Context, descriptor string) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}
	data, err := ctx.Mount.GetPolicy(descriptor)
	if err != nil {
		return nil, err
	}
	log.Printf("got data for %s from %q", descriptor, ctx.Mount.Path)

	return &Policy{Context: ctx, data: data}, nil
}

// GetPolicyFromPath returns the locked policy descriptor for a file on the
// filesystem. The Policy is still locked in this case, so it must be unlocked
// before using certain methods. An error is returned if the metadata is
// inconsistent or the path is not encrypted.
func GetPolicyFromPath(ctx *Context, path string) (*Policy, error) {
	if err := ctx.checkContext(); err != nil {
		return nil, err
	}

	// We double check that the options agree for both the data we get from
	// the path, and the data we get from the mountpoint.
	pathData, err := metadata.GetPolicy(path)
	if err != nil {
		return nil, err
	}
	descriptor := pathData.KeyDescriptor
	log.Printf("found policy %s for %q", descriptor, path)

	mountData, err := ctx.Mount.GetPolicy(descriptor)
	if err != nil {
		log.Printf("getting policy metadata: %v", err)
		return nil, errors.Wrap(ErrMissingPolicyMetadata, path)
	}
	log.Printf("found data for policy %s on %q", descriptor, ctx.Mount.Path)

	if !reflect.DeepEqual(pathData.Options, mountData.Options) {
		log.Printf("options from path: %+v", pathData.Options)
		log.Printf("options from mount: %+v", mountData.Options)
		return nil, errors.Wrapf(ErrPolicyMetadataMismatch, "policy %s", descriptor)
	}
	log.Print("data from filesystem and path agree")

	return &Policy{Context: ctx, data: mountData}, nil
}

// ProtectorOptions creates a slice of ProtectorOptions for the protectors
// protecting this policy.
func (policy *Policy) ProtectorOptions() []*ProtectorOption {
	options := make([]*ProtectorOption, len(policy.data.WrappedPolicyKeys))
	for i, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		options[i] = policy.Context.getProtectorOption(wrappedPolicyKey.ProtectorDescriptor)
	}
	return options
}

// ProtectorDescriptors creates a slice of the Protector descriptors for the
// protectors protecting this policy.
func (policy *Policy) ProtectorDescriptors() []string {
	descriptors := make([]string, len(policy.data.WrappedPolicyKeys))
	for i, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		descriptors[i] = wrappedPolicyKey.ProtectorDescriptor
	}
	return descriptors
}

// Descriptor returns the key descriptor for this policy.
func (policy *Policy) Descriptor() string {
	return policy.data.KeyDescriptor
}

// Description returns the description that will be used when the key for this
// Policy is inserted into the keyring
func (policy *Policy) Description() string {
	return policy.Context.getService() + policy.Descriptor()
}

// Destroy removes a policy from the filesystem. The internal key should still
// be wiped with Lock().
func (policy *Policy) Destroy() error {
	return policy.Context.Mount.RemovePolicy(policy.Descriptor())
}

// Revert destroys a policy if it was created, but does nothing if it was just
// queried from the filesystem.
func (policy *Policy) Revert() error {
	if !policy.created {
		return nil
	}
	return policy.Destroy()
}

func (policy *Policy) String() string {
	return fmt.Sprintf("Policy: %s\nMountpoint: %s\nOptions: %v\nProtectors:%+v",
		policy.Descriptor(), policy.Context.Mount, policy.data.Options,
		policy.ProtectorDescriptors())
}

// Unlock unwraps the Policy's internal key. As a Protector is needed to unlock
// the Policy, callbacks to select the Policy and get the key are needed. This
// method will retry the keyFn as necessary to get the correct key for the
// selected protector. Does nothing if policy is already unlocked.
func (policy *Policy) Unlock(optionFn OptionFunc, keyFn KeyFunc) error {
	if policy.key != nil {
		return nil
	}
	options := policy.ProtectorOptions()

	// The OptionFunc indicates which option and wrapped key we should use.
	idx, err := optionFn(policy.Descriptor(), options)
	if err != nil {
		return err
	}
	option := options[idx]
	if option.LoadError != nil {
		return option.LoadError
	}

	log.Printf("protector %s selected in callback", option.Descriptor())
	protectorKey, err := unwrapProtectorKey(option.ProtectorInfo, keyFn)
	if err != nil {
		return err
	}
	defer protectorKey.Wipe()

	log.Printf("unwrapping policy %s with protector", policy.Descriptor())
	wrappedPolicyKey := policy.data.WrappedPolicyKeys[idx].WrappedKey
	policy.key, err = crypto.Unwrap(protectorKey, wrappedPolicyKey)
	return err
}

// Lock wipes a Policy's internal Key. It should always be called after using a
// Policy. This is often done with a defer statement. There is no effect if
// called multiple times.
func (policy *Policy) Lock() error {
	err := policy.key.Wipe()
	policy.key = nil
	return err
}

// AddProtector updates the data that is wrapping the Policy Key so that the
// provided Protector is now protecting the specified Policy. If an error is
// returned, no data has been changed. If the policy and protector are on
// different filesystems, a link will be created between them. The policy and
// protector must both be unlocked.
func (policy *Policy) AddProtector(protector *Protector) error {
	if _, ok := policy.findWrappedKeyIndex(protector.Descriptor()); ok {
		return ErrAlreadyProtected
	}
	if policy.key == nil || protector.key == nil {
		return ErrLocked
	}

	// If the protector is on a different filesystem, we need to add a link
	// to it on the policy's filesystem.
	if policy.Context.Mount != protector.Context.Mount {
		log.Printf("policy on %s\n protector on %s\n", policy.Context.Mount, protector.Context.Mount)
		err := policy.Context.Mount.AddLinkedProtector(
			protector.Descriptor(), protector.Context.Mount)
		if err != nil {
			return err
		}
	} else {
		log.Printf("policy and protector both on %q", policy.Context.Mount)
	}

	// Create the wrapped policy key
	wrappedKey, err := crypto.Wrap(protector.key, policy.key)
	if err != nil {
		return err
	}

	// Append the wrapped key to the data
	policy.addKey(&metadata.WrappedPolicyKey{
		ProtectorDescriptor: protector.Descriptor(),
		WrappedKey:          wrappedKey,
	})

	if err := policy.commitData(); err != nil {
		// revert the addition on failure
		policy.removeKey(len(policy.data.WrappedPolicyKeys) - 1)
		return err
	}
	return nil
}

// RemoveProtector updates the data that is wrapping the Policy Key so that the
// provided Protector is no longer protecting the specified Policy. If an error
// is returned, no data has been changed. Note that no protector links are
// removed (in the case where the protector and policy are on different
// filesystems). The policy and protector can be locked or unlocked.
func (policy *Policy) RemoveProtector(protector *Protector) error {
	idx, ok := policy.findWrappedKeyIndex(protector.Descriptor())
	if !ok {
		return ErrNotProtected
	}

	if len(policy.data.WrappedPolicyKeys) == 1 {
		return ErrOnlyProtector
	}

	// Remove the wrapped key from the data
	toRemove := policy.removeKey(idx)

	if err := policy.commitData(); err != nil {
		// revert the removal on failure (order is irrelevant)
		policy.addKey(toRemove)
		return err
	}
	return nil
}

// Apply sets the Policy on a specified directory. Currently we impose the
// additional constraint that policies and the directories they are applied to
// must reside on the same filesystem.
func (policy *Policy) Apply(path string) error {
	if pathMount, err := filesystem.FindMount(path); err != nil {
		return err
	} else if pathMount != policy.Context.Mount {
		return ErrDifferentFilesystem
	}

	return metadata.SetPolicy(path, policy.data)
}

// IsProvisioned returns a boolean indicating if the policy has its key in the
// keyring, meaning files and directories using this policy are accessible.
func (policy *Policy) IsProvisioned() bool {
	_, _, err := crypto.FindPolicyKey(policy.Description())
	return err == nil
}

// Provision inserts the Policy key into the kernel keyring. This allows reading
// and writing of files encrypted with this directory. Requires unlocked Policy.
func (policy *Policy) Provision() error {
	if policy.key == nil {
		return ErrLocked
	}
	return crypto.InsertPolicyKey(policy.key, policy.Description())
}

// Deprovision removes the Policy key from the kernel keyring. This prevents
// reading and writing to the directory once the caches are cleared.
func (policy *Policy) Deprovision() error {
	return crypto.RemovePolicyKey(policy.Description())
}

// commitData writes the Policy's current data to the filesystem.
func (policy *Policy) commitData() error {
	return policy.Context.Mount.AddPolicy(policy.data)
}

// findWrappedPolicyKey returns the index of the wrapped policy key
// corresponding to this policy and protector. The returned bool is false if no
// wrapped policy key corresponds to the specified protector, true otherwise.
func (policy *Policy) findWrappedKeyIndex(protectorDescriptor string) (int, bool) {
	for idx, wrappedPolicyKey := range policy.data.WrappedPolicyKeys {
		if wrappedPolicyKey.ProtectorDescriptor == protectorDescriptor {
			return idx, true
		}
	}
	return 0, false
}

// addKey adds the wrapped policy key to end of the wrapped key data.
func (policy *Policy) addKey(toAdd *metadata.WrappedPolicyKey) {
	policy.data.WrappedPolicyKeys = append(policy.data.WrappedPolicyKeys, toAdd)
}

// remove removes the wrapped policy key at the specified index. This
// does not preserve the order of the wrapped policy key array. If no index is
// specified the last key is removed.
func (policy *Policy) removeKey(index int) *metadata.WrappedPolicyKey {
	lastIdx := len(policy.data.WrappedPolicyKeys) - 1
	toRemove := policy.data.WrappedPolicyKeys[index]

	// See https://github.com/golang/go/wiki/SliceTricks
	policy.data.WrappedPolicyKeys[index] = policy.data.WrappedPolicyKeys[lastIdx]
	policy.data.WrappedPolicyKeys[lastIdx] = nil
	policy.data.WrappedPolicyKeys = policy.data.WrappedPolicyKeys[:lastIdx]

	return toRemove
}
<|endoftext|>


<commit_before>package pull

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestPing(t *testing.T) {

	ts := httptest.NewServer(pingHandler(http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", ts.URL, nil)
	req.Header.Set("X-GitHub-Event", "ping")
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusAccepted {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusNotImplemented, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	expected := "Received ping"
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestImplements(t *testing.T) {

	ts := httptest.NewServer(implementsHandler(http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", ts.URL, nil)
	req.Header.Set("X-GitHub-Event", "unsupported")
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusNotImplemented {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusNotImplemented, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	expected := "Unsupported event"
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecret(t *testing.T) {

	secret := "test"
	reqbody := []byte(`{"test":"test"}`)
	hash := "sha1=7c7429c26f63ae28d74597e825c20e1796c167e3"

	expected := "Success"
	ts := httptest.NewServer(secretHandler(secret, http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusTeapot)
			fmt.Fprint(w, expected)
		})))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, bytes.NewBuffer(reqbody))
	req.Header.Set("X-Hub-Signature", hash)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusTeapot {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusTeapot, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecretNoHeader(t *testing.T) {

	expected := "X-Hub-Signature required for HMAC verification"

	ts := httptest.NewServer(secretHandler("secret", http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, nil)
	//req.Header.Set("X-Hub-Signature", hash)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusForbidden {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusForbidden, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecretFailed(t *testing.T) {

	secret := "test"
	reqbody := []byte(`{"test":"test"}`)
	hash := "sha1=mismatchgibberish"

	ts := httptest.NewServer(secretHandler(secret, http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, bytes.NewBuffer(reqbody))
	req.Header.Set("X-Hub-Signature", hash)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusForbidden {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusForbidden, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)
	expected := "HMAC verification failed"

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}
<commit_msg>Clean up function.<commit_after>package pull

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestPing(t *testing.T) {

	ts := httptest.NewServer(pingHandler(http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", ts.URL, nil)
	req.Header.Set("X-GitHub-Event", "ping")
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusAccepted {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusNotImplemented, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	expected := "Received ping"
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestImplements(t *testing.T) {

	ts := httptest.NewServer(implementsHandler(http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", ts.URL, nil)
	req.Header.Set("X-GitHub-Event", "unsupported")
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusNotImplemented {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusNotImplemented, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	expected := "Unsupported event"
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecret(t *testing.T) {

	secret := "test"
	reqbody := []byte(`{"test":"test"}`)
	hash := "sha1=7c7429c26f63ae28d74597e825c20e1796c167e3"

	expected := "Success"
	ts := httptest.NewServer(secretHandler(secret, http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusTeapot)
			fmt.Fprint(w, expected)
		})))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, bytes.NewBuffer(reqbody))
	req.Header.Set("X-Hub-Signature", hash)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusTeapot {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusTeapot, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecretNoHeader(t *testing.T) {

	ts := httptest.NewServer(secretHandler("secret", http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, nil)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusForbidden {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusForbidden, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)
	expected := "X-Hub-Signature required for HMAC verification"

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}

func TestSecretFailed(t *testing.T) {

	secret := "test"
	reqbody := []byte(`{"test":"test"}`)
	hash := "sha1=mismatchgibberish"

	ts := httptest.NewServer(secretHandler(secret, http.NotFoundHandler()))
	defer ts.Close()

	client := &http.Client{}

	req, _ := http.NewRequest("GET", ts.URL, bytes.NewBuffer(reqbody))
	req.Header.Set("X-Hub-Signature", hash)
	resp, err := client.Do(req)

	if err != nil {
		t.Fatalf("HTTP Request failed: %s", err)
	}

	if resp.StatusCode != http.StatusForbidden {
		t.Errorf("Expected HTTP %d, got HTTP %d", http.StatusForbidden, resp.StatusCode)
	}

	body, _ := ioutil.ReadAll(resp.Body)
	actual := string(body)
	expected := "HMAC verification failed"

	if actual != expected {
		t.Errorf("Expected %s got %s", expected, actual)
	}
}
<|endoftext|>


<commit_before>package sqalx

import (
	"database/sql"
	"errors"
	"strings"

	"github.com/jmoiron/sqlx"
	uuid "github.com/satori/go.uuid"
)

var (
	// ErrNotInTransaction is returned when using Commit
	// outside of a transaction.
	ErrNotInTransaction = errors.New("not in transaction")

	// ErrIncompatibleOption is returned when using an option incompatible
	// with the selected driver.
	ErrIncompatibleOption = errors.New("incompatible option")
)

// A Node is a database driver that can manage nested transactions.
type Node interface {
	Driver

	// Close the underlying sqlx connection.
	Close() error
	// Begin a new transaction.
	Beginx() (Node, error)
	// Rollback the associated transaction.
	Rollback() error
	// Commit the assiociated transaction.
	Commit() error
	// Tx returns the underlying transaction.
	Tx() *sqlx.Tx
}

// A Driver can query the database. It can either be a *sqlx.DB or a *sqlx.Tx
// and therefore is limited to the methods they have in common.
type Driver interface {
	sqlx.Execer
	sqlx.Queryer
	sqlx.Preparer
	BindNamed(query string, arg interface{}) (string, []interface{}, error)
	DriverName() string
	Get(dest interface{}, query string, args ...interface{}) error
	MustExec(query string, args ...interface{}) sql.Result
	NamedExec(query string, arg interface{}) (sql.Result, error)
	NamedQuery(query string, arg interface{}) (*sqlx.Rows, error)
	PrepareNamed(query string) (*sqlx.NamedStmt, error)
	Preparex(query string) (*sqlx.Stmt, error)
	Rebind(query string) string
	Select(dest interface{}, query string, args ...interface{}) error
}

// New creates a new Node with the given DB.
func New(db *sqlx.DB, options ...Option) (Node, error) {
	n := node{
		db:     db,
		Driver: db,
	}

	for _, opt := range options {
		err := opt(&n)
		if err != nil {
			return nil, err
		}
	}

	return &n, nil
}

// NewFromTransaction creates a new Node from the given transaction.
func NewFromTransaction(tx *sqlx.Tx, options ...Option) (Node, error) {
	n := node{
		tx:     tx,
		Driver: tx,
	}

	for _, opt := range options {
		err := opt(&n)
		if err != nil {
			return nil, err
		}
	}

	return &n, nil
}

// Connect to a database.
func Connect(driverName, dataSourceName string, options ...Option) (Node, error) {
	db, err := sqlx.Connect(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}

	node, err := New(db, options...)
	if err != nil {
		// the connection has been opened within this function, we must close it
		// on error.
		db.Close()
		return nil, err
	}

	return node, nil
}

type node struct {
	Driver
	db               *sqlx.DB
	tx               *sqlx.Tx
	savePointID      string
	savePointEnabled bool
	nested           bool
}

func (n *node) Close() error {
	return n.db.Close()
}

func (n node) Beginx() (Node, error) {
	var err error

	switch {
	case n.tx == nil:
		// new actual transaction
		n.tx, err = n.db.Beginx()
		n.Driver = n.tx
	case n.savePointEnabled:
		// already in a transaction: using savepoints
		n.nested = true
		// savepoints name must start with a char and cannot contain dashes (-)
		n.savePointID = "sp_" + strings.Replace(uuid.NewV1().String(), "-", "_", -1)
		_, err = n.tx.Exec("SAVEPOINT " + n.savePointID)
	default:
		// already in a transaction: reusing current transaction
		n.nested = true
	}

	if err != nil {
		return nil, err
	}

	return &n, nil
}

func (n *node) Rollback() error {
	if n.tx == nil {
		return nil
	}

	var err error

	if n.savePointEnabled && n.savePointID != "" {
		_, err = n.tx.Exec("ROLLBACK TO SAVEPOINT " + n.savePointID)
	} else if !n.nested {
		err = n.tx.Rollback()
	}

	if err != nil {
		return err
	}

	n.tx = nil
	n.Driver = nil

	return nil
}

func (n *node) Commit() error {
	if n.tx == nil {
		return ErrNotInTransaction
	}

	var err error

	if n.savePointID != "" {
		_, err = n.tx.Exec("RELEASE SAVEPOINT " + n.savePointID)
	} else if !n.nested {
		err = n.tx.Commit()
	}

	if err != nil {
		return err
	}

	n.tx = nil
	n.Driver = nil

	return nil
}

// Tx returns the underlying transaction.
func (n *node) Tx() *sqlx.Tx {
	return n.tx
}

// Option to configure sqalx
type Option func(*node) error

// SavePoint option enables PostgreSQL Savepoints for nested transactions.
func SavePoint(enabled bool) Option {
	return func(n *node) error {
		if enabled && n.Driver.DriverName() != "postgres" {
			return ErrIncompatibleOption
		}
		n.savePointEnabled = enabled
		return nil
	}
}
<commit_msg>Using latest uuid version<commit_after>package sqalx

import (
	"database/sql"
	"errors"
	"strings"

	"github.com/jmoiron/sqlx"
	uuid "github.com/satori/go.uuid"
)

var (
	// ErrNotInTransaction is returned when using Commit
	// outside of a transaction.
	ErrNotInTransaction = errors.New("not in transaction")

	// ErrIncompatibleOption is returned when using an option incompatible
	// with the selected driver.
	ErrIncompatibleOption = errors.New("incompatible option")
)

// A Node is a database driver that can manage nested transactions.
type Node interface {
	Driver

	// Close the underlying sqlx connection.
	Close() error
	// Begin a new transaction.
	Beginx() (Node, error)
	// Rollback the associated transaction.
	Rollback() error
	// Commit the assiociated transaction.
	Commit() error
	// Tx returns the underlying transaction.
	Tx() *sqlx.Tx
}

// A Driver can query the database. It can either be a *sqlx.DB or a *sqlx.Tx
// and therefore is limited to the methods they have in common.
type Driver interface {
	sqlx.Execer
	sqlx.Queryer
	sqlx.Preparer
	BindNamed(query string, arg interface{}) (string, []interface{}, error)
	DriverName() string
	Get(dest interface{}, query string, args ...interface{}) error
	MustExec(query string, args ...interface{}) sql.Result
	NamedExec(query string, arg interface{}) (sql.Result, error)
	NamedQuery(query string, arg interface{}) (*sqlx.Rows, error)
	PrepareNamed(query string) (*sqlx.NamedStmt, error)
	Preparex(query string) (*sqlx.Stmt, error)
	Rebind(query string) string
	Select(dest interface{}, query string, args ...interface{}) error
}

// New creates a new Node with the given DB.
func New(db *sqlx.DB, options ...Option) (Node, error) {
	n := node{
		db:     db,
		Driver: db,
	}

	for _, opt := range options {
		err := opt(&n)
		if err != nil {
			return nil, err
		}
	}

	return &n, nil
}

// NewFromTransaction creates a new Node from the given transaction.
func NewFromTransaction(tx *sqlx.Tx, options ...Option) (Node, error) {
	n := node{
		tx:     tx,
		Driver: tx,
	}

	for _, opt := range options {
		err := opt(&n)
		if err != nil {
			return nil, err
		}
	}

	return &n, nil
}

// Connect to a database.
func Connect(driverName, dataSourceName string, options ...Option) (Node, error) {
	db, err := sqlx.Connect(driverName, dataSourceName)
	if err != nil {
		return nil, err
	}

	node, err := New(db, options...)
	if err != nil {
		// the connection has been opened within this function, we must close it
		// on error.
		db.Close()
		return nil, err
	}

	return node, nil
}

type node struct {
	Driver
	db               *sqlx.DB
	tx               *sqlx.Tx
	savePointID      string
	savePointEnabled bool
	nested           bool
}

func (n *node) Close() error {
	return n.db.Close()
}

func (n node) Beginx() (Node, error) {
	var err error

	switch {
	case n.tx == nil:
		// new actual transaction
		n.tx, err = n.db.Beginx()
		n.Driver = n.tx
	case n.savePointEnabled:
		// already in a transaction: using savepoints
		n.nested = true
		// savepoints name must start with a char and cannot contain dashes (-)
		n.savePointID = "sp_" + strings.Replace(uuid.Must(uuid.NewV1()).String(), "-", "_", -1)
		_, err = n.tx.Exec("SAVEPOINT " + n.savePointID)
	default:
		// already in a transaction: reusing current transaction
		n.nested = true
	}

	if err != nil {
		return nil, err
	}

	return &n, nil
}

func (n *node) Rollback() error {
	if n.tx == nil {
		return nil
	}

	var err error

	if n.savePointEnabled && n.savePointID != "" {
		_, err = n.tx.Exec("ROLLBACK TO SAVEPOINT " + n.savePointID)
	} else if !n.nested {
		err = n.tx.Rollback()
	}

	if err != nil {
		return err
	}

	n.tx = nil
	n.Driver = nil

	return nil
}

func (n *node) Commit() error {
	if n.tx == nil {
		return ErrNotInTransaction
	}

	var err error

	if n.savePointID != "" {
		_, err = n.tx.Exec("RELEASE SAVEPOINT " + n.savePointID)
	} else if !n.nested {
		err = n.tx.Commit()
	}

	if err != nil {
		return err
	}

	n.tx = nil
	n.Driver = nil

	return nil
}

// Tx returns the underlying transaction.
func (n *node) Tx() *sqlx.Tx {
	return n.tx
}

// Option to configure sqalx
type Option func(*node) error

// SavePoint option enables PostgreSQL Savepoints for nested transactions.
func SavePoint(enabled bool) Option {
	return func(n *node) error {
		if enabled && n.Driver.DriverName() != "postgres" {
			return ErrIncompatibleOption
		}
		n.savePointEnabled = enabled
		return nil
	}
}
<|endoftext|>


<commit_before>package gohavestorage

import (
  "testing"
  "encoding/json"
)
var Key = ""
var Account = ""
var Table = "TestTable"

func CreateTable(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.CreateTable(Table)
}

func TestQueryTables(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.QueryTables()
}

func InsertEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json1, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json1)

  entity.RowKey = "456"
  json2, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json2)

  entity.RowKey = "789"
  json3, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json3)
}

func TestQueryEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntity(Table, "ABC", "123", "")
}

func TestQueryEntityWithSelect(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntity(Table, "ABC", "123", "RowKey,Property1,Property3")
}

func TestQueryEntities(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "", "", "")
}

func TestQueryEntitiesWithSelect(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "", "")
}

func TestQueryEntitiesWithTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "", "", "1")
}

func TestQueryEntitiesWithSelectAndTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "", "1")
}

func TestQueryEntitiesWithSelectAndFilterAndTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "RowKey%20gt%20'123'", "1")
}

func TestDeleteEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.DeleteEntity(Table, "ABC", "123")
}

func TestUpdateEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.UpdateEntity(Table, "ABC", "456", json)
}

func TestMergeEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.MergeEntity(Table, "ABC", "456", json)
}

func TestInsertOrMergeEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
    entity.PartitionKey = "ABC"
    entity.RowKey = "123"
    entity.Property1 = "Value1"
    entity.Property2 = "Value2"
    entity.Property3 = "Value3"

    json, _ := json.Marshal(entity)
    tableStorageProxy.InsertOrMergeEntity(Table, "ABC", "456", json)
  }

func TestInsertOrReplaceEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.InsertOrReplaceEntity(Table, "ABC", "456", json)
}

func DeleteTable(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.DeleteTable(Table)
}

type TestEntity struct {
  PartitionKey string
  RowKey string
  Property1 string
  Property2 string
  Property3 string
}
<commit_msg>All tests enabled<commit_after>package gohavestorage

import (
  "testing"
  "encoding/json"
)
var Key = ""
var Account = ""
var Table = "TestTable"

func TestCreateTable(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.CreateTable(Table)
}

func TestQueryTables(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.QueryTables()
}

func TestInsertEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json1, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json1)

  entity.RowKey = "456"
  json2, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json2)

  entity.RowKey = "789"
  json3, _ := json.Marshal(entity)
  tableStorageProxy.InsertEntity(Table, json3)
}

func TestQueryEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntity(Table, "ABC", "123", "")
}

func TestQueryEntityWithSelect(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntity(Table, "ABC", "123", "RowKey,Property1,Property3")
}

func TestQueryEntities(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "", "", "")
}

func TestQueryEntitiesWithSelect(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "", "")
}

func TestQueryEntitiesWithTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "", "", "1")
}

func TestQueryEntitiesWithSelectAndTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "", "1")
}

func TestQueryEntitiesWithSelectAndFilterAndTop(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.QueryEntities(Table, "RowKey,Property1,Property3", "RowKey%20gt%20'123'", "1")
}

func TestDeleteEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  tableStorageProxy.DeleteEntity(Table, "ABC", "123")
}

func TestUpdateEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.UpdateEntity(Table, "ABC", "456", json)
}

func TestMergeEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.MergeEntity(Table, "ABC", "456", json)
}

func TestInsertOrMergeEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
    entity.PartitionKey = "ABC"
    entity.RowKey = "123"
    entity.Property1 = "Value1"
    entity.Property2 = "Value2"
    entity.Property3 = "Value3"

    json, _ := json.Marshal(entity)
    tableStorageProxy.InsertOrMergeEntity(Table, "ABC", "456", json)
  }

func TestInsertOrReplaceEntity(t *testing.T) {
  goHaveStorage := New(Account, Key)
  tableStorageProxy := goHaveStorage.NewTableStorageProxy()

  entity := &TestEntity{}
  entity.PartitionKey = "ABC"
  entity.RowKey = "123"
  entity.Property1 = "Value1"
  entity.Property2 = "Value2"
  entity.Property3 = "Value3"

  json, _ := json.Marshal(entity)
  tableStorageProxy.InsertOrReplaceEntity(Table, "ABC", "456", json)
}

func TestDeleteTable(t *testing.T) {
	goHaveStorage := New(Account, Key)
	tableStorageProxy := goHaveStorage.NewTableStorageProxy()
	tableStorageProxy.DeleteTable(Table)
}

type TestEntity struct {
  PartitionKey string
  RowKey string
  Property1 string
  Property2 string
  Property3 string
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"path/filepath"
	"strings"
)

func purgeUnmanagedContent(allBasedirs map[string]bool, allEnvironments map[string]bool) {
	if !stringSliceContains(config.PurgeLevels, "deployment") {
		if !stringSliceContains(config.PurgeLevels, "environment") {
			// nothing allowed to purge
			return
		}
	}

	for source, sa := range config.Sources {
		prefix := resolveSourcePrefix(source, sa)

		if len(environmentParam) > 0 {
			if !strings.HasPrefix(environmentParam, prefix) {
				Debugf("Skipping purging unmanaged content for source '" + source + "', because -environment parameter is set to " + environmentParam)
				continue
			}
		}

		// Clean up unknown environment directories
		if len(branchParam) == 0 {
			for basedir, _ := range allBasedirs {
				globPath := filepath.Join(basedir, prefix+"*")
				Debugf("Glob'ing with path " + globPath)
				environments, _ := filepath.Glob(globPath)

				whitelistEnvironments := []string{}
				if len(config.DeploymentPurgeWhitelist) > 0 {
					for _, wlpattern := range config.DeploymentPurgeWhitelist {
						whitelistGlobPath := filepath.Join(basedir, wlpattern)
						Debugf("deployment_purge_whitelist Glob'ing with path " + whitelistGlobPath)
						we, _ := filepath.Glob(whitelistGlobPath)
						whitelistEnvironments = append(whitelistEnvironments, we...)
					}
				}

				for _, env := range environments {
					envPath := strings.Split(env, "/")
					envName := envPath[len(envPath)-1]
					if len(environmentParam) > 0 {
						if envName != environmentParam {
							Debugf("Skipping purging unmanaged content for Puppet environment '" + envName + "', because -environment parameter is set to " + environmentParam)
							continue
						}
					}
					if stringSliceContains(config.PurgeLevels, "environment") {
						if allEnvironments[envName] {
							checkForStaleContent(env)
						}
					}
					if stringSliceContains(config.PurgeLevels, "deployment") {
						Debugf("Checking if environment should exist: " + envName)
						if allEnvironments[envName] {
							Debugf("Not purging environment " + envName)
						} else if stringSliceContains(whitelistEnvironments, filepath.Join(basedir, envName)) {
							Debugf("Not purging environment " + envName + " due to deployment_purge_whitelist match")
						} else {
							Infof("Removing unmanaged environment " + envName)
							if !dryRun {
								purgeDir(filepath.Join(basedir, envName), "purgeStaleContent()")
							}
						}
					}
				}
			}
		} else {
			if stringSliceContains(config.PurgeLevels, "environment") {
				// check for purgeable content inside -branch folder
				checkForStaleContent(filepath.Join(sa.Basedir, prefix+branchParam))
			}
		}
	}
}

func checkForStaleContent(workDir string) {
	// add purge whitelist
	if len(config.PurgeWhitelist) > 0 {
		Debugf("additional purge whitelist items: " + strings.Join(config.PurgeWhitelist, " "))
		for _, wlItem := range config.PurgeWhitelist {
			desiredContent = append(desiredContent, filepath.Join(workDir, wlItem))
		}
	}

	checkForStaleContent := func(path string, info os.FileInfo, err error) error {
		//Debugf("filepath.Walk'ing found path: " + path)
		stale := true
		for _, desiredFile := range desiredContent {
			//Debugf("comparing found path: " + path + " with managed path: " + desiredFile)
			//if strings.HasPrefix(path, desiredFile) || path == workDir {
			if path == desiredFile || path == workDir {
				stale = false
			}
		}

		if stale {
			Infof("Removing unmanaged path " + path)
			purgeDir(path, "checkForStaleContent()")
		}
		return nil
	}

	c := make(chan error)
	Debugf("filepath.Walk'ing directory " + workDir)
	go func() { c <- filepath.Walk(workDir, checkForStaleContent) }()
	<-c // Walk done
}
<commit_msg>add unchangedModuleDirs<commit_after>package main

import (
	"os"
	"path/filepath"
	"strings"
)

func purgeUnmanagedContent(allBasedirs map[string]bool, allEnvironments map[string]bool) {
	if !stringSliceContains(config.PurgeLevels, "deployment") {
		if !stringSliceContains(config.PurgeLevels, "environment") {
			// nothing allowed to purge
			return
		}
	}

	for source, sa := range config.Sources {
		prefix := resolveSourcePrefix(source, sa)

		if len(environmentParam) > 0 {
			if !strings.HasPrefix(environmentParam, prefix) {
				Debugf("Skipping purging unmanaged content for source '" + source + "', because -environment parameter is set to " + environmentParam)
				continue
			}
		}

		// Clean up unknown environment directories
		if len(branchParam) == 0 {
			for basedir, _ := range allBasedirs {
				globPath := filepath.Join(basedir, prefix+"*")
				Debugf("Glob'ing with path " + globPath)
				environments, _ := filepath.Glob(globPath)

				whitelistEnvironments := []string{}
				if len(config.DeploymentPurgeWhitelist) > 0 {
					for _, wlpattern := range config.DeploymentPurgeWhitelist {
						whitelistGlobPath := filepath.Join(basedir, wlpattern)
						Debugf("deployment_purge_whitelist Glob'ing with path " + whitelistGlobPath)
						we, _ := filepath.Glob(whitelistGlobPath)
						whitelistEnvironments = append(whitelistEnvironments, we...)
					}
				}

				for _, env := range environments {
					envPath := strings.Split(env, "/")
					envName := envPath[len(envPath)-1]
					if len(environmentParam) > 0 {
						if envName != environmentParam {
							Debugf("Skipping purging unmanaged content for Puppet environment '" + envName + "', because -environment parameter is set to " + environmentParam)
							continue
						}
					}
					if stringSliceContains(config.PurgeLevels, "environment") {
						if allEnvironments[envName] {
							checkForStaleContent(env)
						}
					}
					if stringSliceContains(config.PurgeLevels, "deployment") {
						Debugf("Checking if environment should exist: " + envName)
						if allEnvironments[envName] {
							Debugf("Not purging environment " + envName)
						} else if stringSliceContains(whitelistEnvironments, filepath.Join(basedir, envName)) {
							Debugf("Not purging environment " + envName + " due to deployment_purge_whitelist match")
						} else {
							Infof("Removing unmanaged environment " + envName)
							if !dryRun {
								purgeDir(filepath.Join(basedir, envName), "purgeStaleContent()")
							}
						}
					}
				}
			}
		} else {
			if stringSliceContains(config.PurgeLevels, "environment") {
				// check for purgeable content inside -branch folder
				checkForStaleContent(filepath.Join(sa.Basedir, prefix+branchParam))
			}
		}
	}
}

func checkForStaleContent(workDir string) {
	// add purge whitelist
	if len(config.PurgeWhitelist) > 0 {
		Debugf("additional purge whitelist items: " + strings.Join(config.PurgeWhitelist, " "))
		for _, wlItem := range config.PurgeWhitelist {
			desiredContent = append(desiredContent, filepath.Join(workDir, wlItem))
		}
	}

	checkForStaleContent := func(path string, info os.FileInfo, err error) error {
		//Debugf("filepath.Walk'ing found path: " + path)
		stale := true
		for _, desiredFile := range desiredContent {
			for _, unchangedModuleDir := range unchangedModuleDirs {
				if strings.HasPrefix(path, unchangedModuleDir) {
					stale = false
					break
				}
			}
			if path == desiredFile || path == workDir {
				stale = false
				break
			}
		}

		if stale {
			Infof("Removing unmanaged path " + path)
			purgeDir(path, "checkForStaleContent()")
		}
		return nil
	}

	c := make(chan error)
	Debugf("filepath.Walk'ing directory " + workDir)
	go func() { c <- filepath.Walk(workDir, checkForStaleContent) }()
	<-c // Walk done
}
<|endoftext|>


<commit_before>package agollo

import (
	"github.com/zouyx/agollo/v3/agcache"
	_ "github.com/zouyx/agollo/v3/agcache/memory"
	"github.com/zouyx/agollo/v3/cluster"
	_ "github.com/zouyx/agollo/v3/cluster/roundrobin"
	"github.com/zouyx/agollo/v3/component"
	"github.com/zouyx/agollo/v3/component/log"
	"github.com/zouyx/agollo/v3/component/notify"
	"github.com/zouyx/agollo/v3/component/serverlist"
	"github.com/zouyx/agollo/v3/env"
	"github.com/zouyx/agollo/v3/env/config"
	"github.com/zouyx/agollo/v3/env/file"
	_ "github.com/zouyx/agollo/v3/env/file/json"
	"github.com/zouyx/agollo/v3/extension"
	"github.com/zouyx/agollo/v3/protocol/auth"
	_ "github.com/zouyx/agollo/v3/protocol/auth/sign"
	"github.com/zouyx/agollo/v3/storage"
	_ "github.com/zouyx/agollo/v3/utils/parse/normal"
	_ "github.com/zouyx/agollo/v3/utils/parse/properties"
	_ "github.com/zouyx/agollo/v3/utils/parse/yml"
)

var (
	initAppConfigFunc func() (*config.AppConfig, error)
)

//InitCustomConfig init config by custom
func InitCustomConfig(loadAppConfig func() (*config.AppConfig, error)) {
	initAppConfigFunc = loadAppConfig
}

//start apollo
func Start() error {
	return startAgollo()
}

//SetSignature 设置自定义 http 授权控件
func SetSignature(auth auth.HTTPAuth) {
	if auth != nil {
		extension.SetHTTPAuth(auth)
	}
}

//SetBackupFileHandler 设置自定义备份文件处理组件
func SetBackupFileHandler(file file.FileHandler) {
	if file != nil {
		extension.SetFileHandler(file)
	}
}

//SetLoadBalance 设置自定义负载均衡组件
func SetLoadBalance(loadBalance cluster.LoadBalance) {
	if loadBalance != nil {
		extension.SetLoadBalance(loadBalance)
	}
}

//SetLogger 设置自定义logger组件
func SetLogger(loggerInterface log.LoggerInterface) {
	if loggerInterface != nil {
		log.InitLogger(loggerInterface)
	}
}

//UseEventDispatch  添加为某些key分发event功能
func UseEventDispatch() {
	storage.UseEventDispatch()
}

//SetCache 设置自定义cache组件
func SetCache(cacheFactory agcache.CacheFactory) {
	if cacheFactory != nil {
		extension.SetCacheFactory(cacheFactory)
		storage.InitConfigCache()
	}
}

func startAgollo() error {
	// 有了配置之后才能进行初始化
	if err := env.InitConfig(initAppConfigFunc); err != nil {
		return err
	}

	notify.InitAllNotifications(nil)
	serverlist.InitSyncServerIPList()

	//first sync
	if err := notify.SyncConfigs(); err != nil {
		return err
	}
	log.Debug("init notifySyncConfigServices finished")

	//start long poll sync config
	go component.StartRefreshConfig(&notify.ConfigComponent{})

	log.Info("agollo start finished ! ")

	return nil
}
<commit_msg>fix change listener<commit_after>package agollo

import (
	"container/list"
	"github.com/zouyx/agollo/v3/agcache"
	_ "github.com/zouyx/agollo/v3/agcache/memory"
	"github.com/zouyx/agollo/v3/cluster"
	_ "github.com/zouyx/agollo/v3/cluster/roundrobin"
	"github.com/zouyx/agollo/v3/component"
	"github.com/zouyx/agollo/v3/component/log"
	"github.com/zouyx/agollo/v3/component/notify"
	"github.com/zouyx/agollo/v3/component/serverlist"
	"github.com/zouyx/agollo/v3/env"
	"github.com/zouyx/agollo/v3/env/config"
	"github.com/zouyx/agollo/v3/env/file"
	_ "github.com/zouyx/agollo/v3/env/file/json"
	"github.com/zouyx/agollo/v3/extension"
	"github.com/zouyx/agollo/v3/protocol/auth"
	_ "github.com/zouyx/agollo/v3/protocol/auth/sign"
	"github.com/zouyx/agollo/v3/storage"
	_ "github.com/zouyx/agollo/v3/utils/parse/normal"
	_ "github.com/zouyx/agollo/v3/utils/parse/properties"
	_ "github.com/zouyx/agollo/v3/utils/parse/yml"
)

var (
	initAppConfigFunc func() (*config.AppConfig, error)
)

//InitCustomConfig init config by custom
func InitCustomConfig(loadAppConfig func() (*config.AppConfig, error)) {
	initAppConfigFunc = loadAppConfig
}

//start apollo
func Start() error {
	return startAgollo()
}

//SetSignature 设置自定义 http 授权控件
func SetSignature(auth auth.HTTPAuth) {
	if auth != nil {
		extension.SetHTTPAuth(auth)
	}
}

//SetBackupFileHandler 设置自定义备份文件处理组件
func SetBackupFileHandler(file file.FileHandler) {
	if file != nil {
		extension.SetFileHandler(file)
	}
}

//SetLoadBalance 设置自定义负载均衡组件
func SetLoadBalance(loadBalance cluster.LoadBalance) {
	if loadBalance != nil {
		extension.SetLoadBalance(loadBalance)
	}
}

//SetLogger 设置自定义logger组件
func SetLogger(loggerInterface log.LoggerInterface) {
	if loggerInterface != nil {
		log.InitLogger(loggerInterface)
	}
}

//UseEventDispatch  添加为某些key分发event功能
func UseEventDispatch() {
	storage.UseEventDispatch()
}

//SetCache 设置自定义cache组件
func SetCache(cacheFactory agcache.CacheFactory) {
	if cacheFactory != nil {
		extension.SetCacheFactory(cacheFactory)
		storage.InitConfigCache()
	}
}

//AddChangeListener 增加变更监控
func AddChangeListener(listener storage.ChangeListener) {
	storage.AddChangeListener(listener)
}

//RemoveChangeListener 增加变更监控
func RemoveChangeListener(listener storage.ChangeListener) {
	storage.AddChangeListener(listener)
}

//GetChangeListeners 获取配置修改监听器列表
func GetChangeListeners() *list.List {
	return storage.GetChangeListeners()
}

func startAgollo() error {
	// 有了配置之后才能进行初始化
	if err := env.InitConfig(initAppConfigFunc); err != nil {
		return err
	}

	notify.InitAllNotifications(nil)
	serverlist.InitSyncServerIPList()

	//first sync
	if err := notify.SyncConfigs(); err != nil {
		return err
	}
	log.Debug("init notifySyncConfigServices finished")

	//start long poll sync config
	go component.StartRefreshConfig(&notify.ConfigComponent{})

	log.Info("agollo start finished ! ")

	return nil
}
<|endoftext|>


<commit_before>package metrics

import (
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/hashicorp/go-immutable-radix"
)

// Config is used to configure metrics settings
type Config struct {
	ServiceName          string        // Prefixed with keys to separate services
	HostName             string        // Hostname to use. If not provided and EnableHostname, it will be os.Hostname
	EnableHostname       bool          // Enable prefixing gauge values with hostname
	EnableHostnameLabel  bool          // Enable adding hostname to labels
	EnableServiceLabel   bool          // Enable adding service to labels
	EnableRuntimeMetrics bool          // Enables profiling of runtime metrics (GC, Goroutines, Memory)
	EnableTypePrefix     bool          // Prefixes key with a type ("counter", "gauge", "timer")
	TimerGranularity     time.Duration // Granularity of timers.
	ProfileInterval      time.Duration // Interval to profile runtime metrics

	AllowedPrefixes []string // A list of metric prefixes to allow, with '.' as the separator
	BlockedPrefixes []string // A list of metric prefixes to block, with '.' as the separator
	AllowedLabels   []string // A list of metric labels to allow, with '.' as the separator
	BlockedLabels   []string // A list of metric labels to block, with '.' as the separator
	FilterDefault   bool     // Whether to allow metrics by default
}

// Metrics represents an instance of a metrics sink that can
// be used to emit
type Metrics struct {
	Config
	lastNumGC     uint32
	sink          MetricSink
	filter        *iradix.Tree
	allowedLabels map[string]bool
	blockedLabels map[string]bool
	filterLock    sync.RWMutex // Lock filters and allowedLabels/blockedLabels access
}

// Shared global metrics instance
var globalMetrics atomic.Value // *Metrics

func init() {
	// Initialize to a blackhole sink to avoid errors
	globalMetrics.Store(&Metrics{sink: &BlackholeSink{}})
}

// DefaultConfig provides a sane default configuration
func DefaultConfig(serviceName string) *Config {
	c := &Config{
		ServiceName:          serviceName, // Use client provided service
		HostName:             "",
		EnableHostname:       true,             // Enable hostname prefix
		EnableRuntimeMetrics: true,             // Enable runtime profiling
		EnableTypePrefix:     false,            // Disable type prefix
		TimerGranularity:     time.Millisecond, // Timers are in milliseconds
		ProfileInterval:      time.Second,      // Poll runtime every second
		FilterDefault:        true,             // Don't filter metrics by default
	}

	// Try to get the hostname
	name, _ := os.Hostname()
	c.HostName = name
	return c
}

// New is used to create a new instance of Metrics
func New(conf *Config, sink MetricSink) (*Metrics, error) {
	met := &Metrics{}
	met.Config = *conf
	met.sink = sink
	met.UpdateFilterAndLabels(conf.AllowedPrefixes, conf.BlockedPrefixes, conf.AllowedLabels, conf.BlockedLabels)

	// Start the runtime collector
	if conf.EnableRuntimeMetrics {
		go met.collectStats()
	}
	return met, nil
}

// NewGlobal is the same as New, but it assigns the metrics object to be
// used globally as well as returning it.
func NewGlobal(conf *Config, sink MetricSink) (*Metrics, error) {
	metrics, err := New(conf, sink)
	if err == nil {
		globalMetrics.Store(metrics)
	}
	return metrics, err
}

// Proxy all the methods to the globalMetrics instance
func SetGauge(key []string, val float32) {
	globalMetrics.Load().(*Metrics).SetGauge(key, val)
}

func SetGaugeWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).SetGaugeWithLabels(key, val, labels)
}

func EmitKey(key []string, val float32) {
	globalMetrics.Load().(*Metrics).EmitKey(key, val)
}

func IncrCounter(key []string, val float32) {
	globalMetrics.Load().(*Metrics).IncrCounter(key, val)
}

func IncrCounterWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).IncrCounterWithLabels(key, val, labels)
}

func AddSample(key []string, val float32) {
	globalMetrics.Load().(*Metrics).AddSample(key, val)
}

func AddSampleWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).AddSampleWithLabels(key, val, labels)
}

func MeasureSince(key []string, start time.Time) {
	globalMetrics.Load().(*Metrics).MeasureSince(key, start)
}

func MeasureSinceWithLabels(key []string, start time.Time, labels []Label) {
	globalMetrics.Load().(*Metrics).MeasureSinceWithLabels(key, start, labels)
}

func UpdateFilter(allow, block []string) {
	globalMetrics.Load().(*Metrics).UpdateFilter(allow, block)
}

// UpdateFilterAndLabels set allow/block prefixes of metrics while allowedLabels
// and blockedLabels - when not nil - allow filtering of labels in order to
// block/allow globally labels (especially useful when having large number of
// values for a given label). See README.md for more information about usage.
func UpdateFilterAndLabels(allow, block, allowedLabels, blockedLabels []string) {
	globalMetrics.Load().(*Metrics).UpdateFilterAndLabels(allow, block, allowedLabels, blockedLabels)
}
<commit_msg>Expose the global default metrics<commit_after>package metrics

import (
	"os"
	"sync"
	"sync/atomic"
	"time"

	iradix "github.com/hashicorp/go-immutable-radix"
)

// Config is used to configure metrics settings
type Config struct {
	ServiceName          string        // Prefixed with keys to separate services
	HostName             string        // Hostname to use. If not provided and EnableHostname, it will be os.Hostname
	EnableHostname       bool          // Enable prefixing gauge values with hostname
	EnableHostnameLabel  bool          // Enable adding hostname to labels
	EnableServiceLabel   bool          // Enable adding service to labels
	EnableRuntimeMetrics bool          // Enables profiling of runtime metrics (GC, Goroutines, Memory)
	EnableTypePrefix     bool          // Prefixes key with a type ("counter", "gauge", "timer")
	TimerGranularity     time.Duration // Granularity of timers.
	ProfileInterval      time.Duration // Interval to profile runtime metrics

	AllowedPrefixes []string // A list of metric prefixes to allow, with '.' as the separator
	BlockedPrefixes []string // A list of metric prefixes to block, with '.' as the separator
	AllowedLabels   []string // A list of metric labels to allow, with '.' as the separator
	BlockedLabels   []string // A list of metric labels to block, with '.' as the separator
	FilterDefault   bool     // Whether to allow metrics by default
}

// Metrics represents an instance of a metrics sink that can
// be used to emit
type Metrics struct {
	Config
	lastNumGC     uint32
	sink          MetricSink
	filter        *iradix.Tree
	allowedLabels map[string]bool
	blockedLabels map[string]bool
	filterLock    sync.RWMutex // Lock filters and allowedLabels/blockedLabels access
}

// Shared global metrics instance
var globalMetrics atomic.Value // *Metrics

func init() {
	// Initialize to a blackhole sink to avoid errors
	globalMetrics.Store(&Metrics{sink: &BlackholeSink{}})
}

// Default returns the shared global metrics instance.
func Default() *Metrics {
	return globalMetrics.Load().(*Metrics)
}

// DefaultConfig provides a sane default configuration
func DefaultConfig(serviceName string) *Config {
	c := &Config{
		ServiceName:          serviceName, // Use client provided service
		HostName:             "",
		EnableHostname:       true,             // Enable hostname prefix
		EnableRuntimeMetrics: true,             // Enable runtime profiling
		EnableTypePrefix:     false,            // Disable type prefix
		TimerGranularity:     time.Millisecond, // Timers are in milliseconds
		ProfileInterval:      time.Second,      // Poll runtime every second
		FilterDefault:        true,             // Don't filter metrics by default
	}

	// Try to get the hostname
	name, _ := os.Hostname()
	c.HostName = name
	return c
}

// New is used to create a new instance of Metrics
func New(conf *Config, sink MetricSink) (*Metrics, error) {
	met := &Metrics{}
	met.Config = *conf
	met.sink = sink
	met.UpdateFilterAndLabels(conf.AllowedPrefixes, conf.BlockedPrefixes, conf.AllowedLabels, conf.BlockedLabels)

	// Start the runtime collector
	if conf.EnableRuntimeMetrics {
		go met.collectStats()
	}
	return met, nil
}

// NewGlobal is the same as New, but it assigns the metrics object to be
// used globally as well as returning it.
func NewGlobal(conf *Config, sink MetricSink) (*Metrics, error) {
	metrics, err := New(conf, sink)
	if err == nil {
		globalMetrics.Store(metrics)
	}
	return metrics, err
}

// Proxy all the methods to the globalMetrics instance
func SetGauge(key []string, val float32) {
	globalMetrics.Load().(*Metrics).SetGauge(key, val)
}

func SetGaugeWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).SetGaugeWithLabels(key, val, labels)
}

func EmitKey(key []string, val float32) {
	globalMetrics.Load().(*Metrics).EmitKey(key, val)
}

func IncrCounter(key []string, val float32) {
	globalMetrics.Load().(*Metrics).IncrCounter(key, val)
}

func IncrCounterWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).IncrCounterWithLabels(key, val, labels)
}

func AddSample(key []string, val float32) {
	globalMetrics.Load().(*Metrics).AddSample(key, val)
}

func AddSampleWithLabels(key []string, val float32, labels []Label) {
	globalMetrics.Load().(*Metrics).AddSampleWithLabels(key, val, labels)
}

func MeasureSince(key []string, start time.Time) {
	globalMetrics.Load().(*Metrics).MeasureSince(key, start)
}

func MeasureSinceWithLabels(key []string, start time.Time, labels []Label) {
	globalMetrics.Load().(*Metrics).MeasureSinceWithLabels(key, start, labels)
}

func UpdateFilter(allow, block []string) {
	globalMetrics.Load().(*Metrics).UpdateFilter(allow, block)
}

// UpdateFilterAndLabels set allow/block prefixes of metrics while allowedLabels
// and blockedLabels - when not nil - allow filtering of labels in order to
// block/allow globally labels (especially useful when having large number of
// values for a given label). See README.md for more information about usage.
func UpdateFilterAndLabels(allow, block, allowedLabels, blockedLabels []string) {
	globalMetrics.Load().(*Metrics).UpdateFilterAndLabels(allow, block, allowedLabels, blockedLabels)
}
<|endoftext|>


<commit_before>package mailgun

import (
	"github.com/mbanzon/simplehttp"
	"strconv"
	"time"
)

type Stat struct {
	Event      string `json:"event"`
	TotalCount int    `json:"total_count"`
	CreatedAt  string `json:"created_at"`
	Id         string `json:"id"`
	Tags       string `json:"tags"`
}

type statsEnvelope struct {
	TotalCount int    `json:"total_count"`
	Items      []Stat `json:"items"`
}

func (m *mailgunImpl) GetStats(limit int, skip int, startDate time.Time, event ...string) (int, []Stat, error) {
	r := simplehttp.NewGetRequest(generateApiUrl(m, statsEndpoint))

	if limit != -1 {
		r.AddParameter("limit", strconv.Itoa(limit))
	}
	if skip != -1 {
		r.AddParameter("skip", strconv.Itoa(skip))
	}

	r.AddParameter("start-date", startDate.Format("2006-02-01"))

	for _, e := range event {
		r.AddParameter("event", e)
	}

	var res statsEnvelope
	err := r.MakeJSONRequest(&res)
	if err != nil {
		return -1, nil, err
	} else {
		return res.TotalCount, res.Items, nil
	}
}
<commit_msg>Implementing the stats api.<commit_after>package mailgun

import (
	"github.com/mbanzon/simplehttp"
	"strconv"
	"time"
)

type Stat struct {
	Event      string      `json:"event"`
	TotalCount int         `json:"total_count"`
	CreatedAt  string      `json:"created_at"`
	Id         string      `json:"id"`
	Tags       interface{} `json:"tags"`
}

type statsEnvelope struct {
	TotalCount int    `json:"total_count"`
	Items      []Stat `json:"items"`
}

func (m *mailgunImpl) GetStats(limit int, skip int, startDate time.Time, event ...string) (int, []Stat, error) {
	r := simplehttp.NewGetRequest(generateApiUrl(m, statsEndpoint))

	if limit != -1 {
		r.AddParameter("limit", strconv.Itoa(limit))
	}
	if skip != -1 {
		r.AddParameter("skip", strconv.Itoa(skip))
	}

	r.AddParameter("start-date", startDate.Format("2006-02-01"))

	for _, e := range event {
		r.AddParameter("event", e)
	}
	r.SetBasicAuth(basicAuthUser, m.ApiKey())

	var res statsEnvelope
	err := r.MakeJSONRequest(&res)
	if err != nil {
		return -1, nil, err
	} else {
		return res.TotalCount, res.Items, nil
	}
}
<|endoftext|>


<commit_before>package gopenid

type Store interface {
	StoreAssociation(*Association) error
	GetAssociation(string, bool) (*Association, error)
	DeleteAssociation(*Association) error
	IsKnownNonce(nonce string) (bool, error)
	StoreNonce(nonce string) error
}
<commit_msg>delete name of arguments<commit_after>package gopenid

type Store interface {
	StoreAssociation(*Association) error
	GetAssociation(string, bool) (*Association, error)
	DeleteAssociation(*Association) error
	IsKnownNonce(string) (bool, error)
	StoreNonce(string) error
}
<|endoftext|>


<commit_before>package drax

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/cpuguy83/drax/api"
	"github.com/cpuguy83/drax/api/client"
	libkvstore "github.com/docker/libkv/store"
	"github.com/hashicorp/raft"
)

type store struct {
	mu      sync.RWMutex
	ttlLock sync.Mutex
	data    *db
	r       *Raft
}

type ttl struct {
	TTL         time.Duration
	CreateTime  time.Time
	CreateIndex uint64
}

type db struct {
	KV   map[string]*libkvstore.KVPair
	TTLs map[string]*ttl
}

func newDB() *db {
	return &db{KV: make(map[string]*libkvstore.KVPair), TTLs: make(map[string]*ttl)}
}

func (s *store) newClient() client.Client {
	leader := s.r.getLeader()
	return client.New(leader, s.r.tlsConfig, defaultTimeout)
}

func newStore() *store {
	return &store{data: newDB()}
}

func (s *store) Get(key string) (*libkvstore.KVPair, error) {
	if s.r.IsLeader() {
		s.mu.RLock()
		defer s.mu.RUnlock()
		return s.get(key)
	}
	return s.newClient().Get(key)
}

func (s *store) get(key string) (*libkvstore.KVPair, error) {
	kv, ok := s.data.KV[key]
	if !ok {
		return nil, libkvstore.ErrKeyNotFound
	}
	return kv, nil
}

func (s *store) Put(key string, value []byte, options *libkvstore.WriteOptions) error {
	if !s.r.IsLeader() {
		return s.newClient().Put(key, value, options)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	req := &api.Request{
		Action: api.Put,
		Key:    key,
		Value:  value,
	}
	if options != nil {
		req.TTL = options.TTL
	}
	return s.apply(req)
}

func (s *store) Delete(key string) error {
	if !s.r.IsLeader() {
		return s.newClient().Delete(key)
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	return s.apply(&api.Request{
		Action: api.Delete,
		Key:    key,
	})
}

func (s *store) Exists(key string) (bool, error) {
	if !s.r.IsLeader() {
		return s.newClient().Exists(key)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.data.KV[key]
	return ok, nil
}

func (s *store) List(prefix string) ([]*libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().List(prefix)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()

	var out []*libkvstore.KVPair

	for k, v := range s.data.KV {
		if !strings.HasPrefix(k, prefix) {
			continue
		}
		var kv libkvstore.KVPair
		kv = *v
		out = append(out, &kv)
	}

	if len(out) == 0 {
		return nil, libkvstore.ErrKeyNotFound
	}
	return out, nil
}

func (s *store) DeleteTree(dir string) error {
	if !s.r.IsLeader() {
		return s.newClient().DeleteTree(dir)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.apply(&api.Request{
		Action: api.DeleteTree,
		Key:    dir,
	})
}

func (s *store) Watch(key string, stopCh <-chan struct{}) (<-chan *libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().Watch(key, stopCh)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) WatchTree(dir string, stopCh <-chan struct{}) (<-chan []*libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().WatchTree(dir, stopCh)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) NewLock(key string, options *libkvstore.LockOptions) (libkvstore.Locker, error) {
	if !s.r.IsLeader() {
		return s.newClient().NewLock(key, options)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) AtomicPut(key string, value []byte, previous *libkvstore.KVPair, options *libkvstore.WriteOptions) (bool, *libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().AtomicPut(key, value, previous, options)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	kv, err := s.get(key)
	if err != nil {
		if previous != nil && err == libkvstore.ErrKeyNotFound {
			return false, nil, libkvstore.ErrKeyModified
		}
		return false, nil, err
	}

	if previous != nil && kv.LastIndex != previous.LastIndex {
		return false, nil, libkvstore.ErrKeyModified
	}

	req = &api.Request{
		Action: api.Put,
		Key:    key,
		Value:  value,
	}
	if options != nil {
		req.TTL = options.TTL
	}
	if err := s.apply(req); err != nil {
		return false, nil, err
	}

	kv, err = s.get(key)
	if err != nil {
		return false, nil, libkvstore.ErrKeyNotFound
	}
	return true, kv, nil
}

func (s *store) AtomicDelete(key string, previous *libkvstore.KVPair) (bool, error) {
	if !s.r.IsLeader() {
		return s.newClient().AtomicDelete(key, previous)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	if previous == nil {
		return false, libkvstore.ErrPreviousNotSpecified
	}

	kv, err := s.get(key)
	if err != nil {
		if err == libkvstore.ErrKeyModified {
			return false, err
		}
	}
	if kv.LastIndex != previous.LastIndex {
		return false, libkvstore.ErrKeyModified
	}
	if err := s.apply(&api.Request{
		Action: api.Delete,
		Key:    key,
	}); err != nil {
		return false, err
	}
	return true, nil
}

func (s *store) Close() {
	return
}

func (s *store) forwardToLeader(ax *api.Request) error {
	leader := s.r.Leader()
	if leader == "" {
		return raft.ErrNotLeader
	}

	return raft.ErrNotLeader
}

func (s *store) apply(ax *api.Request) error {
	if !s.r.IsLeader() {
		return s.forwardToLeader(ax)
	}
	buf := bytes.NewBuffer(nil)
	if err := api.Encode(ax, buf); err != nil {
		return err
	}
	return s.r.Apply(buf.Bytes())
}

func (s *store) waitLeader() {
	leaderCh := s.r.LeaderCh()
	logrus.Debug("store: waiting for leader")
	var state raft.RaftState
	for {
		select {
		case si := <-leaderCh:
			state = si.(raft.RaftState)
		case <-s.r.ShutdownCh():
			return
		}

		if state != raft.Leader {
			continue
		}
		logrus.Debug("store: handling leader")
		s.handleLeader(leaderCh)
		logrus.Debugf("store: waiting for leader")
	}
}

func (s *store) handleLeader(leaderCh <-chan interface{}) {
	for {
		select {
		case state := <-leaderCh:
			if state != raft.Leader {
				return
			}
		case <-s.r.ShutdownCh():
			return
		default:
		}
		time.Sleep(1 * time.Second)

		s.ttlLock.Lock()
		var keys []string
		for k, ttl := range s.data.TTLs {
			if ttlDue(ttl) {
				keys = append(keys, k)
			}
		}
		if len(keys) > 0 {
			logrus.Debugf("reaping TTL's for %v", keys)
			s.reapKeys(keys)
		}
		s.ttlLock.Unlock()
	}
}

func (s *store) reapKeys(keys []string) {
	if err := s.apply(&api.Request{
		Action: reapKeys,
		Args:   keys,
	}); err != nil {
		logrus.Debugf("error reaping keys: %v", err)
	}
}

func ttlDue(t *ttl) bool {
	now := time.Now()
	return now.After(t.CreateTime.Add(t.TTL))
}

type storeFSM store

// TODO: handle watches
func (s *storeFSM) Apply(l *raft.Log) interface{} {
	var ax api.Request
	if err := api.Decode(&ax, bytes.NewBuffer(l.Data)); err != nil {
		return err
	}

	switch ax.Action {
	case api.Delete:
		delete(s.data.KV, ax.Key)
		delete(s.data.TTLs, ax.Key)
	case api.Put:
		s.data.KV[ax.Key] = &libkvstore.KVPair{Key: ax.Key, Value: ax.Value, LastIndex: l.Index}
		if ax.TTL != 0 {
			s.ttlLock.Lock()
			s.data.TTLs[ax.Key] = &ttl{CreateTime: time.Now(), TTL: ax.TTL, CreateIndex: l.Index}
			s.ttlLock.Unlock()
		}
	case api.DeleteTree:
		for k := range s.data.KV {
			if !strings.HasPrefix(k, ax.Key) {
				continue
			}
			delete(s.data.KV, k)
			delete(s.data.TTLs, k)
		}
	case reapKeys:
		s.mu.Lock()
		for _, k := range ax.Args {
			delete(s.data.KV, k)
			delete(s.data.TTLs, k)
		}
		s.mu.Unlock()
	default:
		return fmt.Errorf("unknown api.Request")
	}
	return nil
}

func (s *storeFSM) Snapshot() (raft.FSMSnapshot, error) {
	return s, nil
}

func (s *storeFSM) Restore(r io.ReadCloser) error {
	defer r.Close()
	s.data = newDB()
	return api.Decode(s.data, r)
}

func (s *storeFSM) Persist(sink raft.SnapshotSink) error {
	defer sink.Close()
	s.mu.Lock()
	defer s.mu.Unlock()
	return api.Encode(s.data, sink)
}

func (*storeFSM) Release() {}
<commit_msg>remove check in apply helper for leader<commit_after>package drax

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/cpuguy83/drax/api"
	"github.com/cpuguy83/drax/api/client"
	libkvstore "github.com/docker/libkv/store"
	"github.com/hashicorp/raft"
)

type store struct {
	mu      sync.RWMutex
	ttlLock sync.Mutex
	data    *db
	r       *Raft
}

type ttl struct {
	TTL         time.Duration
	CreateTime  time.Time
	CreateIndex uint64
}

type db struct {
	KV   map[string]*libkvstore.KVPair
	TTLs map[string]*ttl
}

func newDB() *db {
	return &db{KV: make(map[string]*libkvstore.KVPair), TTLs: make(map[string]*ttl)}
}

func (s *store) newClient() client.Client {
	leader := s.r.getLeader()
	return client.New(leader, s.r.tlsConfig, defaultTimeout)
}

func newStore() *store {
	return &store{data: newDB()}
}

func (s *store) Get(key string) (*libkvstore.KVPair, error) {
	if s.r.IsLeader() {
		s.mu.RLock()
		defer s.mu.RUnlock()
		return s.get(key)
	}
	return s.newClient().Get(key)
}

func (s *store) get(key string) (*libkvstore.KVPair, error) {
	kv, ok := s.data.KV[key]
	if !ok {
		return nil, libkvstore.ErrKeyNotFound
	}
	return kv, nil
}

func (s *store) Put(key string, value []byte, options *libkvstore.WriteOptions) error {
	if !s.r.IsLeader() {
		return s.newClient().Put(key, value, options)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	req := &api.Request{
		Action: api.Put,
		Key:    key,
		Value:  value,
	}
	if options != nil {
		req.TTL = options.TTL
	}
	return s.apply(req)
}

func (s *store) Delete(key string) error {
	if !s.r.IsLeader() {
		return s.newClient().Delete(key)
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	return s.apply(&api.Request{
		Action: api.Delete,
		Key:    key,
	})
}

func (s *store) Exists(key string) (bool, error) {
	if !s.r.IsLeader() {
		return s.newClient().Exists(key)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.data.KV[key]
	return ok, nil
}

func (s *store) List(prefix string) ([]*libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().List(prefix)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()

	var out []*libkvstore.KVPair

	for k, v := range s.data.KV {
		if !strings.HasPrefix(k, prefix) {
			continue
		}
		var kv libkvstore.KVPair
		kv = *v
		out = append(out, &kv)
	}

	if len(out) == 0 {
		return nil, libkvstore.ErrKeyNotFound
	}
	return out, nil
}

func (s *store) DeleteTree(dir string) error {
	if !s.r.IsLeader() {
		return s.newClient().DeleteTree(dir)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.apply(&api.Request{
		Action: api.DeleteTree,
		Key:    dir,
	})
}

func (s *store) Watch(key string, stopCh <-chan struct{}) (<-chan *libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().Watch(key, stopCh)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) WatchTree(dir string, stopCh <-chan struct{}) (<-chan []*libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().WatchTree(dir, stopCh)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) NewLock(key string, options *libkvstore.LockOptions) (libkvstore.Locker, error) {
	if !s.r.IsLeader() {
		return s.newClient().NewLock(key, options)
	}
	return nil, libkvstore.ErrCallNotSupported
}

func (s *store) AtomicPut(key string, value []byte, previous *libkvstore.KVPair, options *libkvstore.WriteOptions) (bool, *libkvstore.KVPair, error) {
	if !s.r.IsLeader() {
		return s.newClient().AtomicPut(key, value, previous, options)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	kv, err := s.get(key)
	if err != nil {
		if previous != nil && err == libkvstore.ErrKeyNotFound {
			return false, nil, libkvstore.ErrKeyModified
		}
		return false, nil, err
	}

	if previous != nil && kv.LastIndex != previous.LastIndex {
		return false, nil, libkvstore.ErrKeyModified
	}

	req = &api.Request{
		Action: api.Put,
		Key:    key,
		Value:  value,
	}
	if options != nil {
		req.TTL = options.TTL
	}
	if err := s.apply(req); err != nil {
		return false, nil, err
	}

	kv, err = s.get(key)
	if err != nil {
		return false, nil, libkvstore.ErrKeyNotFound
	}
	return true, kv, nil
}

func (s *store) AtomicDelete(key string, previous *libkvstore.KVPair) (bool, error) {
	if !s.r.IsLeader() {
		return s.newClient().AtomicDelete(key, previous)
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	if previous == nil {
		return false, libkvstore.ErrPreviousNotSpecified
	}

	kv, err := s.get(key)
	if err != nil {
		if err == libkvstore.ErrKeyModified {
			return false, err
		}
	}
	if kv.LastIndex != previous.LastIndex {
		return false, libkvstore.ErrKeyModified
	}
	if err := s.apply(&api.Request{
		Action: api.Delete,
		Key:    key,
	}); err != nil {
		return false, err
	}
	return true, nil
}

func (s *store) Close() {
	return
}

func (s *store) apply(ax *api.Request) error {
	buf := bytes.NewBuffer(nil)
	if err := api.Encode(ax, buf); err != nil {
		return err
	}
	return s.r.Apply(buf.Bytes())
}

func (s *store) waitLeader() {
	leaderCh := s.r.LeaderCh()
	logrus.Debug("store: waiting for leader")
	var state raft.RaftState
	for {
		select {
		case si := <-leaderCh:
			state = si.(raft.RaftState)
		case <-s.r.ShutdownCh():
			return
		}

		if state != raft.Leader {
			continue
		}
		logrus.Debug("store: handling leader")
		s.handleLeader(leaderCh)
		logrus.Debugf("store: waiting for leader")
	}
}

func (s *store) handleLeader(leaderCh <-chan interface{}) {
	for {
		select {
		case state := <-leaderCh:
			if state != raft.Leader {
				return
			}
		case <-s.r.ShutdownCh():
			return
		default:
		}
		time.Sleep(1 * time.Second)

		s.ttlLock.Lock()
		var keys []string
		for k, ttl := range s.data.TTLs {
			if ttlDue(ttl) {
				keys = append(keys, k)
			}
		}
		if len(keys) > 0 {
			logrus.Debugf("reaping TTL's for %v", keys)
			s.reapKeys(keys)
		}
		s.ttlLock.Unlock()
	}
}

func (s *store) reapKeys(keys []string) {
	if err := s.apply(&api.Request{
		Action: reapKeys,
		Args:   keys,
	}); err != nil {
		logrus.Debugf("error reaping keys: %v", err)
	}
}

func ttlDue(t *ttl) bool {
	now := time.Now()
	return now.After(t.CreateTime.Add(t.TTL))
}

type storeFSM store

// TODO: handle watches
func (s *storeFSM) Apply(l *raft.Log) interface{} {
	var ax api.Request
	if err := api.Decode(&ax, bytes.NewBuffer(l.Data)); err != nil {
		return err
	}

	switch ax.Action {
	case api.Delete:
		delete(s.data.KV, ax.Key)
		delete(s.data.TTLs, ax.Key)
	case api.Put:
		s.data.KV[ax.Key] = &libkvstore.KVPair{Key: ax.Key, Value: ax.Value, LastIndex: l.Index}
		if ax.TTL != 0 {
			s.ttlLock.Lock()
			s.data.TTLs[ax.Key] = &ttl{CreateTime: time.Now(), TTL: ax.TTL, CreateIndex: l.Index}
			s.ttlLock.Unlock()
		}
	case api.DeleteTree:
		for k := range s.data.KV {
			if !strings.HasPrefix(k, ax.Key) {
				continue
			}
			delete(s.data.KV, k)
			delete(s.data.TTLs, k)
		}
	case reapKeys:
		s.mu.Lock()
		for _, k := range ax.Args {
			delete(s.data.KV, k)
			delete(s.data.TTLs, k)
		}
		s.mu.Unlock()
	default:
		return fmt.Errorf("unknown api.Request")
	}
	return nil
}

func (s *storeFSM) Snapshot() (raft.FSMSnapshot, error) {
	return s, nil
}

func (s *storeFSM) Restore(r io.ReadCloser) error {
	defer r.Close()
	s.data = newDB()
	return api.Decode(s.data, r)
}

func (s *storeFSM) Persist(sink raft.SnapshotSink) error {
	defer sink.Close()
	s.mu.Lock()
	defer s.mu.Unlock()
	return api.Encode(s.data, sink)
}

func (*storeFSM) Release() {}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/ChronixDB/chronix.go/chronix"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/storage/local/chunk"
)

type chronixStore struct {
	chronix chronix.Client
}

func (s *chronixStore) Put(metric model.Metric, descs []*chunk.Desc) error {
	for _, desc := range descs {
		ts, err := transcodeChunk(metric, desc)
		if err != nil {
			return fmt.Errorf("error transcoding chunk: %v", err)
		}
		if err := s.chronix.Store([]*chronix.TimeSeries{ts}, false); err != nil {
			return fmt.Errorf("error storing chunk: %v", err)
		}
	}
	return nil
}

func transcodeChunk(metric model.Metric, desc *chunk.Desc) (*chronix.TimeSeries, error) {
	ts := &chronix.TimeSeries{
		Metric:     string(metric[model.MetricNameLabel]),
		Attributes: map[string]string{},
	}
	for k, v := range metric {
		if k == model.MetricNameLabel {
			continue
		}
		ts.Attributes[string(k)] = string(v)
	}

	it := desc.C.NewIterator()
	for it.Scan() {
		sp := it.Value()
		ts.Points = append(ts.Points, chronix.Point{
			Value:     float64(sp.Value),
			Timestamp: sp.Timestamp.Unix(),
		})
	}
	if it.Err() != nil {
		return nil, it.Err()
	}
	return ts, nil
}
<commit_msg>Store Chronix timestamps as milliseconds<commit_after>package main

import (
	"fmt"

	"github.com/ChronixDB/chronix.go/chronix"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/storage/local/chunk"
)

type chronixStore struct {
	chronix chronix.Client
}

func (s *chronixStore) Put(metric model.Metric, descs []*chunk.Desc) error {
	for _, desc := range descs {
		ts, err := transcodeChunk(metric, desc)
		if err != nil {
			return fmt.Errorf("error transcoding chunk: %v", err)
		}
		if err := s.chronix.Store([]*chronix.TimeSeries{ts}, false); err != nil {
			return fmt.Errorf("error storing chunk: %v", err)
		}
	}
	return nil
}

func transcodeChunk(metric model.Metric, desc *chunk.Desc) (*chronix.TimeSeries, error) {
	ts := &chronix.TimeSeries{
		Metric:     string(metric[model.MetricNameLabel]),
		Attributes: map[string]string{},
	}
	for k, v := range metric {
		if k == model.MetricNameLabel {
			continue
		}
		ts.Attributes[string(k)] = string(v)
	}

	it := desc.C.NewIterator()
	for it.Scan() {
		sp := it.Value()
		ts.Points = append(ts.Points, chronix.Point{
			Value:     float64(sp.Value),
			Timestamp: sp.Timestamp.UnixNano() / 1e6,
		})
	}
	if it.Err() != nil {
		return nil, it.Err()
	}
	return ts, nil
}
<|endoftext|>