<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package schedulercache

import (
	"fmt"
	"sync"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/wait"
)

var (
	cleanAssumedPeriod = 1 * time.Second
)

// New returns a Cache implementation.
// It automatically starts a go routine that manages expiration of assumed pods.
// "ttl" is how long the assumed pod will get expired.
// "stop" is the channel that would close the background goroutine.
func New(ttl time.Duration, stop <-chan struct{}) Cache {
	cache := newSchedulerCache(ttl, cleanAssumedPeriod, stop)
	cache.run()
	return cache
}

type schedulerCache struct {
	stop   <-chan struct{}
	ttl    time.Duration
	period time.Duration

	// This mutex guards all fields within this cache struct.
	mu sync.Mutex
	// a set of assumed pod keys.
	// The key could further be used to get an entry in podStates.
	assumedPods map[string]bool
	// a map from pod key to podState.
	podStates map[string]*podState
	nodes     map[string]*NodeInfo
}

type podState struct {
	pod *api.Pod
	// Used by assumedPod to determinate expiration.
	deadline *time.Time
}

func newSchedulerCache(ttl, period time.Duration, stop <-chan struct{}) *schedulerCache {
	return &schedulerCache{
		ttl:    ttl,
		period: period,
		stop:   stop,

		nodes:       make(map[string]*NodeInfo),
		assumedPods: make(map[string]bool),
		podStates:   make(map[string]*podState),
	}
}

func (cache *schedulerCache) UpdateNodeNameToInfoMap(nodeNameToInfo map[string]*NodeInfo) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()
	for name, info := range cache.nodes {
		if current, ok := nodeNameToInfo[name]; !ok || current.generation != info.generation {
			nodeNameToInfo[name] = info.Clone()
		}
	}
	for name := range nodeNameToInfo {
		if _, ok := cache.nodes[name]; !ok {
			delete(nodeNameToInfo, name)
		}
	}
	return nil
}

func (cache *schedulerCache) List(selector labels.Selector) ([]*api.Pod, error) {
	cache.mu.Lock()
	defer cache.mu.Unlock()
	var pods []*api.Pod
	for _, info := range cache.nodes {
		for _, pod := range info.pods {
			if selector.Matches(labels.Set(pod.Labels)) {
				pods = append(pods, pod)
			}
		}
	}
	return pods, nil
}

func (cache *schedulerCache) AssumePod(pod *api.Pod) error {
	return cache.assumePod(pod, time.Now())
}

// assumePod exists for making test deterministic by taking time as input argument.
func (cache *schedulerCache) assumePod(pod *api.Pod, now time.Time) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key, err := getPodKey(pod)
	if err != nil {
		return err
	}
	if _, ok := cache.podStates[key]; ok {
		return fmt.Errorf("pod state wasn't initial but get assumed. Pod key: %v", key)
	}

	cache.addPod(pod)
	dl := now.Add(cache.ttl)
	ps := &podState{
		pod:      pod,
		deadline: &dl,
	}
	cache.podStates[key] = ps
	cache.assumedPods[key] = true
	return nil
}

func (cache *schedulerCache) ForgetPod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	// Only assumed pod can be forgotten.
	case ok && cache.assumedPods[key]:
		err := cache.removePod(pod)
		if err != nil {
			return err
		}
		delete(cache.assumedPods, key)
		delete(cache.podStates, key)
	default:
		return fmt.Errorf("pod state wasn't assumed but get forgotten. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) AddPod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	case ok && cache.assumedPods[key]:
		delete(cache.assumedPods, key)
		cache.podStates[key].deadline = nil
	case !ok:
		// Pod was expired. We should add it back.
		cache.addPod(pod)
		ps := &podState{
			pod: pod,
		}
		cache.podStates[key] = ps
	default:
		return fmt.Errorf("pod was already in added state. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) UpdatePod(oldPod, newPod *api.Pod) error {
	key, err := getPodKey(oldPod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	// An assumed pod won't have Update/Remove event. It needs to have Add event
	// before Update event, in which case the state would change from Assumed to Added.
	case ok && !cache.assumedPods[key]:
		if err := cache.updatePod(oldPod, newPod); err != nil {
			return err
		}
	default:
		return fmt.Errorf("pod state wasn't added but get updated. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) updatePod(oldPod, newPod *api.Pod) error {
	if err := cache.removePod(oldPod); err != nil {
		return err
	}
	cache.addPod(newPod)
	return nil
}

func (cache *schedulerCache) addPod(pod *api.Pod) {
	n, ok := cache.nodes[pod.Spec.NodeName]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[pod.Spec.NodeName] = n
	}
	n.addPod(pod)
}

func (cache *schedulerCache) removePod(pod *api.Pod) error {
	n := cache.nodes[pod.Spec.NodeName]
	if err := n.removePod(pod); err != nil {
		return err
	}
	if len(n.pods) == 0 && n.node == nil {
		delete(cache.nodes, pod.Spec.NodeName)
	}
	return nil
}

func (cache *schedulerCache) RemovePod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	// An assumed pod won't have Delete/Remove event. It needs to have Add event
	// before Remove event, in which case the state would change from Assumed to Added.
	case ok && !cache.assumedPods[key]:
		err := cache.removePod(pod)
		if err != nil {
			return err
		}
		delete(cache.podStates, key)
	default:
		return fmt.Errorf("pod state wasn't added but get removed. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) AddNode(node *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[node.Name]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[node.Name] = n
	}
	return n.SetNode(node)
}

func (cache *schedulerCache) UpdateNode(oldNode, newNode *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[newNode.Name]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[newNode.Name] = n
	}
	return n.SetNode(newNode)
}

func (cache *schedulerCache) RemoveNode(node *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n := cache.nodes[node.Name]
	if err := n.RemoveNode(node); err != nil {
		return err
	}
	// We remove NodeInfo for this node only if there aren't any pods on this node.
	// We can't do it unconditionally, because notifications about pods are delivered
	// in a different watch, and thus can potentially be observed later, even though
	// they happened before node removal.
	if len(n.pods) == 0 && n.node == nil {
		delete(cache.nodes, node.Name)
	}
	return nil
}

func (cache *schedulerCache) run() {
	go wait.Until(cache.cleanupExpiredAssumedPods, cache.period, cache.stop)
}

func (cache *schedulerCache) cleanupExpiredAssumedPods() {
	cache.cleanupAssumedPods(time.Now())
}

// cleanupAssumedPods exists for making test deterministic by taking time as input argument.
func (cache *schedulerCache) cleanupAssumedPods(now time.Time) {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	// The size of assumedPods should be small
	for key := range cache.assumedPods {
		ps, ok := cache.podStates[key]
		if !ok {
			panic("Key found in assumed set but not in podStates. Potentially a logical error.")
		}
		if now.After(*ps.deadline) {
			if err := cache.expirePod(key, ps); err != nil {
				glog.Errorf(" expirePod failed for %s: %v", key, err)
			}
		}
	}
}

func (cache *schedulerCache) expirePod(key string, ps *podState) error {
	if err := cache.removePod(ps.pod); err != nil {
		return err
	}
	delete(cache.assumedPods, key)
	delete(cache.podStates, key)
	return nil
}
<commit_msg>scheduler: cache.delete deletes the pod from node specified in the cached state<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

package schedulercache

import (
	"fmt"
	"sync"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/wait"
)

var (
	cleanAssumedPeriod = 1 * time.Second
)

// New returns a Cache implementation.
// It automatically starts a go routine that manages expiration of assumed pods.
// "ttl" is how long the assumed pod will get expired.
// "stop" is the channel that would close the background goroutine.
func New(ttl time.Duration, stop <-chan struct{}) Cache {
	cache := newSchedulerCache(ttl, cleanAssumedPeriod, stop)
	cache.run()
	return cache
}

type schedulerCache struct {
	stop   <-chan struct{}
	ttl    time.Duration
	period time.Duration

	// This mutex guards all fields within this cache struct.
	mu sync.Mutex
	// a set of assumed pod keys.
	// The key could further be used to get an entry in podStates.
	assumedPods map[string]bool
	// a map from pod key to podState.
	podStates map[string]*podState
	nodes     map[string]*NodeInfo
}

type podState struct {
	pod *api.Pod
	// Used by assumedPod to determinate expiration.
	deadline *time.Time
}

func newSchedulerCache(ttl, period time.Duration, stop <-chan struct{}) *schedulerCache {
	return &schedulerCache{
		ttl:    ttl,
		period: period,
		stop:   stop,

		nodes:       make(map[string]*NodeInfo),
		assumedPods: make(map[string]bool),
		podStates:   make(map[string]*podState),
	}
}

func (cache *schedulerCache) UpdateNodeNameToInfoMap(nodeNameToInfo map[string]*NodeInfo) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()
	for name, info := range cache.nodes {
		if current, ok := nodeNameToInfo[name]; !ok || current.generation != info.generation {
			nodeNameToInfo[name] = info.Clone()
		}
	}
	for name := range nodeNameToInfo {
		if _, ok := cache.nodes[name]; !ok {
			delete(nodeNameToInfo, name)
		}
	}
	return nil
}

func (cache *schedulerCache) List(selector labels.Selector) ([]*api.Pod, error) {
	cache.mu.Lock()
	defer cache.mu.Unlock()
	var pods []*api.Pod
	for _, info := range cache.nodes {
		for _, pod := range info.pods {
			if selector.Matches(labels.Set(pod.Labels)) {
				pods = append(pods, pod)
			}
		}
	}
	return pods, nil
}

func (cache *schedulerCache) AssumePod(pod *api.Pod) error {
	return cache.assumePod(pod, time.Now())
}

// assumePod exists for making test deterministic by taking time as input argument.
func (cache *schedulerCache) assumePod(pod *api.Pod, now time.Time) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	key, err := getPodKey(pod)
	if err != nil {
		return err
	}
	if _, ok := cache.podStates[key]; ok {
		return fmt.Errorf("pod state wasn't initial but get assumed. Pod key: %v", key)
	}

	cache.addPod(pod)
	dl := now.Add(cache.ttl)
	ps := &podState{
		pod:      pod,
		deadline: &dl,
	}
	cache.podStates[key] = ps
	cache.assumedPods[key] = true
	return nil
}

func (cache *schedulerCache) ForgetPod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	// Only assumed pod can be forgotten.
	case ok && cache.assumedPods[key]:
		err := cache.removePod(pod)
		if err != nil {
			return err
		}
		delete(cache.assumedPods, key)
		delete(cache.podStates, key)
	default:
		return fmt.Errorf("pod state wasn't assumed but get forgotten. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) AddPod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	case ok && cache.assumedPods[key]:
		delete(cache.assumedPods, key)
		cache.podStates[key].deadline = nil
	case !ok:
		// Pod was expired. We should add it back.
		cache.addPod(pod)
		ps := &podState{
			pod: pod,
		}
		cache.podStates[key] = ps
	default:
		return fmt.Errorf("pod was already in added state. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) UpdatePod(oldPod, newPod *api.Pod) error {
	key, err := getPodKey(oldPod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	_, ok := cache.podStates[key]
	switch {
	// An assumed pod won't have Update/Remove event. It needs to have Add event
	// before Update event, in which case the state would change from Assumed to Added.
	case ok && !cache.assumedPods[key]:
		if err := cache.updatePod(oldPod, newPod); err != nil {
			return err
		}
	default:
		return fmt.Errorf("pod state wasn't added but get updated. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) updatePod(oldPod, newPod *api.Pod) error {
	if err := cache.removePod(oldPod); err != nil {
		return err
	}
	cache.addPod(newPod)
	return nil
}

func (cache *schedulerCache) addPod(pod *api.Pod) {
	n, ok := cache.nodes[pod.Spec.NodeName]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[pod.Spec.NodeName] = n
	}
	n.addPod(pod)
}

func (cache *schedulerCache) removePod(pod *api.Pod) error {
	n := cache.nodes[pod.Spec.NodeName]
	if err := n.removePod(pod); err != nil {
		return err
	}
	if len(n.pods) == 0 && n.node == nil {
		delete(cache.nodes, pod.Spec.NodeName)
	}
	return nil
}

func (cache *schedulerCache) RemovePod(pod *api.Pod) error {
	key, err := getPodKey(pod)
	if err != nil {
		return err
	}

	cache.mu.Lock()
	defer cache.mu.Unlock()

	cachedstate, ok := cache.podStates[key]
	switch {
	// An assumed pod won't have Delete/Remove event. It needs to have Add event
	// before Remove event, in which case the state would change from Assumed to Added.
	case ok && !cache.assumedPods[key]:
		err := cache.removePod(cachedstate.pod)
		if err != nil {
			return err
		}
		delete(cache.podStates, key)
	default:
		return fmt.Errorf("pod state wasn't added but get removed. Pod key: %v", key)
	}
	return nil
}

func (cache *schedulerCache) AddNode(node *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[node.Name]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[node.Name] = n
	}
	return n.SetNode(node)
}

func (cache *schedulerCache) UpdateNode(oldNode, newNode *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n, ok := cache.nodes[newNode.Name]
	if !ok {
		n = NewNodeInfo()
		cache.nodes[newNode.Name] = n
	}
	return n.SetNode(newNode)
}

func (cache *schedulerCache) RemoveNode(node *api.Node) error {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	n := cache.nodes[node.Name]
	if err := n.RemoveNode(node); err != nil {
		return err
	}
	// We remove NodeInfo for this node only if there aren't any pods on this node.
	// We can't do it unconditionally, because notifications about pods are delivered
	// in a different watch, and thus can potentially be observed later, even though
	// they happened before node removal.
	if len(n.pods) == 0 && n.node == nil {
		delete(cache.nodes, node.Name)
	}
	return nil
}

func (cache *schedulerCache) run() {
	go wait.Until(cache.cleanupExpiredAssumedPods, cache.period, cache.stop)
}

func (cache *schedulerCache) cleanupExpiredAssumedPods() {
	cache.cleanupAssumedPods(time.Now())
}

// cleanupAssumedPods exists for making test deterministic by taking time as input argument.
func (cache *schedulerCache) cleanupAssumedPods(now time.Time) {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	// The size of assumedPods should be small
	for key := range cache.assumedPods {
		ps, ok := cache.podStates[key]
		if !ok {
			panic("Key found in assumed set but not in podStates. Potentially a logical error.")
		}
		if now.After(*ps.deadline) {
			if err := cache.expirePod(key, ps); err != nil {
				glog.Errorf(" expirePod failed for %s: %v", key, err)
			}
		}
	}
}

func (cache *schedulerCache) expirePod(key string, ps *podState) error {
	if err := cache.removePod(ps.pod); err != nil {
		return err
	}
	delete(cache.assumedPods, key)
	delete(cache.podStates, key)
	return nil
}
<|endoftext|>


<commit_before>// +build bench

package restxml_test

import (
	"testing"

	"bytes"
	"encoding/xml"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/awstesting"
	"github.com/aws/aws-sdk-go/private/protocol/restxml"
	"github.com/aws/aws-sdk-go/service/cloudfront"
)

func BenchmarkRESTXMLBuild_Complex_cloudfrontCreateDistribution(b *testing.B) {
	params := restxmlBuildCreateDistroParms

	op := &request.Operation{
		Name:       "CreateDistribution",
		HTTPMethod: "POST",
		HTTPPath:   "/2015-04-17/distribution/{DistributionId}/invalidation",
	}

	benchRESTXMLBuild(b, op, params)
}

func BenchmarkRESTXMLBuild_Simple_cloudfrontDeleteStreamingDistribution(b *testing.B) {
	params := &cloudfront.DeleteDistributionInput{
		Id:      aws.String("string"), // Required
		IfMatch: aws.String("string"),
	}
	op := &request.Operation{
		Name:       "DeleteStreamingDistribution",
		HTTPMethod: "DELETE",
		HTTPPath:   "/2015-04-17/streaming-distribution/{Id}",
	}
	benchRESTXMLBuild(b, op, params)
}

func BenchmarkEncodingXMLMarshal_Simple_cloudfrontDeleteStreamingDistribution(b *testing.B) {
	params := &cloudfront.DeleteDistributionInput{
		Id:      aws.String("string"), // Required
		IfMatch: aws.String("string"),
	}

	for i := 0; i < b.N; i++ {
		buf := &bytes.Buffer{}
		encoder := xml.NewEncoder(buf)
		if err := encoder.Encode(params); err != nil {
			b.Fatal("Unexpected error", err)
		}
	}
}

func benchRESTXMLBuild(b *testing.B, op *request.Operation, params interface{}) {
	svc := awstesting.NewClient()
	svc.ServiceName = "cloudfront"
	svc.APIVersion = "2015-04-17"

	for i := 0; i < b.N; i++ {
		r := svc.NewRequest(op, params, nil)
		restxml.Build(r)
		if r.Error != nil {
			b.Fatal("Unexpected error", r.Error)
		}
	}
}

var restxmlBuildCreateDistroParms = &cloudfront.CreateDistributionInput{
	DistributionConfig: &cloudfront.DistributionConfig{ // Required
		CallerReference: aws.String("string"), // Required
		Comment:         aws.String("string"), // Required
		DefaultCacheBehavior: &cloudfront.DefaultCacheBehavior{ // Required
			ForwardedValues: &cloudfront.ForwardedValues{ // Required
				Cookies: &cloudfront.CookiePreference{ // Required
					Forward: aws.String("ItemSelection"), // Required
					WhitelistedNames: &cloudfront.CookieNames{
						Quantity: aws.Int64(1), // Required
						Items: []*string{
							aws.String("string"), // Required
							// More values...
						},
					},
				},
				QueryString: aws.Bool(true), // Required
				Headers: &cloudfront.Headers{
					Quantity: aws.Int64(1), // Required
					Items: []*string{
						aws.String("string"), // Required
						// More values...
					},
				},
			},
			MinTTL:         aws.Int64(1),         // Required
			TargetOriginId: aws.String("string"), // Required
			TrustedSigners: &cloudfront.TrustedSigners{ // Required
				Enabled:  aws.Bool(true), // Required
				Quantity: aws.Int64(1),   // Required
				Items: []*string{
					aws.String("string"), // Required
					// More values...
				},
			},
			ViewerProtocolPolicy: aws.String("ViewerProtocolPolicy"), // Required
			AllowedMethods: &cloudfront.AllowedMethods{
				Items: []*string{ // Required
					aws.String("Method"), // Required
					// More values...
				},
				Quantity: aws.Int64(1), // Required
				CachedMethods: &cloudfront.CachedMethods{
					Items: []*string{ // Required
						aws.String("Method"), // Required
						// More values...
					},
					Quantity: aws.Int64(1), // Required
				},
			},
			DefaultTTL:      aws.Int64(1),
			MaxTTL:          aws.Int64(1),
			SmoothStreaming: aws.Bool(true),
		},
		Enabled: aws.Bool(true), // Required
		Origins: &cloudfront.Origins{ // Required
			Quantity: aws.Int64(1), // Required
			Items: []*cloudfront.Origin{
				{ // Required
					DomainName: aws.String("string"), // Required
					Id:         aws.String("string"), // Required
					CustomOriginConfig: &cloudfront.CustomOriginConfig{
						HTTPPort:             aws.Int64(1),                       // Required
						HTTPSPort:            aws.Int64(1),                       // Required
						OriginProtocolPolicy: aws.String("OriginProtocolPolicy"), // Required
					},
					OriginPath: aws.String("string"),
					S3OriginConfig: &cloudfront.S3OriginConfig{
						OriginAccessIdentity: aws.String("string"), // Required
					},
				},
				// More values...
			},
		},
		Aliases: &cloudfront.Aliases{
			Quantity: aws.Int64(1), // Required
			Items: []*string{
				aws.String("string"), // Required
				// More values...
			},
		},
		CacheBehaviors: &cloudfront.CacheBehaviors{
			Quantity: aws.Int64(1), // Required
			Items: []*cloudfront.CacheBehavior{
				{ // Required
					ForwardedValues: &cloudfront.ForwardedValues{ // Required
						Cookies: &cloudfront.CookiePreference{ // Required
							Forward: aws.String("ItemSelection"), // Required
							WhitelistedNames: &cloudfront.CookieNames{
								Quantity: aws.Int64(1), // Required
								Items: []*string{
									aws.String("string"), // Required
									// More values...
								},
							},
						},
						QueryString: aws.Bool(true), // Required
						Headers: &cloudfront.Headers{
							Quantity: aws.Int64(1), // Required
							Items: []*string{
								aws.String("string"), // Required
								// More values...
							},
						},
					},
					MinTTL:         aws.Int64(1),         // Required
					PathPattern:    aws.String("string"), // Required
					TargetOriginId: aws.String("string"), // Required
					TrustedSigners: &cloudfront.TrustedSigners{ // Required
						Enabled:  aws.Bool(true), // Required
						Quantity: aws.Int64(1),   // Required
						Items: []*string{
							aws.String("string"), // Required
							// More values...
						},
					},
					ViewerProtocolPolicy: aws.String("ViewerProtocolPolicy"), // Required
					AllowedMethods: &cloudfront.AllowedMethods{
						Items: []*string{ // Required
							aws.String("Method"), // Required
							// More values...
						},
						Quantity: aws.Int64(1), // Required
						CachedMethods: &cloudfront.CachedMethods{
							Items: []*string{ // Required
								aws.String("Method"), // Required
								// More values...
							},
							Quantity: aws.Int64(1), // Required
						},
					},
					DefaultTTL:      aws.Int64(1),
					MaxTTL:          aws.Int64(1),
					SmoothStreaming: aws.Bool(true),
				},
				// More values...
			},
		},
		CustomErrorResponses: &cloudfront.CustomErrorResponses{
			Quantity: aws.Int64(1), // Required
			Items: []*cloudfront.CustomErrorResponse{
				{ // Required
					ErrorCode:          aws.Int64(1), // Required
					ErrorCachingMinTTL: aws.Int64(1),
					ResponseCode:       aws.String("string"),
					ResponsePagePath:   aws.String("string"),
				},
				// More values...
			},
		},
		DefaultRootObject: aws.String("string"),
		Logging: &cloudfront.LoggingConfig{
			Bucket:         aws.String("string"), // Required
			Enabled:        aws.Bool(true),       // Required
			IncludeCookies: aws.Bool(true),       // Required
			Prefix:         aws.String("string"), // Required
		},
		PriceClass: aws.String("PriceClass"),
		Restrictions: &cloudfront.Restrictions{
			GeoRestriction: &cloudfront.GeoRestriction{ // Required
				Quantity:        aws.Int64(1),                     // Required
				RestrictionType: aws.String("GeoRestrictionType"), // Required
				Items: []*string{
					aws.String("string"), // Required
					// More values...
				},
			},
		},
		ViewerCertificate: &cloudfront.ViewerCertificate{
			CloudFrontDefaultCertificate: aws.Bool(true),
			IAMCertificateId:             aws.String("string"),
			MinimumProtocolVersion:       aws.String("MinimumProtocolVersion"),
			SSLSupportMethod:             aws.String("SSLSupportMethod"),
		},
	},
}
<commit_msg>private/protocol/restxml: Cleanup and add new benchmarks (#1408)<commit_after>// +build bench

package restxml_test

import (
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"bytes"
	"encoding/xml"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/endpoints"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/private/protocol/restxml"
	"github.com/aws/aws-sdk-go/service/cloudfront"
	"github.com/aws/aws-sdk-go/service/s3"
)

var (
	cloudfrontSvc *cloudfront.CloudFront
	s3Svc         *s3.S3
)

func TestMain(m *testing.M) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))

	sess := session.Must(session.NewSession(&aws.Config{
		Credentials:      credentials.NewStaticCredentials("Key", "Secret", "Token"),
		Endpoint:         aws.String(server.URL),
		S3ForcePathStyle: aws.Bool(true),
		DisableSSL:       aws.Bool(true),
		Region:           aws.String(endpoints.UsWest2RegionID),
	}))
	cloudfrontSvc = cloudfront.New(sess)
	s3Svc = s3.New(sess)

	c := m.Run()
	server.Close()
	os.Exit(c)
}

func BenchmarkRESTXMLBuild_Complex_CFCreateDistro(b *testing.B) {
	params := cloudfrontCreateDistributionInput()

	benchRESTXMLBuild(b, func() *request.Request {
		req, _ := cloudfrontSvc.CreateDistributionRequest(params)
		return req
	})
}

func BenchmarkRESTXMLRequest_Complex_CFCreateDistro(b *testing.B) {
	benchRESTXMLRequest(b, func() *request.Request {
		req, _ := cloudfrontSvc.CreateDistributionRequest(cloudfrontCreateDistributionInput())
		return req
	})
}

func BenchmarkRESTXMLBuild_Simple_CFDeleteDistro(b *testing.B) {
	params := cloudfrontDeleteDistributionInput()

	benchRESTXMLBuild(b, func() *request.Request {
		req, _ := cloudfrontSvc.DeleteDistributionRequest(params)
		return req
	})
}

func BenchmarkRESTXMLRequest_Simple_CFDeleteDistro(b *testing.B) {
	benchRESTXMLRequest(b, func() *request.Request {
		req, _ := cloudfrontSvc.DeleteDistributionRequest(cloudfrontDeleteDistributionInput())
		return req
	})
}

func BenchmarkRESTXMLBuild_REST_S3HeadObject(b *testing.B) {
	params := s3HeadObjectInput()

	benchRESTXMLBuild(b, func() *request.Request {
		req, _ := s3Svc.HeadObjectRequest(params)
		return req
	})
}

func BenchmarkRESTXMLRequest_REST_S3HeadObject(b *testing.B) {
	benchRESTXMLRequest(b, func() *request.Request {
		req, _ := s3Svc.HeadObjectRequest(s3HeadObjectInput())
		return req
	})
}

func BenchmarkRESTXMLBuild_XML_S3PutObjectAcl(b *testing.B) {
	params := s3PutObjectAclInput()

	benchRESTXMLBuild(b, func() *request.Request {
		req, _ := s3Svc.PutObjectAclRequest(params)
		return req
	})
}

func BenchmarkRESTXMLRequest_XML_S3PutObjectAcl(b *testing.B) {
	benchRESTXMLRequest(b, func() *request.Request {
		req, _ := s3Svc.PutObjectAclRequest(s3PutObjectAclInput())
		return req
	})
}

func BenchmarkEncodingXML_Simple(b *testing.B) {
	params := cloudfrontDeleteDistributionInput()

	for i := 0; i < b.N; i++ {
		buf := &bytes.Buffer{}
		encoder := xml.NewEncoder(buf)
		if err := encoder.Encode(params); err != nil {
			b.Fatal("Unexpected error", err)
		}
	}
}

func benchRESTXMLBuild(b *testing.B, reqFn func() *request.Request) {
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		req := reqFn()
		restxml.Build(req)
		if req.Error != nil {
			b.Fatal("Unexpected error", req.Error)
		}
	}
}

func benchRESTXMLRequest(b *testing.B, reqFn func() *request.Request) {
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		err := reqFn().Send()
		if err != nil {
			b.Fatal("Unexpected error", err)
		}
	}
}

func cloudfrontCreateDistributionInput() *cloudfront.CreateDistributionInput {
	return &cloudfront.CreateDistributionInput{
		DistributionConfig: &cloudfront.DistributionConfig{ // Required
			CallerReference: aws.String("string"), // Required
			Comment:         aws.String("string"), // Required
			DefaultCacheBehavior: &cloudfront.DefaultCacheBehavior{ // Required
				ForwardedValues: &cloudfront.ForwardedValues{ // Required
					Cookies: &cloudfront.CookiePreference{ // Required
						Forward: aws.String("ItemSelection"), // Required
						WhitelistedNames: &cloudfront.CookieNames{
							Quantity: aws.Int64(1), // Required
							Items: []*string{
								aws.String("string"), // Required
								// More values...
							},
						},
					},
					QueryString: aws.Bool(true), // Required
					Headers: &cloudfront.Headers{
						Quantity: aws.Int64(1), // Required
						Items: []*string{
							aws.String("string"), // Required
							// More values...
						},
					},
				},
				MinTTL:         aws.Int64(1),         // Required
				TargetOriginId: aws.String("string"), // Required
				TrustedSigners: &cloudfront.TrustedSigners{ // Required
					Enabled:  aws.Bool(true), // Required
					Quantity: aws.Int64(1),   // Required
					Items: []*string{
						aws.String("string"), // Required
						// More values...
					},
				},
				ViewerProtocolPolicy: aws.String("ViewerProtocolPolicy"), // Required
				AllowedMethods: &cloudfront.AllowedMethods{
					Items: []*string{ // Required
						aws.String("Method"), // Required
						// More values...
					},
					Quantity: aws.Int64(1), // Required
					CachedMethods: &cloudfront.CachedMethods{
						Items: []*string{ // Required
							aws.String("Method"), // Required
							// More values...
						},
						Quantity: aws.Int64(1), // Required
					},
				},
				DefaultTTL:      aws.Int64(1),
				MaxTTL:          aws.Int64(1),
				SmoothStreaming: aws.Bool(true),
			},
			Enabled: aws.Bool(true), // Required
			Origins: &cloudfront.Origins{ // Required
				Quantity: aws.Int64(1), // Required
				Items: []*cloudfront.Origin{
					{ // Required
						DomainName: aws.String("string"), // Required
						Id:         aws.String("string"), // Required
						CustomOriginConfig: &cloudfront.CustomOriginConfig{
							HTTPPort:             aws.Int64(1),                       // Required
							HTTPSPort:            aws.Int64(1),                       // Required
							OriginProtocolPolicy: aws.String("OriginProtocolPolicy"), // Required
						},
						OriginPath: aws.String("string"),
						S3OriginConfig: &cloudfront.S3OriginConfig{
							OriginAccessIdentity: aws.String("string"), // Required
						},
					},
					// More values...
				},
			},
			Aliases: &cloudfront.Aliases{
				Quantity: aws.Int64(1), // Required
				Items: []*string{
					aws.String("string"), // Required
					// More values...
				},
			},
			CacheBehaviors: &cloudfront.CacheBehaviors{
				Quantity: aws.Int64(1), // Required
				Items: []*cloudfront.CacheBehavior{
					{ // Required
						ForwardedValues: &cloudfront.ForwardedValues{ // Required
							Cookies: &cloudfront.CookiePreference{ // Required
								Forward: aws.String("ItemSelection"), // Required
								WhitelistedNames: &cloudfront.CookieNames{
									Quantity: aws.Int64(1), // Required
									Items: []*string{
										aws.String("string"), // Required
										// More values...
									},
								},
							},
							QueryString: aws.Bool(true), // Required
							Headers: &cloudfront.Headers{
								Quantity: aws.Int64(1), // Required
								Items: []*string{
									aws.String("string"), // Required
									// More values...
								},
							},
						},
						MinTTL:         aws.Int64(1),         // Required
						PathPattern:    aws.String("string"), // Required
						TargetOriginId: aws.String("string"), // Required
						TrustedSigners: &cloudfront.TrustedSigners{ // Required
							Enabled:  aws.Bool(true), // Required
							Quantity: aws.Int64(1),   // Required
							Items: []*string{
								aws.String("string"), // Required
								// More values...
							},
						},
						ViewerProtocolPolicy: aws.String("ViewerProtocolPolicy"), // Required
						AllowedMethods: &cloudfront.AllowedMethods{
							Items: []*string{ // Required
								aws.String("Method"), // Required
								// More values...
							},
							Quantity: aws.Int64(1), // Required
							CachedMethods: &cloudfront.CachedMethods{
								Items: []*string{ // Required
									aws.String("Method"), // Required
									// More values...
								},
								Quantity: aws.Int64(1), // Required
							},
						},
						DefaultTTL:      aws.Int64(1),
						MaxTTL:          aws.Int64(1),
						SmoothStreaming: aws.Bool(true),
					},
					// More values...
				},
			},
			CustomErrorResponses: &cloudfront.CustomErrorResponses{
				Quantity: aws.Int64(1), // Required
				Items: []*cloudfront.CustomErrorResponse{
					{ // Required
						ErrorCode:          aws.Int64(1), // Required
						ErrorCachingMinTTL: aws.Int64(1),
						ResponseCode:       aws.String("string"),
						ResponsePagePath:   aws.String("string"),
					},
					// More values...
				},
			},
			DefaultRootObject: aws.String("string"),
			Logging: &cloudfront.LoggingConfig{
				Bucket:         aws.String("string"), // Required
				Enabled:        aws.Bool(true),       // Required
				IncludeCookies: aws.Bool(true),       // Required
				Prefix:         aws.String("string"), // Required
			},
			PriceClass: aws.String("PriceClass"),
			Restrictions: &cloudfront.Restrictions{
				GeoRestriction: &cloudfront.GeoRestriction{ // Required
					Quantity:        aws.Int64(1),                     // Required
					RestrictionType: aws.String("GeoRestrictionType"), // Required
					Items: []*string{
						aws.String("string"), // Required
						// More values...
					},
				},
			},
			ViewerCertificate: &cloudfront.ViewerCertificate{
				CloudFrontDefaultCertificate: aws.Bool(true),
				IAMCertificateId:             aws.String("string"),
				MinimumProtocolVersion:       aws.String("MinimumProtocolVersion"),
				SSLSupportMethod:             aws.String("SSLSupportMethod"),
			},
		},
	}
}

func cloudfrontDeleteDistributionInput() *cloudfront.DeleteDistributionInput {
	return &cloudfront.DeleteDistributionInput{
		Id:      aws.String("string"), // Required
		IfMatch: aws.String("string"),
	}
}

func s3HeadObjectInput() *s3.HeadObjectInput {
	return &s3.HeadObjectInput{
		Bucket:    aws.String("somebucketname"),
		Key:       aws.String("keyname"),
		VersionId: aws.String("someVersion"),
		IfMatch:   aws.String("IfMatch"),
	}
}

func s3PutObjectAclInput() *s3.PutObjectAclInput {
	return &s3.PutObjectAclInput{
		Bucket: aws.String("somebucketname"),
		Key:    aws.String("keyname"),
		AccessControlPolicy: &s3.AccessControlPolicy{
			Grants: []*s3.Grant{
				{
					Grantee: &s3.Grantee{
						DisplayName:  aws.String("someName"),
						EmailAddress: aws.String("someAddr"),
						ID:           aws.String("someID"),
						Type:         aws.String(s3.TypeCanonicalUser),
						URI:          aws.String("someURI"),
					},
					Permission: aws.String(s3.PermissionWrite),
				},
			},
			Owner: &s3.Owner{
				DisplayName: aws.String("howdy"),
				ID:          aws.String("someID"),
			},
		},
	}
}
<|endoftext|>


<commit_before>package transport

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"errors"
	"hash"
	"io"
	"strconv"
	"time"

	pond "github.com/unixninja92/Div-III-Server/protos"
	"github.com/golang/protobuf/proto"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/nacl/secretbox"
)

// blockSize is the size of the blocks of data that we'll send and receive when
// working in streaming mode. Each block is prefixed by two length bytes (which
// aren't counted in blockSize) and includes secretbox.Overhead bytes of MAC
// tag (which are).
const blockSize = 4096 - 2

type Conn struct {
	conn                     io.ReadWriteCloser
	isServer                 bool
	identity, identityPublic [32]byte
	Peer                     [32]byte

	writeKey, readKey           [32]byte
	writeKeyValid, readKeyValid bool
	writeSequence, readSequence [24]byte

	// readBuffer is used to receive bytes from the network when this Conn
	// is used to stream data.
	readBuffer []byte
	// decryptBuffer is used to store decrypted payloads when this Conn is
	// used to stream data and the caller's buffer isn't large enough to
	// decrypt into directly.
	decryptBuffer []byte
	// readPending aliases into decryptBuffer when a partial decryption had
	// to be returned to a caller because of buffer size limitations.
	readPending []byte

	// writeBuffer is used to hold encrypted payloads when this Conn is
	// used for streaming data.
	writeBuffer []byte
}

func NewServer(conn io.ReadWriteCloser, identity *[32]byte) *Conn {
	c := &Conn{
		conn:     conn,
		isServer: true,
	}
	copy(c.identity[:], identity[:])
	return c
}

func NewClient(conn io.ReadWriteCloser, myIdentity, myIdentityPublic, serverPublic *[32]byte) *Conn {
	c := &Conn{
		conn: conn,
	}
	copy(c.identity[:], myIdentity[:])
	copy(c.identityPublic[:], myIdentityPublic[:])
	copy(c.Peer[:], serverPublic[:])
	return c
}

func incSequence(seq *[24]byte) {
	n := uint32(1)

	for i := 0; i < 8; i++ {
		n += uint32(seq[i])
		seq[i] = byte(n)
		n >>= 8
	}
}

type deadlineable interface {
	SetDeadline(time.Time)
}

func (c *Conn) SetDeadline(t time.Time) {
	if d, ok := c.conn.(deadlineable); ok {
		d.SetDeadline(t)
	}
}

func (c *Conn) Read(out []byte) (n int, err error) {
	if len(c.readPending) > 0 {
		n = copy(out, c.readPending)
		c.readPending = c.readPending[n:]
		return
	}

	if c.readBuffer == nil {
		c.readBuffer = make([]byte, blockSize+2)
	}

	if _, err := io.ReadFull(c.conn, c.readBuffer[:2]); err != nil {
		return 0, err
	}
	n = int(c.readBuffer[0]) | int(c.readBuffer[1])<<8
	if n > len(c.readBuffer) {
		return 0, errors.New("transport: peer's message too large for Read")
	}
	if _, err := io.ReadFull(c.conn, c.readBuffer[:n]); err != nil {
		return 0, err
	}

	var ok bool
	if len(out) >= n-secretbox.Overhead {
		// We can decrypt directly into the output buffer.
		out, ok = secretbox.Open(out[:0], c.readBuffer[:n], &c.readSequence, &c.readKey)
		n = len(out)
	} else {
		// We need to decrypt into a side buffer and copy a prefix of
		// the result into the caller's buffer.
		c.decryptBuffer, ok = secretbox.Open(c.decryptBuffer[:0], c.readBuffer[:n], &c.readSequence, &c.readKey)
		n = copy(out, c.decryptBuffer)
		c.readPending = c.decryptBuffer[n:]
	}
	incSequence(&c.readSequence)
	if !ok {
		c.readPending = c.readPending[:0]
		return 0, errors.New("transport: bad MAC")
	}

	return
}

func (c *Conn) Write(buf []byte) (n int, err error) {
	if c.writeBuffer == nil {
		c.writeBuffer = make([]byte, blockSize+2)
	}

	for len(buf) > 0 {
		m := len(buf)
		if m > blockSize-secretbox.Overhead {
			m = blockSize - secretbox.Overhead
		}
		l := len(secretbox.Seal(c.writeBuffer[2:2], buf[:m], &c.writeSequence, &c.writeKey))
		c.writeBuffer[0] = byte(l)
		c.writeBuffer[1] = byte(l >> 8)
		if _, err = c.conn.Write(c.writeBuffer[:2+l]); err != nil {
			return n, err
		}
		n += m
		buf = buf[m:]
		incSequence(&c.writeSequence)
	}

	return
}

func (c *Conn) ReadProto(out proto.Message) error {
	buf := make([]byte, pond.TransportSize+3+secretbox.Overhead)
	n, err := c.read(buf)
	if err != nil {
		return err
	}
	if n != pond.TransportSize+3 {
		return errors.New("transport: message wrong length")
	}

	n = int(buf[0]) | int(buf[1])<<8 | int(buf[2])<<16 //| int(buf[3])<<32
	buf = buf[3:]
	if n > len(buf) {
		return errors.New("transport: corrupt message")
	}
	return proto.Unmarshal(buf[:n], out)
}

func (c *Conn) WriteProto(msg proto.Message) error {
	data, err := proto.Marshal(msg)
	if err != nil {
		return err
	}
	if len(data) > pond.TransportSize {
		return errors.New("transport: message too large")
	}

	buf := make([]byte, pond.TransportSize+3)
	buf[0] = byte(len(data))
	buf[1] = byte(len(data) >> 8)
	buf[2] = byte(len(data) >> 16)
	// buf[3] = byte(len(data) >> 32)
	copy(buf[3:], data)
	_, err = c.write(buf)
	return err
}

func (c *Conn) Close() (err error) {
	if !c.isServer {
		_, err = c.write(nil)
	}

	if closeErr := c.conn.Close(); err == nil {
		err = closeErr
	}

	return
}

func (c *Conn) WaitForClose() error {
	if !c.isServer {
		panic("non-server waited for connection close")
	}
	n, err := c.read(make([]byte, 128))
	if err != nil {
		return err
	}
	if n != 0 {
		return errors.New("transport: non-close message received when expecting close")
	}

	return nil
}

func (c *Conn) read(data []byte) (n int, err error) {
	var lengthBytes [3]byte

	if _, err := io.ReadFull(c.conn, lengthBytes[:]); err != nil {
		return 0, err
	}

	theirLength := int(lengthBytes[0]) + int(lengthBytes[1])<<8 + int(lengthBytes[2])<<16 //+ int(lengthBytes[3])<<32
	if theirLength > len(data) {
		return 0, errors.New("tranport: given buffer too small (" + strconv.Itoa(len(data)) + " vs " + strconv.Itoa(theirLength) + ")")
	}

	data = data[:theirLength]
	if _, err := io.ReadFull(c.conn, data); err != nil {
		return 0, err
	}

	decrypted, err := c.decrypt(data)
	if err != nil {
		return 0, err
	}
	copy(data, decrypted)

	return len(decrypted), nil
}

func (c *Conn) write(data []byte) (n int, err error) {
	encrypted := c.encrypt(data)

	var lengthBytes [3]byte
	lengthBytes[0] = byte(len(encrypted))
	lengthBytes[1] = byte(len(encrypted) >> 8)
	lengthBytes[2] = byte(len(encrypted) >> 16)
	// lengthBytes[3] = byte(len(encrypted) >> 32)

	if _, err := c.conn.Write(lengthBytes[:]); err != nil {
		return 0, err
	}
	if _, err := c.conn.Write(encrypted); err != nil {
		return 0, err
	}

	return len(data), nil
}

func (c *Conn) encrypt(data []byte) []byte {
	if !c.writeKeyValid {
		return data
	}

	encrypted := secretbox.Seal(nil, data, &c.writeSequence, &c.writeKey)
	incSequence(&c.writeSequence)
	return encrypted
}

func (c *Conn) decrypt(data []byte) ([]byte, error) {
	if !c.readKeyValid {
		return data, nil
	}

	decrypted, ok := secretbox.Open(nil, data, &c.readSequence, &c.readKey)
	incSequence(&c.readSequence)
	if !ok {
		return nil, errors.New("transport: bad MAC")
	}
	return decrypted, nil
}

var serverKeysMagic = []byte("server keys snap")
var clientKeysMagic = []byte("client keys snap")

func (c *Conn) setupKeys(ephemeralShared *[32]byte) {
	var writeMagic, readMagic []byte
	if c.isServer {
		writeMagic, readMagic = serverKeysMagic, clientKeysMagic
	} else {
		writeMagic, readMagic = clientKeysMagic, serverKeysMagic
	}

	h := sha256.New()
	h.Write(writeMagic)
	h.Write(ephemeralShared[:])
	h.Sum(c.writeKey[:0])
	c.writeKeyValid = true

	h.Reset()
	h.Write(readMagic)
	h.Write(ephemeralShared[:])
	h.Sum(c.readKey[:0])
	c.readKeyValid = true
}

var serverProofMagic = []byte("server proof snap")
var clientProofMagic = []byte("client proof snap")

var shortMessageError = errors.New("transport: received short handshake message")

func (c *Conn) Handshake() error {
	var ephemeralPrivate, ephemeralPublic, ephemeralShared [32]byte
	if _, err := io.ReadFull(rand.Reader, ephemeralPrivate[:]); err != nil {
		return err
	}
	curve25519.ScalarBaseMult(&ephemeralPublic, &ephemeralPrivate)

	if _, err := c.write(ephemeralPublic[:]); err != nil {
		return err
	}

	var theirEphemeralPublic [32]byte
	if n, err := c.read(theirEphemeralPublic[:]); err != nil || n != len(theirEphemeralPublic) {
		if err == nil {
			err = shortMessageError
		}
		return err
	}

	handshakeHash := sha256.New()
	if c.isServer {
		handshakeHash.Write(theirEphemeralPublic[:])
		handshakeHash.Write(ephemeralPublic[:])
	} else {
		handshakeHash.Write(ephemeralPublic[:])
		handshakeHash.Write(theirEphemeralPublic[:])
	}

	curve25519.ScalarMult(&ephemeralShared, &ephemeralPrivate, &theirEphemeralPublic)
	c.setupKeys(&ephemeralShared)

	if c.isServer {
		return c.handshakeServer(handshakeHash, &theirEphemeralPublic)
	}
	return c.handshakeClient(handshakeHash, &ephemeralPrivate)
}

func (c *Conn) handshakeClient(handshakeHash hash.Hash, ephemeralPrivate *[32]byte) error {
	var ephemeralIdentityShared [32]byte
	curve25519.ScalarMult(&ephemeralIdentityShared, ephemeralPrivate, &c.Peer)

	digest := handshakeHash.Sum(nil)
	h := hmac.New(sha256.New, ephemeralIdentityShared[:])
	h.Write(serverProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	digestReceived := make([]byte, len(digest)+secretbox.Overhead)
	n, err := c.read(digestReceived)
	if err != nil {
		return err
	}
	if n != len(digest) {
		return shortMessageError
	}
	digestReceived = digestReceived[:n]

	if subtle.ConstantTimeCompare(digest, digestReceived) != 1 {
		return errors.New("transport: server identity incorrect")
	}

	var identityShared [32]byte
	curve25519.ScalarMult(&identityShared, &c.identity, &c.Peer)

	handshakeHash.Write(digest)
	digest = handshakeHash.Sum(digest[:0])

	h = hmac.New(sha256.New, identityShared[:])
	h.Write(clientProofMagic)
	h.Write(digest)

	finalMessage := make([]byte, 32+sha256.Size)
	copy(finalMessage, c.identityPublic[:])
	h.Sum(finalMessage[32:32])

	if _, err := c.write(finalMessage); err != nil {
		return err
	}

	return nil
}

func (c *Conn) handshakeServer(handshakeHash hash.Hash, theirEphemeralPublic *[32]byte) error {
	var ephemeralIdentityShared [32]byte
	curve25519.ScalarMult(&ephemeralIdentityShared, &c.identity, theirEphemeralPublic)

	digest := handshakeHash.Sum(nil)
	h := hmac.New(sha256.New, ephemeralIdentityShared[:])
	h.Write(serverProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	if _, err := c.write(digest); err != nil {
		return err
	}

	handshakeHash.Write(digest)
	digest = handshakeHash.Sum(digest[:0])

	finalMessage := make([]byte, 32+sha256.Size+secretbox.Overhead)
	n, err := c.read(finalMessage)
	if err != nil {
		return err
	}
	if n != 32+sha256.Size {
		return shortMessageError
	}
	finalMessage = finalMessage[:n]

	copy(c.Peer[:], finalMessage[:32])
	var identityShared [32]byte
	curve25519.ScalarMult(&identityShared, &c.identity, &c.Peer)

	h = hmac.New(sha256.New, identityShared[:])
	h.Write(clientProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	if subtle.ConstantTimeCompare(digest, finalMessage[32:]) != 1 {
		return errors.New("transport: bad proof from client")
	}

	return nil
}
<commit_msg>Fixed transport lenght issues<commit_after>package transport

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"errors"
	"hash"
	"io"
	"strconv"
	"time"

	pond "github.com/unixninja92/Div-III-Server/protos"
	"github.com/golang/protobuf/proto"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/nacl/secretbox"
)

// blockSize is the size of the blocks of data that we'll send and receive when
// working in streaming mode. Each block is prefixed by two length bytes (which
// aren't counted in blockSize) and includes secretbox.Overhead bytes of MAC
// tag (which are).
const blockSize = 4096 - 2

type Conn struct {
	conn                     io.ReadWriteCloser
	isServer                 bool
	identity, identityPublic [32]byte
	Peer                     [32]byte

	writeKey, readKey           [32]byte
	writeKeyValid, readKeyValid bool
	writeSequence, readSequence [24]byte

	// readBuffer is used to receive bytes from the network when this Conn
	// is used to stream data.
	readBuffer []byte
	// decryptBuffer is used to store decrypted payloads when this Conn is
	// used to stream data and the caller's buffer isn't large enough to
	// decrypt into directly.
	decryptBuffer []byte
	// readPending aliases into decryptBuffer when a partial decryption had
	// to be returned to a caller because of buffer size limitations.
	readPending []byte

	// writeBuffer is used to hold encrypted payloads when this Conn is
	// used for streaming data.
	writeBuffer []byte
}

func NewServer(conn io.ReadWriteCloser, identity *[32]byte) *Conn {
	c := &Conn{
		conn:     conn,
		isServer: true,
	}
	copy(c.identity[:], identity[:])
	return c
}

func NewClient(conn io.ReadWriteCloser, myIdentity, myIdentityPublic, serverPublic *[32]byte) *Conn {
	c := &Conn{
		conn: conn,
	}
	copy(c.identity[:], myIdentity[:])
	copy(c.identityPublic[:], myIdentityPublic[:])
	copy(c.Peer[:], serverPublic[:])
	return c
}

func incSequence(seq *[24]byte) {
	n := uint32(1)

	for i := 0; i < 8; i++ {
		n += uint32(seq[i])
		seq[i] = byte(n)
		n >>= 8
	}
}

type deadlineable interface {
	SetDeadline(time.Time)
}

func (c *Conn) SetDeadline(t time.Time) {
	if d, ok := c.conn.(deadlineable); ok {
		d.SetDeadline(t)
	}
}

func (c *Conn) Read(out []byte) (n int, err error) {
	if len(c.readPending) > 0 {
		n = copy(out, c.readPending)
		c.readPending = c.readPending[n:]
		return
	}

	if c.readBuffer == nil {
		c.readBuffer = make([]byte, blockSize+2)
	}

	if _, err := io.ReadFull(c.conn, c.readBuffer[:2]); err != nil {
		return 0, err
	}
	n = int(c.readBuffer[0]) | int(c.readBuffer[1])<<8
	if n > len(c.readBuffer) {
		return 0, errors.New("transport: peer's message too large for Read")
	}
	if _, err := io.ReadFull(c.conn, c.readBuffer[:n]); err != nil {
		return 0, err
	}

	var ok bool
	if len(out) >= n-secretbox.Overhead {
		// We can decrypt directly into the output buffer.
		out, ok = secretbox.Open(out[:0], c.readBuffer[:n], &c.readSequence, &c.readKey)
		n = len(out)
	} else {
		// We need to decrypt into a side buffer and copy a prefix of
		// the result into the caller's buffer.
		c.decryptBuffer, ok = secretbox.Open(c.decryptBuffer[:0], c.readBuffer[:n], &c.readSequence, &c.readKey)
		n = copy(out, c.decryptBuffer)
		c.readPending = c.decryptBuffer[n:]
	}
	incSequence(&c.readSequence)
	if !ok {
		c.readPending = c.readPending[:0]
		return 0, errors.New("transport: bad MAC")
	}

	return
}

func (c *Conn) Write(buf []byte) (n int, err error) {
	if c.writeBuffer == nil {
		c.writeBuffer = make([]byte, blockSize+2)
	}

	for len(buf) > 0 {
		m := len(buf)
		if m > blockSize-secretbox.Overhead {
			m = blockSize - secretbox.Overhead
		}
		l := len(secretbox.Seal(c.writeBuffer[2:2], buf[:m], &c.writeSequence, &c.writeKey))
		c.writeBuffer[0] = byte(l)
		c.writeBuffer[1] = byte(l >> 8)
		if _, err = c.conn.Write(c.writeBuffer[:2+l]); err != nil {
			return n, err
		}
		n += m
		buf = buf[m:]
		incSequence(&c.writeSequence)
	}

	return
}

func (c *Conn) ReadProto(out proto.Message) error {
	buf := make([]byte, pond.TransportSize+4+secretbox.Overhead)
	n, err := c.read(buf)
	if err != nil {
		return err
	}
	if n != pond.TransportSize+4 {
		return errors.New("transport: message wrong length")
	}

	n = int(buf[0]) | int(buf[1])<<8 | int(buf[2])<<12 | int(buf[3])<<16
	buf = buf[4:]
	if n > len(buf) {
		return errors.New("transport: corrupt message")
	}
	return proto.Unmarshal(buf[:n], out)
}

func (c *Conn) WriteProto(msg proto.Message) error {
	data, err := proto.Marshal(msg)
	if err != nil {
		return err
	}
	if len(data) > pond.TransportSize {
		return errors.New("transport: message too large")
	}

	buf := make([]byte, pond.TransportSize+4)
	buf[0] = byte(len(data))
	buf[1] = byte(len(data) >> 8)
	buf[2] = byte(len(data) >> 12)
	buf[3] = byte(len(data) >> 16)
	copy(buf[4:], data)
	_, err = c.write(buf)
	return err
}

func (c *Conn) Close() (err error) {
	if !c.isServer {
		_, err = c.write(nil)
	}

	if closeErr := c.conn.Close(); err == nil {
		err = closeErr
	}

	return
}

func (c *Conn) WaitForClose() error {
	if !c.isServer {
		panic("non-server waited for connection close")
	}
	n, err := c.read(make([]byte, 128))
	if err != nil {
		return err
	}
	if n != 0 {
		return errors.New("transport: non-close message received when expecting close")
	}

	return nil
}

func (c *Conn) read(data []byte) (n int, err error) {
	var lengthBytes [4]byte

	if _, err := io.ReadFull(c.conn, lengthBytes[:]); err != nil {
		return 0, err
	}

	theirLength := int(lengthBytes[0]) + int(lengthBytes[1])<<8 + int(lengthBytes[2])<<12 + int(lengthBytes[3])<<16
	if theirLength > len(data) {
		return 0, errors.New("tranport: given buffer too small (" + strconv.Itoa(len(data)) + " vs " + strconv.Itoa(theirLength) + ")")
	}

	data = data[:theirLength]
	if _, err := io.ReadFull(c.conn, data); err != nil {
		return 0, err
	}

	decrypted, err := c.decrypt(data)
	if err != nil {
		return 0, err
	}
	copy(data, decrypted)

	return len(decrypted), nil
}

func (c *Conn) write(data []byte) (n int, err error) {
	encrypted := c.encrypt(data)

	var lengthBytes [4]byte
	lengthBytes[0] = byte(len(encrypted))
	lengthBytes[1] = byte(len(encrypted) >> 8)
	lengthBytes[2] = byte(len(encrypted) >> 12)
	lengthBytes[3] = byte(len(encrypted) >> 16)

	if _, err := c.conn.Write(lengthBytes[:]); err != nil {
		return 0, err
	}
	if _, err := c.conn.Write(encrypted); err != nil {
		return 0, err
	}

	return len(data), nil
}

func (c *Conn) encrypt(data []byte) []byte {
	if !c.writeKeyValid {
		return data
	}

	encrypted := secretbox.Seal(nil, data, &c.writeSequence, &c.writeKey)
	incSequence(&c.writeSequence)
	return encrypted
}

func (c *Conn) decrypt(data []byte) ([]byte, error) {
	if !c.readKeyValid {
		return data, nil
	}

	decrypted, ok := secretbox.Open(nil, data, &c.readSequence, &c.readKey)
	incSequence(&c.readSequence)
	if !ok {
		return nil, errors.New("transport: bad MAC")
	}
	return decrypted, nil
}

var serverKeysMagic = []byte("server keys snap")
var clientKeysMagic = []byte("client keys snap")

func (c *Conn) setupKeys(ephemeralShared *[32]byte) {
	var writeMagic, readMagic []byte
	if c.isServer {
		writeMagic, readMagic = serverKeysMagic, clientKeysMagic
	} else {
		writeMagic, readMagic = clientKeysMagic, serverKeysMagic
	}

	h := sha256.New()
	h.Write(writeMagic)
	h.Write(ephemeralShared[:])
	h.Sum(c.writeKey[:0])
	c.writeKeyValid = true

	h.Reset()
	h.Write(readMagic)
	h.Write(ephemeralShared[:])
	h.Sum(c.readKey[:0])
	c.readKeyValid = true
}

var serverProofMagic = []byte("server proof snap")
var clientProofMagic = []byte("client proof snap")

var shortMessageError = errors.New("transport: received short handshake message")

func (c *Conn) Handshake() error {
	var ephemeralPrivate, ephemeralPublic, ephemeralShared [32]byte
	if _, err := io.ReadFull(rand.Reader, ephemeralPrivate[:]); err != nil {
		return err
	}
	curve25519.ScalarBaseMult(&ephemeralPublic, &ephemeralPrivate)

	if _, err := c.write(ephemeralPublic[:]); err != nil {
		return err
	}

	var theirEphemeralPublic [32]byte
	if n, err := c.read(theirEphemeralPublic[:]); err != nil || n != len(theirEphemeralPublic) {
		if err == nil {
			err = shortMessageError
		}
		return err
	}

	handshakeHash := sha256.New()
	if c.isServer {
		handshakeHash.Write(theirEphemeralPublic[:])
		handshakeHash.Write(ephemeralPublic[:])
	} else {
		handshakeHash.Write(ephemeralPublic[:])
		handshakeHash.Write(theirEphemeralPublic[:])
	}

	curve25519.ScalarMult(&ephemeralShared, &ephemeralPrivate, &theirEphemeralPublic)
	c.setupKeys(&ephemeralShared)

	if c.isServer {
		return c.handshakeServer(handshakeHash, &theirEphemeralPublic)
	}
	return c.handshakeClient(handshakeHash, &ephemeralPrivate)
}

func (c *Conn) handshakeClient(handshakeHash hash.Hash, ephemeralPrivate *[32]byte) error {
	var ephemeralIdentityShared [32]byte
	curve25519.ScalarMult(&ephemeralIdentityShared, ephemeralPrivate, &c.Peer)

	digest := handshakeHash.Sum(nil)
	h := hmac.New(sha256.New, ephemeralIdentityShared[:])
	h.Write(serverProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	digestReceived := make([]byte, len(digest)+secretbox.Overhead)
	n, err := c.read(digestReceived)
	if err != nil {
		return err
	}
	if n != len(digest) {
		return shortMessageError
	}
	digestReceived = digestReceived[:n]

	if subtle.ConstantTimeCompare(digest, digestReceived) != 1 {
		return errors.New("transport: server identity incorrect")
	}

	var identityShared [32]byte
	curve25519.ScalarMult(&identityShared, &c.identity, &c.Peer)

	handshakeHash.Write(digest)
	digest = handshakeHash.Sum(digest[:0])

	h = hmac.New(sha256.New, identityShared[:])
	h.Write(clientProofMagic)
	h.Write(digest)

	finalMessage := make([]byte, 32+sha256.Size)
	copy(finalMessage, c.identityPublic[:])
	h.Sum(finalMessage[32:32])

	if _, err := c.write(finalMessage); err != nil {
		return err
	}

	return nil
}

func (c *Conn) handshakeServer(handshakeHash hash.Hash, theirEphemeralPublic *[32]byte) error {
	var ephemeralIdentityShared [32]byte
	curve25519.ScalarMult(&ephemeralIdentityShared, &c.identity, theirEphemeralPublic)

	digest := handshakeHash.Sum(nil)
	h := hmac.New(sha256.New, ephemeralIdentityShared[:])
	h.Write(serverProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	if _, err := c.write(digest); err != nil {
		return err
	}

	handshakeHash.Write(digest)
	digest = handshakeHash.Sum(digest[:0])

	finalMessage := make([]byte, 32+sha256.Size+secretbox.Overhead)
	n, err := c.read(finalMessage)
	if err != nil {
		return err
	}
	if n != 32+sha256.Size {
		return shortMessageError
	}
	finalMessage = finalMessage[:n]

	copy(c.Peer[:], finalMessage[:32])
	var identityShared [32]byte
	curve25519.ScalarMult(&identityShared, &c.identity, &c.Peer)

	h = hmac.New(sha256.New, identityShared[:])
	h.Write(clientProofMagic)
	h.Write(digest)
	digest = h.Sum(digest[:0])

	if subtle.ConstantTimeCompare(digest, finalMessage[32:]) != 1 {
		return errors.New("transport: bad proof from client")
	}

	return nil
}
<|endoftext|>


<commit_before>package hwaflib

func (ctx *Context) Version() string {
	version := "20131126"
	return version
}

func (ctx *Context) Revision() string {
	revision := "c824fb3"
	return revision
}

// EOF

<commit_msg>version: 20131128<commit_after>package hwaflib

func (ctx *Context) Version() string {
	version := "20131128"
	return version
}

func (ctx *Context) Revision() string {
	revision := "7664e8f"
	return revision
}

// EOF

<|endoftext|>


<commit_before>package generic

import (
	"encoding/json"
	"strconv"
	"testing"
)

func BenchmarkConvJSON(b *testing.B) {
	var n int64 = 1000
	for i := 0; i < b.N; i++ {
		json.Marshal(n)
	}
}

func BenchmarkConvStrconv(b *testing.B) {
	var n int64 = 1000
	for i := 0; i < b.N; i++ {
		_ = []byte(strconv.FormatInt(n, 10))
	}
}

func BenchmarkUnmarshalInterface(b *testing.B) {
	jb := []byte(`"\"Fran\""`)
	var in interface{}
	for i := 0; i < b.N; i++ {
		json.Unmarshal(jb, &in)
	}
	// fmt.Println(in)
}

func BenchmarkUnmarshalRawMessage(b *testing.B) {
	jb := []byte(`"\"Fran\""`)
	for i := 0; i < b.N; i++ {
		s := string(jb)
		strconv.Unquote(s)
		// fmt.Println(s)
	}
}
<commit_msg>remove test file<commit_after><|endoftext|>


<commit_before>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type ShutdownConfig

package common

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/packer/builder/vsphere/driver"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type ShutdownConfig struct {
	// Specify a VM guest shutdown command. This command will be executed using
	// the `communicator`. Otherwise the VMware guest tools are used to gracefully
	// shutdown the VM guest.
	Command string `mapstructure:"shutdown_command"`
	// Amount of time to wait for graceful VM shutdown.
	// Defaults to 5m or five minutes.
	// This will likely need to be modified if the `communicator` is 'none'.
	Timeout time.Duration `mapstructure:"shutdown_timeout"`
	// Packer normally halts the virtual machine after all provisioners have
	// run when no `shutdown_command` is defined. If this is set to `true`, Packer
	// *will not* halt the virtual machine but will assume that you will send the stop
	// signal yourself through a script or the final provisioner.
	// Packer will wait for a default of five minutes until the virtual machine is shutdown.
	// The timeout can be changed using `shutdown_timeout` option.
	DisableShutdown bool `mapstructure:"disable_shutdown"`
}

func (c *ShutdownConfig) Prepare(comm communicator.Config) (warnings []string, errs []error) {

	if c.Timeout == 0 {
		c.Timeout = 5 * time.Minute
	}

	if comm.Type == "none" && c.Command != "" {
		warnings = append(warnings, "The parameter `shutdown_command` is ignored as it requires a `communicator`.")
	}

	return
}

type StepShutdown struct {
	Config *ShutdownConfig
}

func (s *StepShutdown) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	vm := state.Get("vm").(*driver.VirtualMachineDriver)

	if off, _ := vm.IsPoweredOff(); off {
		// Probably power off initiated by last provisioner, though disable_shutdown is not set
		ui.Say("VM is already powered off")
		return multistep.ActionContinue
	}

	if state.Get("communicator") == nil {

		msg := fmt.Sprintf("Please shutdown virtual machine within %s.", s.Config.Timeout)
		ui.Message(msg)

	} else if s.Config.DisableShutdown {
		ui.Say("Automatic shutdown disabled. Please shutdown virtual machine.")
	} else if s.Config.Command != "" {
		// Communicator is not needed unless shutdown_command is populated
		comm := state.Get("communicator").(packer.Communicator)

		ui.Say("Executing shutdown command...")
		log.Printf("Shutdown command: %s", s.Config.Command)

		var stdout, stderr bytes.Buffer
		cmd := &packer.RemoteCmd{
			Command: s.Config.Command,
			Stdout:  &stdout,
			Stderr:  &stderr,
		}
		err := comm.Start(ctx, cmd)
		if err != nil {
			state.Put("error", fmt.Errorf("Failed to send shutdown command: %s", err))
			return multistep.ActionHalt
		}
	} else {
		ui.Say("Shutting down VM...")

		err := vm.StartShutdown()
		if err != nil {
			state.Put("error", fmt.Errorf("Cannot shut down VM: %v", err))
			return multistep.ActionHalt
		}
	}

	log.Printf("Waiting max %s for shutdown to complete", s.Config.Timeout)
	err := vm.WaitForShutdown(ctx, s.Config.Timeout)
	if err != nil {
		state.Put("error", err)
		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

func (s *StepShutdown) Cleanup(state multistep.StateBag) {}
<commit_msg>syntax change<commit_after>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type ShutdownConfig

package common

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/packer/builder/vsphere/driver"
	"github.com/hashicorp/packer/helper/communicator"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type ShutdownConfig struct {
	// Specify a VM guest shutdown command. This command will be executed using
	// the `communicator`. Otherwise the VMware guest tools are used to gracefully
	// shutdown the VM guest.
	Command string `mapstructure:"shutdown_command"`
	// Amount of time to wait for graceful VM shutdown.
	// Defaults to 5m or five minutes.
	// This will likely need to be modified if the `communicator` is 'none'.
	Timeout time.Duration `mapstructure:"shutdown_timeout"`
	// Packer normally halts the virtual machine after all provisioners have
	// run when no `shutdown_command` is defined. If this is set to `true`, Packer
	// *will not* halt the virtual machine but will assume that you will send the stop
	// signal yourself through a script or the final provisioner.
	// Packer will wait for a default of five minutes until the virtual machine is shutdown.
	// The timeout can be changed using `shutdown_timeout` option.
	DisableShutdown bool `mapstructure:"disable_shutdown"`
}

func (c *ShutdownConfig) Prepare(comm communicator.Config) (warnings []string, errs []error) {

	if c.Timeout == 0 {
		c.Timeout = 5 * time.Minute
	}

	if comm.Type == "none" && c.Command != "" {
		warnings = append(warnings, "The parameter `shutdown_command` is ignored as it requires a `communicator`.")
	}

	return
}

type StepShutdown struct {
	Config *ShutdownConfig
}

func (s *StepShutdown) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	vm := state.Get("vm").(*driver.VirtualMachineDriver)

	if off, _ := vm.IsPoweredOff(); off {
		// Probably power off initiated by last provisioner, though disable_shutdown is not set
		ui.Say("VM is already powered off")
		return multistep.ActionContinue
	}

	comm, _ := state.Get("communicator").(packer.Communicator)
	if comm == nil {

		msg := fmt.Sprintf("Please shutdown virtual machine within %s.", s.Config.Timeout)
		ui.Message(msg)

	} else if s.Config.DisableShutdown {
		ui.Say("Automatic shutdown disabled. Please shutdown virtual machine.")
	} else if s.Config.Command != "" {
		// Communicator is not needed unless shutdown_command is populated
		comm := state.Get("communicator").(packer.Communicator)

		ui.Say("Executing shutdown command...")
		log.Printf("Shutdown command: %s", s.Config.Command)

		var stdout, stderr bytes.Buffer
		cmd := &packer.RemoteCmd{
			Command: s.Config.Command,
			Stdout:  &stdout,
			Stderr:  &stderr,
		}
		err := comm.Start(ctx, cmd)
		if err != nil {
			state.Put("error", fmt.Errorf("Failed to send shutdown command: %s", err))
			return multistep.ActionHalt
		}
	} else {
		ui.Say("Shutting down VM...")

		err := vm.StartShutdown()
		if err != nil {
			state.Put("error", fmt.Errorf("Cannot shut down VM: %v", err))
			return multistep.ActionHalt
		}
	}

	log.Printf("Waiting max %s for shutdown to complete", s.Config.Timeout)
	err := vm.WaitForShutdown(ctx, s.Config.Timeout)
	if err != nil {
		state.Put("error", err)
		return multistep.ActionHalt
	}

	return multistep.ActionContinue
}

func (s *StepShutdown) Cleanup(state multistep.StateBag) {}
<|endoftext|>


<commit_before>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

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

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// A directory of manifest files that will be uploaded to the remote
	// machine.
	ManifestDir string `mapstructure:"manifest_dir"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	FacterVars      string
	HieraConfigPath string
	ModulePath      string
	ManifestFile    string
	ManifestDir     string
	Sudo            bool
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate: true,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if ne .HieraConfigPath \"\"}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"{{if ne .ManifestDir \"\"}}--manifestdir='{{.ManifestDir}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	// Validation
	var errs *packer.MultiError
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.HieraConfigPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestDir != "" {
		info, err := os.Stat(p.config.ManifestDir)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir must point to a directory"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		_, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory", i))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload manifest dir if set
	remoteManifestDir := ""
	if p.config.ManifestDir != "" {
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestDir))
		remoteManifestDir = fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestDir)
		if err != nil {
			return fmt.Errorf("Error uploading manifest dir: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	p.config.ctx.Data = &ExecuteTemplate{
		FacterVars:      strings.Join(facterVars, " "),
		HieraConfigPath: remoteHieraConfigPath,
		ManifestDir:     remoteManifestDir,
		ManifestFile:    remoteManifestFile,
		ModulePath:      strings.Join(modulePaths, ":"),
		Sudo:            !p.config.PreventSudo,
	}
	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f, nil); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// Upload the main manifest
	f, err := os.Open(p.config.ManifestFile)
	if err != nil {
		return "", err
	}
	defer f.Close()

	manifestFilename := p.config.ManifestFile
	if fi, err := os.Stat(p.config.ManifestFile); err != nil {
		return "", fmt.Errorf("Error inspecting manifest file: %s", err)
	} else if !fi.IsDir() {
		manifestFilename = filepath.Base(manifestFilename)
	}

	remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
	if err := comm.Upload(remoteManifestFile, f, nil); err != nil {
		return "", err
	}

	return remoteManifestFile, nil
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<commit_msg>provisioner/puppet-masterless: deprecation warning<commit_after>// This package implements a provisioner for Packer that executes
// Puppet on the remote machine, configured to apply a local manifest
// versus connecting to a Puppet master.
package puppetmasterless

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

	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`
	ctx                 interpolate.Context

	// The command used to execute Puppet.
	ExecuteCommand string `mapstructure:"execute_command"`

	// Additional facts to set when executing Puppet
	Facter map[string]string

	// Path to a hiera configuration file to upload and use.
	HieraConfigPath string `mapstructure:"hiera_config_path"`

	// An array of local paths of modules to upload.
	ModulePaths []string `mapstructure:"module_paths"`

	// The main manifest file to apply to kick off the entire thing.
	ManifestFile string `mapstructure:"manifest_file"`

	// A directory of manifest files that will be uploaded to the remote
	// machine.
	ManifestDir string `mapstructure:"manifest_dir"`

	// If true, `sudo` will NOT be used to execute Puppet.
	PreventSudo bool `mapstructure:"prevent_sudo"`

	// The directory where files will be uploaded. Packer requires write
	// permissions in this directory.
	StagingDir string `mapstructure:"staging_directory"`
}

type Provisioner struct {
	config Config
}

type ExecuteTemplate struct {
	FacterVars      string
	HieraConfigPath string
	ModulePath      string
	ManifestFile    string
	ManifestDir     string
	Sudo            bool
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate: true,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	// Set some defaults
	if p.config.ExecuteCommand == "" {
		p.config.ExecuteCommand = "{{.FacterVars}} {{if .Sudo}} sudo -E {{end}}" +
			"puppet apply --verbose --modulepath='{{.ModulePath}}' " +
			"{{if ne .HieraConfigPath \"\"}}--hiera_config='{{.HieraConfigPath}}' {{end}}" +
			"{{if ne .ManifestDir \"\"}}--manifestdir='{{.ManifestDir}}' {{end}}" +
			"--detailed-exitcodes " +
			"{{.ManifestFile}}"
	}

	if p.config.StagingDir == "" {
		p.config.StagingDir = "/tmp/packer-puppet-masterless"
	}

	// Validation
	var errs *packer.MultiError
	if p.config.HieraConfigPath != "" {
		info, err := os.Stat(p.config.HieraConfigPath)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path is invalid: %s", err))
		} else if info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("hiera_config_path must point to a file"))
		}
	}

	if p.config.ManifestDir != "" {
		info, err := os.Stat(p.config.ManifestDir)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir is invalid: %s", err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_dir must point to a directory"))
		}
	}

	if p.config.ManifestFile == "" {
		errs = packer.MultiErrorAppend(errs,
			fmt.Errorf("A manifest_file must be specified."))
	} else {
		_, err := os.Stat(p.config.ManifestFile)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("manifest_file is invalid: %s", err))
		}
	}

	for i, path := range p.config.ModulePaths {
		info, err := os.Stat(path)
		if err != nil {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] is invalid: %s", i, err))
		} else if !info.IsDir() {
			errs = packer.MultiErrorAppend(errs,
				fmt.Errorf("module_path[%d] must point to a directory", i))
		}
	}

	if errs != nil && len(errs.Errors) > 0 {
		return errs
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	ui.Say("Provisioning with Puppet...")
	ui.Message("Creating Puppet staging directory...")
	if err := p.createDir(ui, comm, p.config.StagingDir); err != nil {
		return fmt.Errorf("Error creating staging directory: %s", err)
	}

	// Upload hiera config if set
	remoteHieraConfigPath := ""
	if p.config.HieraConfigPath != "" {
		var err error
		remoteHieraConfigPath, err = p.uploadHieraConfig(ui, comm)
		if err != nil {
			return fmt.Errorf("Error uploading hiera config: %s", err)
		}
	}

	// Upload manifest dir if set
	remoteManifestDir := ""
	if p.config.ManifestDir != "" {
		ui.Message(fmt.Sprintf(
			"Uploading manifest directory from: %s", p.config.ManifestDir))
		remoteManifestDir = fmt.Sprintf("%s/manifests", p.config.StagingDir)
		err := p.uploadDirectory(ui, comm, remoteManifestDir, p.config.ManifestDir)
		if err != nil {
			return fmt.Errorf("Error uploading manifest dir: %s", err)
		}
	}

	// Upload all modules
	modulePaths := make([]string, 0, len(p.config.ModulePaths))
	for i, path := range p.config.ModulePaths {
		ui.Message(fmt.Sprintf("Uploading local modules from: %s", path))
		targetPath := fmt.Sprintf("%s/module-%d", p.config.StagingDir, i)
		if err := p.uploadDirectory(ui, comm, targetPath, path); err != nil {
			return fmt.Errorf("Error uploading modules: %s", err)
		}

		modulePaths = append(modulePaths, targetPath)
	}

	// Upload manifests
	remoteManifestFile, err := p.uploadManifests(ui, comm)
	if err != nil {
		return fmt.Errorf("Error uploading manifests: %s", err)
	}

	// Compile the facter variables
	facterVars := make([]string, 0, len(p.config.Facter))
	for k, v := range p.config.Facter {
		facterVars = append(facterVars, fmt.Sprintf("FACTER_%s='%s'", k, v))
	}

	// Execute Puppet
	p.config.ctx.Data = &ExecuteTemplate{
		FacterVars:      strings.Join(facterVars, " "),
		HieraConfigPath: remoteHieraConfigPath,
		ManifestDir:     remoteManifestDir,
		ManifestFile:    remoteManifestFile,
		ModulePath:      strings.Join(modulePaths, ":"),
		Sudo:            !p.config.PreventSudo,
	}
	command, err := interpolate.Render(p.config.ExecuteCommand, &p.config.ctx)
	if err != nil {
		return err
	}

	cmd := &packer.RemoteCmd{
		Command: command,
	}

	ui.Message(fmt.Sprintf("Running Puppet: %s", command))
	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 && cmd.ExitStatus != 2 {
		return fmt.Errorf("Puppet exited with a non-zero exit status: %d", cmd.ExitStatus)
	}

	return nil
}

func (p *Provisioner) Cancel() {
	// Just hard quit. It isn't a big deal if what we're doing keeps
	// running on the other side.
	os.Exit(0)
}

func (p *Provisioner) uploadHieraConfig(ui packer.Ui, comm packer.Communicator) (string, error) {
	ui.Message("Uploading hiera configuration...")
	f, err := os.Open(p.config.HieraConfigPath)
	if err != nil {
		return "", err
	}
	defer f.Close()

	path := fmt.Sprintf("%s/hiera.yaml", p.config.StagingDir)
	if err := comm.Upload(path, f, nil); err != nil {
		return "", err
	}

	return path, nil
}

func (p *Provisioner) uploadManifests(ui packer.Ui, comm packer.Communicator) (string, error) {
	// Create the remote manifests directory...
	ui.Message("Uploading manifests...")
	remoteManifestsPath := fmt.Sprintf("%s/manifests", p.config.StagingDir)
	if err := p.createDir(ui, comm, remoteManifestsPath); err != nil {
		return "", fmt.Errorf("Error creating manifests directory: %s", err)
	}

	// Upload the main manifest
	f, err := os.Open(p.config.ManifestFile)
	if err != nil {
		return "", err
	}
	defer f.Close()

	manifestFilename := p.config.ManifestFile
	if fi, err := os.Stat(p.config.ManifestFile); err != nil {
		return "", fmt.Errorf("Error inspecting manifest file: %s", err)
	} else if !fi.IsDir() {
		manifestFilename = filepath.Base(manifestFilename)
	} else {
		ui.Say("WARNING: manifest_file should be a file. Use manifest_dir for directories")
	}

	remoteManifestFile := fmt.Sprintf("%s/%s", remoteManifestsPath, manifestFilename)
	if err := comm.Upload(remoteManifestFile, f, nil); err != nil {
		return "", err
	}

	return remoteManifestFile, nil
}

func (p *Provisioner) createDir(ui packer.Ui, comm packer.Communicator, dir string) error {
	cmd := &packer.RemoteCmd{
		Command: fmt.Sprintf("mkdir -p '%s'", dir),
	}

	if err := cmd.StartWithUi(comm, ui); err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Non-zero exit status.")
	}

	return nil
}

func (p *Provisioner) uploadDirectory(ui packer.Ui, comm packer.Communicator, dst string, src string) error {
	if err := p.createDir(ui, comm, dst); err != nil {
		return err
	}

	// Make sure there is a trailing "/" so that the directory isn't
	// created on the other side.
	if src[len(src)-1] != '/' {
		src = src + "/"
	}

	return comm.UploadDir(dst, src, nil)
}
<|endoftext|>


<commit_before>/*
                       WARNING WARNING WARNING

                Attention all potential contributors

   This testfile is not in the best state. We've been slowly transitioning
   from the built in "testing" package to using Ginkgo. As you can see, we've
   changed the format, but a lot of the setup, test body, descriptions, etc
   are either hardcoded, completely lacking, or misleading.

   For example:

   Describe("Testing with ginkgo"...)      // This is not a great description
   It("TestDoesSoemthing"...)              // This is a horrible description

   Describe("create-user command"...       // Describe the actual object under test
   It("creates a user when provided ..."   // this is more descriptive

   For good examples of writing Ginkgo tests for the cli, refer to

   src/github.com/cloudfoundry/cli/cf/commands/application/delete_app_test.go
   src/github.com/cloudfoundry/cli/cf/terminal/ui_test.go
   src/github.com/cloudfoundry/loggregator_consumer/consumer_test.go
*/

package application_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/application"
	"github.com/cloudfoundry/cli/cf/models"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

func callRestart(args []string, requirementsFactory *testreq.FakeReqFactory, starter ApplicationStarter, stopper ApplicationStopper) (ui *testterm.FakeUI) {
	ui = new(testterm.FakeUI)

	cmd := NewRestart(ui, starter, stopper)
	testcmd.RunCommand(cmd, args, requirementsFactory)
	return
}

var _ = Describe("Testing with ginkgo", func() {
	It("TestRestartCommandFailsWithUsage", func() {
		requirementsFactory := &testreq.FakeReqFactory{}
		starter := &testcmd.FakeAppStarter{}
		stopper := &testcmd.FakeAppStopper{}
		ui := callRestart([]string{}, requirementsFactory, starter, stopper)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callRestart([]string{"my-app"}, requirementsFactory, starter, stopper)
		Expect(ui.FailedWithUsage).To(BeFalse())
	})
	It("TestRestartRequirements", func() {

		app := models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"
		starter := &testcmd.FakeAppStarter{}
		stopper := &testcmd.FakeAppStopper{}

		requirementsFactory := &testreq.FakeReqFactory{Application: app, LoginSuccess: true, TargetedSpaceSuccess: true}
		callRestart([]string{"my-app"}, requirementsFactory, starter, stopper)
		Expect(testcmd.CommandDidPassRequirements).To(BeTrue())

		requirementsFactory = &testreq.FakeReqFactory{Application: app, LoginSuccess: false, TargetedSpaceSuccess: true}
		callRestart([]string{"my-app"}, requirementsFactory, starter, stopper)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())

		requirementsFactory = &testreq.FakeReqFactory{Application: app, LoginSuccess: true, TargetedSpaceSuccess: false}
		callRestart([]string{"my-app"}, requirementsFactory, starter, stopper)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
	})
	It("TestRestartApplication", func() {

		app := models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"
		requirementsFactory := &testreq.FakeReqFactory{Application: app, LoginSuccess: true, TargetedSpaceSuccess: true}
		starter := &testcmd.FakeAppStarter{}
		stopper := &testcmd.FakeAppStopper{}
		callRestart([]string{"my-app"}, requirementsFactory, starter, stopper)

		Expect(stopper.AppToStop).To(Equal(app))
		Expect(starter.AppToStart).To(Equal(app))
	})
})
<commit_msg>Cleanup 'cf restart' tests<commit_after>package application_test

import (
	"github.com/cloudfoundry/cli/cf/models"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"

	. "github.com/cloudfoundry/cli/cf/commands/application"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("restart command", func() {
	var (
		ui                  *testterm.FakeUI
		requirementsFactory *testreq.FakeReqFactory
		starter             *testcmd.FakeAppStarter
		stopper             *testcmd.FakeAppStopper
		app                 models.Application
	)

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		requirementsFactory = &testreq.FakeReqFactory{}
		starter = &testcmd.FakeAppStarter{}
		stopper = &testcmd.FakeAppStopper{}

		app = models.Application{}
		app.Name = "my-app"
		app.Guid = "my-app-guid"
	})

	runCommand := func(args ...string) {
		testcmd.RunCommand(NewRestart(ui, starter, stopper), args, requirementsFactory)
	}

	Describe("requirements", func() {
		It("fails with usage when an app name is not given", func() {
			requirementsFactory.LoginSuccess = true
			runCommand()
			Expect(ui.FailedWithUsage).To(BeTrue())
		})

		It("fails when not logged in", func() {
			requirementsFactory.Application = app
			requirementsFactory.TargetedSpaceSuccess = true
			runCommand()
			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})

		It("fails when a space is not targeted", func() {
			requirementsFactory.Application = app
			requirementsFactory.LoginSuccess = true
			runCommand()
			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})
	})

	Context("when logged in, targeting a space, and an app name is provided", func() {
		BeforeEach(func() {
			requirementsFactory.Application = app
			requirementsFactory.LoginSuccess = true
			requirementsFactory.TargetedSpaceSuccess = true
		})

		It("restarts the given app", func() {
			runCommand("my-app")

			Expect(stopper.AppToStop).To(Equal(app))
			Expect(starter.AppToStart).To(Equal(app))
			Expect(requirementsFactory.ApplicationName).To(Equal("my-app"))
		})
	})
})
<|endoftext|>


<commit_before>/*
                       WARNING WARNING WARNING

                Attention all potential contributors

   This testfile is not in the best state. We've been slowly transitioning
   from the built in "testing" package to using Ginkgo. As you can see, we've
   changed the format, but a lot of the setup, test body, descriptions, etc
   are either hardcoded, completely lacking, or misleading.

   For example:

   Describe("Testing with ginkgo"...)      // This is not a great description
   It("TestDoesSoemthing"...)              // This is a horrible description

   Describe("create-user command"...       // Describe the actual object under test
   It("creates a user when provided ..."   // this is more descriptive

   For good examples of writing Ginkgo tests for the cli, refer to

   src/github.com/cloudfoundry/cli/cf/commands/application/delete_app_test.go
   src/github.com/cloudfoundry/cli/cf/terminal/ui_test.go
   src/github.com/cloudfoundry/loggregator_consumer/consumer_test.go
*/

package user_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/user"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

var _ = Describe("Testing with ginkgo", func() {
	It("TestUnsetOrgRoleFailsWithUsage", func() {
		userRepo := &testapi.FakeUserRepository{}
		requirementsFactory := &testreq.FakeReqFactory{}

		ui := callUnsetOrgRole([]string{}, userRepo, requirementsFactory)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callUnsetOrgRole([]string{"username"}, userRepo, requirementsFactory)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callUnsetOrgRole([]string{"username", "org"}, userRepo, requirementsFactory)
		Expect(ui.FailedWithUsage).To(BeTrue())

		ui = callUnsetOrgRole([]string{"username", "org", "role"}, userRepo, requirementsFactory)
		Expect(ui.FailedWithUsage).To(BeFalse())
	})
	It("TestUnsetOrgRoleRequirements", func() {

		userRepo := &testapi.FakeUserRepository{}
		requirementsFactory := &testreq.FakeReqFactory{}
		args := []string{"username", "org", "role"}

		requirementsFactory.LoginSuccess = false
		callUnsetOrgRole(args, userRepo, requirementsFactory)
		Expect(testcmd.CommandDidPassRequirements).To(BeFalse())

		requirementsFactory.LoginSuccess = true
		callUnsetOrgRole(args, userRepo, requirementsFactory)
		Expect(testcmd.CommandDidPassRequirements).To(BeTrue())

		Expect(requirementsFactory.UserUsername).To(Equal("username"))
		Expect(requirementsFactory.OrganizationName).To(Equal("org"))
	})
	It("TestUnsetOrgRole", func() {

		userRepo := &testapi.FakeUserRepository{}
		user := models.UserFields{}
		user.Username = "some-user"
		user.Guid = "some-user-guid"
		org := models.Organization{}
		org.Name = "some-org"
		org.Guid = "some-org-guid"
		requirementsFactory := &testreq.FakeReqFactory{
			LoginSuccess: true,
			UserFields:   user,
			Organization: org,
		}
		args := []string{"my-username", "my-org", "OrgManager"}

		ui := callUnsetOrgRole(args, userRepo, requirementsFactory)

		Expect(ui.Outputs).To(ContainSubstrings(
			[]string{"Removing role", "OrgManager", "my-username", "my-org", "my-user"},
			[]string{"OK"},
		))

		Expect(userRepo.UnsetOrgRoleRole).To(Equal(models.ORG_MANAGER))
		Expect(userRepo.UnsetOrgRoleUserGuid).To(Equal("some-user-guid"))
		Expect(userRepo.UnsetOrgRoleOrganizationGuid).To(Equal("some-org-guid"))
	})
})

func callUnsetOrgRole(args []string, userRepo *testapi.FakeUserRepository, requirementsFactory *testreq.FakeReqFactory) (ui *testterm.FakeUI) {
	ui = &testterm.FakeUI{}
	configRepo := testconfig.NewRepositoryWithDefaults()
	cmd := NewUnsetOrgRole(ui, configRepo, userRepo)
	testcmd.RunCommand(cmd, args, requirementsFactory)
	return
}
<commit_msg>Cleanup unset_org_role test<commit_after>package user_test

import (
	. "github.com/cloudfoundry/cli/cf/commands/user"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/models"
	testapi "github.com/cloudfoundry/cli/testhelpers/api"
	testcmd "github.com/cloudfoundry/cli/testhelpers/commands"
	testconfig "github.com/cloudfoundry/cli/testhelpers/configuration"
	testreq "github.com/cloudfoundry/cli/testhelpers/requirements"
	testterm "github.com/cloudfoundry/cli/testhelpers/terminal"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	. "github.com/cloudfoundry/cli/testhelpers/matchers"
)

var _ = Describe("unset-org-role command", func() {
	var (
		ui                  *testterm.FakeUI
		userRepo            *testapi.FakeUserRepository
		configRepo          configuration.ReadWriter
		requirementsFactory *testreq.FakeReqFactory
	)

	runCommand := func(args ...string) {
		cmd := NewUnsetOrgRole(ui, configRepo, userRepo)
		testcmd.RunCommand(cmd, args, requirementsFactory)
	}

	BeforeEach(func() {
		ui = &testterm.FakeUI{}
		userRepo = &testapi.FakeUserRepository{}
		requirementsFactory = &testreq.FakeReqFactory{}
		configRepo = testconfig.NewRepositoryWithDefaults()
	})

	It("fails with usage when invoked without exactly three args", func() {
		runCommand("username", "org")
		Expect(ui.FailedWithUsage).To(BeTrue())

		runCommand("woah", "too", "many", "args")
		Expect(ui.FailedWithUsage).To(BeTrue())
	})

	Describe("requirements", func() {
		It("fails when not logged in", func() {
			requirementsFactory.LoginSuccess = false
			runCommand("username", "org", "role")
			Expect(testcmd.CommandDidPassRequirements).To(BeFalse())
		})

		It("succeeds when logged in", func() {
			requirementsFactory.LoginSuccess = true
			runCommand("username", "org", "role")
			Expect(testcmd.CommandDidPassRequirements).To(BeTrue())

			Expect(requirementsFactory.UserUsername).To(Equal("username"))
			Expect(requirementsFactory.OrganizationName).To(Equal("org"))
		})
	})

	Context("when logged in", func() {
		BeforeEach(func() {
			requirementsFactory.LoginSuccess = true

			user := models.UserFields{}
			user.Username = "some-user"
			user.Guid = "some-user-guid"
			org := models.Organization{}
			org.Name = "some-org"
			org.Guid = "some-org-guid"

			requirementsFactory.UserFields = user
			requirementsFactory.Organization = org
		})

		It("unsets a user's org role", func() {
			runCommand("my-username", "my-org", "OrgManager")

			Expect(ui.Outputs).To(ContainSubstrings(
				[]string{"Removing role", "OrgManager", "my-username", "my-org", "my-user"},
				[]string{"OK"},
			))

			Expect(userRepo.UnsetOrgRoleRole).To(Equal(models.ORG_MANAGER))
			Expect(userRepo.UnsetOrgRoleUserGuid).To(Equal("some-user-guid"))
			Expect(userRepo.UnsetOrgRoleOrganizationGuid).To(Equal("some-org-guid"))
		})
	})
})
<|endoftext|>