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

package pod

import (
	"reflect"
	"sync"

	"time"

	calicocache "github.com/projectcalico/k8s-policy/pkg/cache"
	"github.com/projectcalico/k8s-policy/pkg/controllers/controller"
	"github.com/projectcalico/k8s-policy/pkg/converter"
	"github.com/projectcalico/libcalico-go/lib/api"
	"github.com/projectcalico/libcalico-go/lib/client"
	"github.com/projectcalico/libcalico-go/lib/errors"
	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/fields"
	uruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

type EndpointCache struct {
	sync.RWMutex
	m map[string]api.WorkloadEndpoint
}

// PodController implements the Controller interface.  It is responsible for monitoring
// kubernetes pods and updating workload endpoints in the Calico datastore if changed.
type PodController struct {
	indexer       cache.Indexer
	informer      cache.Controller
	wepDataCache  calicocache.ResourceCache
	endpointCache EndpointCache
	calicoClient  *client.Client
	k8sClientset  *kubernetes.Clientset
}

// NewPodController Constructor for PodController
func NewPodController(k8sClientset *kubernetes.Clientset, calicoClient *client.Client) controller.Controller {
	podConverter := converter.NewPodConverter()

	// Function returns map of key->WorkloadEndpointData from the Calico datastore.
	listFunc := func() (map[string]interface{}, error) {
		// Get all workloadEndpoints for kubernetes orchestrator from the Calico datastore
		workloadEndpoints, err := calicoClient.WorkloadEndpoints().List(api.WorkloadEndpointMetadata{Orchestrator: "k8s"})
		if err != nil {
			return nil, err
		}

		// Iterate through and collect data from workload endpoints that we care about.
		m := make(map[string]interface{})
		for _, endpoint := range workloadEndpoints.Items {
			d := converter.BuildWorkloadEndpointData(endpoint)
			key := podConverter.GetKey(d)
			m[key] = d
		}
		log.Debugf("Found %d workload endpoints in Calico datastore:", len(m))
		return m, nil
	}

	cacheArgs := calicocache.ResourceCacheArgs{
		ListFunc:   listFunc,
		ObjectType: reflect.TypeOf(converter.WorkloadEndpointData{}),
	}

	wepDataCache := calicocache.NewResourceCache(cacheArgs)
	endpointCache := EndpointCache{m: make(map[string]api.WorkloadEndpoint)}

	// Create a Pod watcher.
	listWatcher := cache.NewListWatchFromClient(k8sClientset.Core().RESTClient(), "pods", "", fields.Everything())

	// Bind the Calico cache to kubernetes cache with the help of an informer. This way we make sure that
	// whenever the kubernetes cache is updated, changes get reflected in the Calico cache as well.
	indexer, informer := cache.NewIndexerInformer(listWatcher, &v1.Pod{}, 0, cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			key, err := cache.MetaNamespaceKeyFunc(obj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got ADD event for pod: %s", key)

			// Ignore updates for host networked pods.
			if isHostNetworked(obj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(obj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", obj)
				return
			}

			// Prime the cache - we only need to make changes to the datastore when the controller
			// receives an update, because initial state is written to the datastore by the CNI plugin.
			wepData := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wepData)
			wepDataCache.Prime(k, wepData)
		},
		UpdateFunc: func(oldObj interface{}, newObj interface{}) {
			key, err := cache.MetaNamespaceKeyFunc(newObj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got UPDATE event for pod: %s", key)
			log.Debugf("Old object: \n%#v\n", oldObj)
			log.Debugf("New object: \n%#v\n", newObj)

			// Ignore updates for not ready / irrelevant pods.
			if !isReadyCalicoPod(newObj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(newObj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", newObj)
				return
			}

			// Update the cache.
			wepData := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wepData)
			wepDataCache.Set(k, wepData)
		},
		DeleteFunc: func(obj interface{}) {
			key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(obj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got DELETE event for pod: %s", key)

			// Ignore updates for host networked pods.
			if isHostNetworked(obj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(obj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", obj)
				return
			}

			// Clean up after the deleted workload endpoint.
			wep := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wep)
			wepDataCache.Clean(k)
			endpointCache.Lock()
			delete(endpointCache.m, k)
			endpointCache.Unlock()
		},
	}, cache.Indexers{})

	return &PodController{indexer, informer, wepDataCache, endpointCache, calicoClient, k8sClientset}
}

// Run starts controller.Internally it starts syncing
// kubernetes and calico caches.
func (c *PodController) Run(threadiness int, reconcilerPeriod string, stopCh chan struct{}) {
	defer uruntime.HandleCrash()

	// Let the workers stop when we are done
	workqueue := c.wepDataCache.GetQueue()
	defer workqueue.ShutDown()

	log.Info("Starting Pod/WorkloadEndpoint controller")

	// Load endpoint cache. Retry when failed.
	log.Debug("Loading endpoint cache at start of day")
	for err := c.populateEndpointCache(); err != nil; {
		log.WithError(err).Errorf("Failed to load workload endpoint cache, retrying in 5s")
		time.Sleep(5 * time.Second)
	}

	// Wait till k8s cache is synced
	go c.informer.Run(stopCh)
	log.Debug("Waiting to sync with Kubernetes API (Pods)")
	for !c.informer.HasSynced() {
	}
	log.Debug("Finished syncing with Kubernetes API (Pods)")

	// Start Calico cache.
	c.wepDataCache.Run(reconcilerPeriod)

	// Start a number of worker threads to read from the queue.
	for i := 0; i < threadiness; i++ {
		go c.runWorker()
	}
	log.Info("Pod/WorkloadEndpoint controller is now running")

	<-stopCh
	log.Info("Stopping Pod controller")
}

func (c *PodController) runWorker() {
	for c.processNextItem() {
	}
}

func (c *PodController) processNextItem() bool {
	// Wait until there is a new item in the work queue.
	workqueue := c.wepDataCache.GetQueue()
	key, quit := workqueue.Get()
	if quit {
		return false
	}

	// Sync the object to the Calico datastore.
	if err := c.syncToCalico(key.(string)); err != nil {
		c.handleErr(err, key.(string))
	}

	// Indicate that we're done processing this key, allowing for safe parallel processing such that
	// two objects with the same key are never processed in parallel.
	workqueue.Done(key)
	return true
}

// syncToCalico syncs the given update to the Calico datastore.
func (c *PodController) syncToCalico(key string) error {
	// Check if the wep data exists in our cache.  If it doesn't, then we don't need to do anything,
	// since CNI handles deletion of workload endpoints.
	if wepData, exists := c.wepDataCache.Get(key); exists {
		// Get workloadEndpoint from cache
		clog := log.WithField("wep", key)
		endpoint, exists := c.endpointCache.m[key]
		if !exists {
			// Load endpoint cache.
			clog.Warnf("No corresponding WorkloadEndpoint in cache, re-loading cache from datastore")
			if err := c.populateEndpointCache(); err != nil {
				clog.WithError(err).Error("Failed to load workload endpoint cache")
				return err
			}

			// See if it is in the cache now.
			c.endpointCache.RLock()
			endpoint, exists = c.endpointCache.m[key]
			c.endpointCache.RUnlock()
			if !exists {
				// No endpoint in datastore - this means the pod hasn't been
				// created by the CNI plugin yet. Just wait until it has been.
				// This can only be hit when pod changes before
				// the pod has been deployed, so should be pretty uncommon.
				clog.Infof("Pod hasn't been created by the CNI plugin yet.")
				return nil
			}
		}

		// Compare to see if the workload endpoint data has changed.
		old := converter.BuildWorkloadEndpointData(endpoint)
		new := wepData.(converter.WorkloadEndpointData)
		if !reflect.DeepEqual(old, new) {
			// The relevant wep data has changed - update the wep and write it to the datastore.
			log.Infof("Writing endpoint %s with updated data %#v to Calico datastore", key, new)
			converter.MergeWorkloadEndpointData(&endpoint, new)
			_, err := c.calicoClient.WorkloadEndpoints().Update(&endpoint)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceUpdateConflict); !ok {
					// Not an update conflict - return the error right away.
					clog.WithError(err).Errorf("failed to update workload endpoint")
					return err
				}

				// We hit an update conflict, re-query the WorkloadEndpoint before we try again.
				clog.Warn("Update conflict, re-querying workload endpoint")
				qwep, gErr := c.calicoClient.WorkloadEndpoints().Get(endpoint.Metadata)
				if gErr != nil {
					log.WithError(err).Errorf("failed to query workload endpoint %s", key)
					return gErr
				}
				clog.Warn("Updated cache with latest wep from datastore.")
				c.endpointCache.Lock()
				c.endpointCache.m[key] = *qwep
				c.endpointCache.Unlock()
				return err
			}

			// Update endpoint cache as well with the modified endpoint.
			updatedWep, err := c.calicoClient.WorkloadEndpoints().Get(endpoint.Metadata)
			if err != nil {
				log.WithError(err).Errorf("failed to query workload endpoint %s", key)
				return err
			}
			c.endpointCache.Lock()
			c.endpointCache.m[key] = *updatedWep
			c.endpointCache.Unlock()
			return nil
		}
	}
	return nil
}

// populateEndpointCache() Loads map of workload endpoint objects from the Calico datastore
func (c *PodController) populateEndpointCache() error {
	// List all workload endpoints for kubernetes orchestrator
	workloadEndpointList, err := c.calicoClient.WorkloadEndpoints().List(api.WorkloadEndpointMetadata{Orchestrator: "k8s"})
	if err != nil {
		return err
	}

	c.endpointCache.Lock()
	for _, wep := range workloadEndpointList.Items {
		c.endpointCache.m[wep.Metadata.Workload] = wep
	}
	c.endpointCache.Unlock()
	return nil
}

// handleErr checks if an error happened and makes sure we will retry later.
func (c *PodController) handleErr(err error, key string) {
	workqueue := c.wepDataCache.GetQueue()
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		log.WithField("key", key).Debug("Error for key is no more, drop from retry queue")
		workqueue.Forget(key)
		return
	}

	// This controller retries 5 times if something goes wrong. After that, it stops trying.
	if workqueue.NumRequeues(key) < 5 {
		log.WithError(err).Errorf("Error syncing pod, will retry: %v: %v", key, err)
		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		workqueue.AddRateLimited(key)
		return
	}
	workqueue.Forget(key)

	// Report to an external entity that, even after several retries, we could not successfully process this key
	uruntime.HandleError(err)
	log.WithError(err).Errorf("Dropping pod %q out of the retry queue: %v", key, err)
}

func isReadyCalicoPod(pod *v1.Pod) bool {
	if isHostNetworked(pod) {
		log.WithField("pod", pod.Name).Debug("Pod is host networked.")
		return false
	} else if !hasIPAddress(pod) {
		log.WithField("pod", pod.Name).Debug("Pod does not have an IP address.")
		return false
	} else if !isScheduled(pod) {
		log.WithField("pod", pod.Name).Debug("Pod is not scheduled.")
		return false
	}
	return true
}

func isScheduled(pod *v1.Pod) bool {
	return pod.Spec.NodeName != ""
}

func isHostNetworked(pod *v1.Pod) bool {
	return pod.Spec.HostNetwork
}

func hasIPAddress(pod *v1.Pod) bool {
	return pod.Status.PodIP != ""
}
<commit_msg>Added missing lock around endpointCache<commit_after>// Copyright (c) 2017 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pod

import (
	"reflect"
	"sync"

	"time"

	calicocache "github.com/projectcalico/k8s-policy/pkg/cache"
	"github.com/projectcalico/k8s-policy/pkg/controllers/controller"
	"github.com/projectcalico/k8s-policy/pkg/converter"
	"github.com/projectcalico/libcalico-go/lib/api"
	"github.com/projectcalico/libcalico-go/lib/client"
	"github.com/projectcalico/libcalico-go/lib/errors"
	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/fields"
	uruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

type EndpointCache struct {
	sync.RWMutex
	m map[string]api.WorkloadEndpoint
}

// PodController implements the Controller interface.  It is responsible for monitoring
// kubernetes pods and updating workload endpoints in the Calico datastore if changed.
type PodController struct {
	indexer       cache.Indexer
	informer      cache.Controller
	wepDataCache  calicocache.ResourceCache
	endpointCache EndpointCache
	calicoClient  *client.Client
	k8sClientset  *kubernetes.Clientset
}

// NewPodController Constructor for PodController
func NewPodController(k8sClientset *kubernetes.Clientset, calicoClient *client.Client) controller.Controller {
	podConverter := converter.NewPodConverter()

	// Function returns map of key->WorkloadEndpointData from the Calico datastore.
	listFunc := func() (map[string]interface{}, error) {
		// Get all workloadEndpoints for kubernetes orchestrator from the Calico datastore
		workloadEndpoints, err := calicoClient.WorkloadEndpoints().List(api.WorkloadEndpointMetadata{Orchestrator: "k8s"})
		if err != nil {
			return nil, err
		}

		// Iterate through and collect data from workload endpoints that we care about.
		m := make(map[string]interface{})
		for _, endpoint := range workloadEndpoints.Items {
			d := converter.BuildWorkloadEndpointData(endpoint)
			key := podConverter.GetKey(d)
			m[key] = d
		}
		log.Debugf("Found %d workload endpoints in Calico datastore:", len(m))
		return m, nil
	}

	cacheArgs := calicocache.ResourceCacheArgs{
		ListFunc:   listFunc,
		ObjectType: reflect.TypeOf(converter.WorkloadEndpointData{}),
	}

	wepDataCache := calicocache.NewResourceCache(cacheArgs)
	endpointCache := EndpointCache{m: make(map[string]api.WorkloadEndpoint)}

	// Create a Pod watcher.
	listWatcher := cache.NewListWatchFromClient(k8sClientset.Core().RESTClient(), "pods", "", fields.Everything())

	// Bind the Calico cache to kubernetes cache with the help of an informer. This way we make sure that
	// whenever the kubernetes cache is updated, changes get reflected in the Calico cache as well.
	indexer, informer := cache.NewIndexerInformer(listWatcher, &v1.Pod{}, 0, cache.ResourceEventHandlerFuncs{
		AddFunc: func(obj interface{}) {
			key, err := cache.MetaNamespaceKeyFunc(obj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got ADD event for pod: %s", key)

			// Ignore updates for host networked pods.
			if isHostNetworked(obj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(obj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", obj)
				return
			}

			// Prime the cache - we only need to make changes to the datastore when the controller
			// receives an update, because initial state is written to the datastore by the CNI plugin.
			wepData := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wepData)
			wepDataCache.Prime(k, wepData)
		},
		UpdateFunc: func(oldObj interface{}, newObj interface{}) {
			key, err := cache.MetaNamespaceKeyFunc(newObj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got UPDATE event for pod: %s", key)
			log.Debugf("Old object: \n%#v\n", oldObj)
			log.Debugf("New object: \n%#v\n", newObj)

			// Ignore updates for not ready / irrelevant pods.
			if !isReadyCalicoPod(newObj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(newObj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", newObj)
				return
			}

			// Update the cache.
			wepData := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wepData)
			wepDataCache.Set(k, wepData)
		},
		DeleteFunc: func(obj interface{}) {
			key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(obj)
			if err != nil {
				log.WithError(err).Error("Failed to generate key")
				return
			}
			log.Debugf("Got DELETE event for pod: %s", key)

			// Ignore updates for host networked pods.
			if isHostNetworked(obj.(*v1.Pod)) {
				log.Debugf("Skipping irrelevant pod %s", key)
				return
			}

			wepInterface, err := podConverter.Convert(obj)
			if err != nil {
				log.WithError(err).Errorf("Error while converting %#v to wep.", obj)
				return
			}

			// Clean up after the deleted workload endpoint.
			wep := wepInterface.(converter.WorkloadEndpointData)
			k := podConverter.GetKey(wep)
			wepDataCache.Clean(k)
			endpointCache.Lock()
			delete(endpointCache.m, k)
			endpointCache.Unlock()
		},
	}, cache.Indexers{})

	return &PodController{indexer, informer, wepDataCache, endpointCache, calicoClient, k8sClientset}
}

// Run starts controller.Internally it starts syncing
// kubernetes and calico caches.
func (c *PodController) Run(threadiness int, reconcilerPeriod string, stopCh chan struct{}) {
	defer uruntime.HandleCrash()

	// Let the workers stop when we are done
	workqueue := c.wepDataCache.GetQueue()
	defer workqueue.ShutDown()

	log.Info("Starting Pod/WorkloadEndpoint controller")

	// Load endpoint cache. Retry when failed.
	log.Debug("Loading endpoint cache at start of day")
	for err := c.populateEndpointCache(); err != nil; {
		log.WithError(err).Errorf("Failed to load workload endpoint cache, retrying in 5s")
		time.Sleep(5 * time.Second)
	}

	// Wait till k8s cache is synced
	go c.informer.Run(stopCh)
	log.Debug("Waiting to sync with Kubernetes API (Pods)")
	for !c.informer.HasSynced() {
	}
	log.Debug("Finished syncing with Kubernetes API (Pods)")

	// Start Calico cache.
	c.wepDataCache.Run(reconcilerPeriod)

	// Start a number of worker threads to read from the queue.
	for i := 0; i < threadiness; i++ {
		go c.runWorker()
	}
	log.Info("Pod/WorkloadEndpoint controller is now running")

	<-stopCh
	log.Info("Stopping Pod controller")
}

func (c *PodController) runWorker() {
	for c.processNextItem() {
	}
}

func (c *PodController) processNextItem() bool {
	// Wait until there is a new item in the work queue.
	workqueue := c.wepDataCache.GetQueue()
	key, quit := workqueue.Get()
	if quit {
		return false
	}

	// Sync the object to the Calico datastore.
	if err := c.syncToCalico(key.(string)); err != nil {
		c.handleErr(err, key.(string))
	}

	// Indicate that we're done processing this key, allowing for safe parallel processing such that
	// two objects with the same key are never processed in parallel.
	workqueue.Done(key)
	return true
}

// syncToCalico syncs the given update to the Calico datastore.
func (c *PodController) syncToCalico(key string) error {
	// Check if the wep data exists in our cache.  If it doesn't, then we don't need to do anything,
	// since CNI handles deletion of workload endpoints.
	if wepData, exists := c.wepDataCache.Get(key); exists {
		// Get workloadEndpoint from cache
		clog := log.WithField("wep", key)
		c.endpointCache.RLock()
		endpoint, exists := c.endpointCache.m[key]
		c.endpointCache.RUnlock()
		if !exists {
			// Load endpoint cache.
			clog.Warnf("No corresponding WorkloadEndpoint in cache, re-loading cache from datastore")
			if err := c.populateEndpointCache(); err != nil {
				clog.WithError(err).Error("Failed to load workload endpoint cache")
				return err
			}

			// See if it is in the cache now.
			c.endpointCache.RLock()
			endpoint, exists = c.endpointCache.m[key]
			c.endpointCache.RUnlock()
			if !exists {
				// No endpoint in datastore - this means the pod hasn't been
				// created by the CNI plugin yet. Just wait until it has been.
				// This can only be hit when pod changes before
				// the pod has been deployed, so should be pretty uncommon.
				clog.Infof("Pod hasn't been created by the CNI plugin yet.")
				return nil
			}
		}

		// Compare to see if the workload endpoint data has changed.
		old := converter.BuildWorkloadEndpointData(endpoint)
		new := wepData.(converter.WorkloadEndpointData)
		if !reflect.DeepEqual(old, new) {
			// The relevant wep data has changed - update the wep and write it to the datastore.
			log.Infof("Writing endpoint %s with updated data %#v to Calico datastore", key, new)
			converter.MergeWorkloadEndpointData(&endpoint, new)
			_, err := c.calicoClient.WorkloadEndpoints().Update(&endpoint)
			if err != nil {
				if _, ok := err.(errors.ErrorResourceUpdateConflict); !ok {
					// Not an update conflict - return the error right away.
					clog.WithError(err).Errorf("failed to update workload endpoint")
					return err
				}

				// We hit an update conflict, re-query the WorkloadEndpoint before we try again.
				clog.Warn("Update conflict, re-querying workload endpoint")
				qwep, gErr := c.calicoClient.WorkloadEndpoints().Get(endpoint.Metadata)
				if gErr != nil {
					log.WithError(err).Errorf("failed to query workload endpoint %s", key)
					return gErr
				}
				clog.Warn("Updated cache with latest wep from datastore.")
				c.endpointCache.Lock()
				c.endpointCache.m[key] = *qwep
				c.endpointCache.Unlock()
				return err
			}

			// Update endpoint cache as well with the modified endpoint.
			updatedWep, err := c.calicoClient.WorkloadEndpoints().Get(endpoint.Metadata)
			if err != nil {
				log.WithError(err).Errorf("failed to query workload endpoint %s", key)
				return err
			}
			c.endpointCache.Lock()
			c.endpointCache.m[key] = *updatedWep
			c.endpointCache.Unlock()
			return nil
		}
	}
	return nil
}

// populateEndpointCache() Loads map of workload endpoint objects from the Calico datastore
func (c *PodController) populateEndpointCache() error {
	// List all workload endpoints for kubernetes orchestrator
	workloadEndpointList, err := c.calicoClient.WorkloadEndpoints().List(api.WorkloadEndpointMetadata{Orchestrator: "k8s"})
	if err != nil {
		return err
	}

	c.endpointCache.Lock()
	for _, wep := range workloadEndpointList.Items {
		c.endpointCache.m[wep.Metadata.Workload] = wep
	}
	c.endpointCache.Unlock()
	return nil
}

// handleErr checks if an error happened and makes sure we will retry later.
func (c *PodController) handleErr(err error, key string) {
	workqueue := c.wepDataCache.GetQueue()
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		log.WithField("key", key).Debug("Error for key is no more, drop from retry queue")
		workqueue.Forget(key)
		return
	}

	// This controller retries 5 times if something goes wrong. After that, it stops trying.
	if workqueue.NumRequeues(key) < 5 {
		log.WithError(err).Errorf("Error syncing pod, will retry: %v: %v", key, err)
		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		workqueue.AddRateLimited(key)
		return
	}
	workqueue.Forget(key)

	// Report to an external entity that, even after several retries, we could not successfully process this key
	uruntime.HandleError(err)
	log.WithError(err).Errorf("Dropping pod %q out of the retry queue: %v", key, err)
}

func isReadyCalicoPod(pod *v1.Pod) bool {
	if isHostNetworked(pod) {
		log.WithField("pod", pod.Name).Debug("Pod is host networked.")
		return false
	} else if !hasIPAddress(pod) {
		log.WithField("pod", pod.Name).Debug("Pod does not have an IP address.")
		return false
	} else if !isScheduled(pod) {
		log.WithField("pod", pod.Name).Debug("Pod is not scheduled.")
		return false
	}
	return true
}

func isScheduled(pod *v1.Pod) bool {
	return pod.Spec.NodeName != ""
}

func isHostNetworked(pod *v1.Pod) bool {
	return pod.Spec.HostNetwork
}

func hasIPAddress(pod *v1.Pod) bool {
	return pod.Status.PodIP != ""
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Skaffold 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 gcb

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	cstorage "cloud.google.com/go/storage"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/gcp"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/version"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/oauth2/google"
	cloudbuild "google.golang.org/api/cloudbuild/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/iterator"
)

// Build builds a list of artifacts with Google Cloud Build.
func (b *Builder) Build(ctx context.Context, out io.Writer, tagger tag.Tagger, artifacts []*latest.Artifact) ([]build.Artifact, error) {
	return build.InParallel(ctx, out, tagger, artifacts, b.buildArtifact)
}

func (b *Builder) buildArtifact(ctx context.Context, out io.Writer, tagger tag.Tagger, artifact *latest.Artifact) (string, error) {
	client, err := google.DefaultClient(ctx, cloudbuild.CloudPlatformScope)
	if err != nil {
		return "", errors.Wrap(err, "getting google client")
	}

	cbclient, err := cloudbuild.New(client)
	if err != nil {
		return "", errors.Wrap(err, "getting builder")
	}
	cbclient.UserAgent = version.UserAgent()

	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return "", errors.Wrap(err, "getting cloud storage client")
	}
	defer c.Close()

	projectID := b.ProjectID
	if projectID == "" {
		guessedProjectID, err := gcp.ExtractProjectID(artifact.ImageName)
		if err != nil {
			return "", errors.Wrap(err, "extracting projectID from image name")
		}

		projectID = guessedProjectID
	}

	cbBucket := fmt.Sprintf("%s%s", projectID, constants.GCSBucketSuffix)
	buildObject := fmt.Sprintf("source/%s-%s.tar.gz", projectID, util.RandomID())

	if err := b.createBucketIfNotExists(ctx, projectID, cbBucket); err != nil {
		return "", errors.Wrap(err, "creating bucket if not exists")
	}
	if err := b.checkBucketProjectCorrect(ctx, projectID, cbBucket); err != nil {
		return "", errors.Wrap(err, "checking bucket is in correct project")
	}

	color.Default.Fprintf(out, "Pushing code to gs://%s/%s\n", cbBucket, buildObject)
	if err := docker.UploadContextToGCS(ctx, artifact.Workspace, artifact.DockerArtifact, cbBucket, buildObject); err != nil {
		return "", errors.Wrap(err, "uploading source tarball")
	}

	desc := b.buildDescription(artifact, cbBucket, buildObject)
	call := cbclient.Projects.Builds.Create(projectID, desc)
	op, err := call.Context(ctx).Do()
	if err != nil {
		return "", errors.Wrap(err, "could not create build")
	}

	remoteID, err := getBuildID(op)
	if err != nil {
		return "", errors.Wrapf(err, "getting build ID from op")
	}
	logsObject := fmt.Sprintf("log-%s.txt", remoteID)
	color.Default.Fprintf(out, "Logs are available at \nhttps://console.cloud.google.com/m/cloudstorage/b/%s/o/%s\n", cbBucket, logsObject)
	var imageID string
	offset := int64(0)
watch:
	for {
		logrus.Debugf("current offset %d", offset)
		cb, err := cbclient.Projects.Builds.Get(projectID, remoteID).Do()
		if err != nil {
			return "", errors.Wrap(err, "getting build status")
		}

		r, err := b.getLogs(ctx, offset, cbBucket, logsObject)
		if err != nil {
			return "", errors.Wrap(err, "getting logs")
		}
		if r != nil {
			written, err := io.Copy(out, r)
			if err != nil {
				return "", errors.Wrap(err, "copying logs to stdout")
			}
			offset += written
			r.Close()
		}
		switch cb.Status {
		case StatusQueued, StatusWorking, StatusUnknown:
		case StatusSuccess:
			imageID, err = getImageID(cb)
			if err != nil {
				return "", errors.Wrap(err, "getting image id from finished build")
			}
			break watch
		case StatusFailure, StatusInternalError, StatusTimeout, StatusCancelled:
			return "", fmt.Errorf("cloud build failed: %s", cb.Status)
		default:
			return "", fmt.Errorf("unknown status: %s", cb.Status)
		}

		time.Sleep(RetryDelay)
	}

	if err := c.Bucket(cbBucket).Object(buildObject).Delete(ctx); err != nil {
		return "", errors.Wrap(err, "cleaning up source tar after build")
	}
	logrus.Infof("Deleted object %s", buildObject)
	builtTag := fmt.Sprintf("%s@%s", artifact.ImageName, imageID)
	logrus.Infof("Image built at %s", builtTag)

	newTag, err := tagger.GenerateFullyQualifiedImageName(artifact.Workspace, tag.Options{
		ImageName: artifact.ImageName,
		Digest:    imageID,
	})
	if err != nil {
		return "", errors.Wrap(err, "generating tag")
	}

	if err := docker.AddTag(builtTag, newTag); err != nil {
		return "", errors.Wrap(err, "tagging image")
	}

	return newTag, nil
}

func getBuildID(op *cloudbuild.Operation) (string, error) {
	if op.Metadata == nil {
		return "", errors.New("missing Metadata in operation")
	}
	var buildMeta cloudbuild.BuildOperationMetadata
	if err := json.Unmarshal([]byte(op.Metadata), &buildMeta); err != nil {
		return "", err
	}
	if buildMeta.Build == nil {
		return "", errors.New("missing Build in operation metadata")
	}
	return buildMeta.Build.Id, nil
}

func getImageID(b *cloudbuild.Build) (string, error) {
	if b.Results == nil || len(b.Results.Images) == 0 {
		return "", errors.New("build failed")
	}
	return b.Results.Images[0].Digest, nil
}

func (b *Builder) getLogs(ctx context.Context, offset int64, bucket, objectName string) (io.ReadCloser, error) {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	r, err := c.Bucket(bucket).Object(objectName).NewRangeReader(ctx, offset, -1)
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			switch gerr.Code {
			// case http.
			case 404, 416, 429, 503:
				logrus.Debugf("Status Code: %d, %s", gerr.Code, gerr.Body)
				return nil, nil
			}
		}
		if err == cstorage.ErrObjectNotExist {
			logrus.Debugf("Logs for %s %s not uploaded yet...", bucket, objectName)
			return nil, nil
		}
		return nil, errors.Wrap(err, "unknown error")
	}
	return r, nil
}

func (b *Builder) checkBucketProjectCorrect(ctx context.Context, projectID, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	it := c.Buckets(ctx, projectID)
	// Set the prefix to the bucket we're looking for to only return that bucket and buckets with that prefix
	// that we'll filter further later on
	it.Prefix = bucket
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			return errors.Wrap(err, "bucket not found")
		}
		if err != nil {
			return errors.Wrap(err, "iterating over buckets")
		}
		// Since we can't filter on bucket name specifically, only prefix, we need to check equality here and not just prefix
		if attrs.Name == bucket {
			return nil
		}
	}
}

func (b *Builder) createBucketIfNotExists(ctx context.Context, projectID, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	_, err = c.Bucket(bucket).Attrs(ctx)

	if err == nil {
		// Bucket exists
		return nil
	}

	if err != cstorage.ErrBucketNotExist {
		return errors.Wrapf(err, "getting bucket %s", bucket)
	}

	err = c.Bucket(bucket).Create(ctx, projectID, &cstorage.BucketAttrs{
		Name: bucket,
	})
	if e, ok := err.(*googleapi.Error); ok {
		if e.Code == http.StatusConflict {
			// 409 errors are ok, there could have been a race condition or eventual consistency.
			logrus.Debugf("Not creating bucket, got a 409 error indicating it already exists.")
			return nil
		}
	}

	if err != nil {
		return err
	}
	logrus.Debugf("Created bucket %s in %s", bucket, projectID)
	return nil
}
<commit_msg>[GCB] Rename ImageID to Digest.<commit_after>/*
Copyright 2018 The Skaffold 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 gcb

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	cstorage "cloud.google.com/go/storage"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/color"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/docker"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/gcp"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/version"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"golang.org/x/oauth2/google"
	cloudbuild "google.golang.org/api/cloudbuild/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/iterator"
)

// Build builds a list of artifacts with Google Cloud Build.
func (b *Builder) Build(ctx context.Context, out io.Writer, tagger tag.Tagger, artifacts []*latest.Artifact) ([]build.Artifact, error) {
	return build.InParallel(ctx, out, tagger, artifacts, b.buildArtifact)
}

func (b *Builder) buildArtifact(ctx context.Context, out io.Writer, tagger tag.Tagger, artifact *latest.Artifact) (string, error) {
	client, err := google.DefaultClient(ctx, cloudbuild.CloudPlatformScope)
	if err != nil {
		return "", errors.Wrap(err, "getting google client")
	}

	cbclient, err := cloudbuild.New(client)
	if err != nil {
		return "", errors.Wrap(err, "getting builder")
	}
	cbclient.UserAgent = version.UserAgent()

	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return "", errors.Wrap(err, "getting cloud storage client")
	}
	defer c.Close()

	projectID := b.ProjectID
	if projectID == "" {
		guessedProjectID, err := gcp.ExtractProjectID(artifact.ImageName)
		if err != nil {
			return "", errors.Wrap(err, "extracting projectID from image name")
		}

		projectID = guessedProjectID
	}

	cbBucket := fmt.Sprintf("%s%s", projectID, constants.GCSBucketSuffix)
	buildObject := fmt.Sprintf("source/%s-%s.tar.gz", projectID, util.RandomID())

	if err := b.createBucketIfNotExists(ctx, projectID, cbBucket); err != nil {
		return "", errors.Wrap(err, "creating bucket if not exists")
	}
	if err := b.checkBucketProjectCorrect(ctx, projectID, cbBucket); err != nil {
		return "", errors.Wrap(err, "checking bucket is in correct project")
	}

	color.Default.Fprintf(out, "Pushing code to gs://%s/%s\n", cbBucket, buildObject)
	if err := docker.UploadContextToGCS(ctx, artifact.Workspace, artifact.DockerArtifact, cbBucket, buildObject); err != nil {
		return "", errors.Wrap(err, "uploading source tarball")
	}

	desc := b.buildDescription(artifact, cbBucket, buildObject)
	call := cbclient.Projects.Builds.Create(projectID, desc)
	op, err := call.Context(ctx).Do()
	if err != nil {
		return "", errors.Wrap(err, "could not create build")
	}

	remoteID, err := getBuildID(op)
	if err != nil {
		return "", errors.Wrapf(err, "getting build ID from op")
	}
	logsObject := fmt.Sprintf("log-%s.txt", remoteID)
	color.Default.Fprintf(out, "Logs are available at \nhttps://console.cloud.google.com/m/cloudstorage/b/%s/o/%s\n", cbBucket, logsObject)

	var digest string
	offset := int64(0)
watch:
	for {
		logrus.Debugf("current offset %d", offset)
		cb, err := cbclient.Projects.Builds.Get(projectID, remoteID).Do()
		if err != nil {
			return "", errors.Wrap(err, "getting build status")
		}

		r, err := b.getLogs(ctx, offset, cbBucket, logsObject)
		if err != nil {
			return "", errors.Wrap(err, "getting logs")
		}
		if r != nil {
			written, err := io.Copy(out, r)
			if err != nil {
				return "", errors.Wrap(err, "copying logs to stdout")
			}
			offset += written
			r.Close()
		}
		switch cb.Status {
		case StatusQueued, StatusWorking, StatusUnknown:
		case StatusSuccess:
			digest, err = getDigest(cb)
			if err != nil {
				return "", errors.Wrap(err, "getting image id from finished build")
			}
			break watch
		case StatusFailure, StatusInternalError, StatusTimeout, StatusCancelled:
			return "", fmt.Errorf("cloud build failed: %s", cb.Status)
		default:
			return "", fmt.Errorf("unknown status: %s", cb.Status)
		}

		time.Sleep(RetryDelay)
	}

	if err := c.Bucket(cbBucket).Object(buildObject).Delete(ctx); err != nil {
		return "", errors.Wrap(err, "cleaning up source tar after build")
	}
	logrus.Infof("Deleted object %s", buildObject)
	builtTag := fmt.Sprintf("%s@%s", artifact.ImageName, digest)
	logrus.Infof("Image built at %s", builtTag)

	newTag, err := tagger.GenerateFullyQualifiedImageName(artifact.Workspace, tag.Options{
		ImageName: artifact.ImageName,
		Digest:    digest,
	})
	if err != nil {
		return "", errors.Wrap(err, "generating tag")
	}

	if err := docker.AddTag(builtTag, newTag); err != nil {
		return "", errors.Wrap(err, "tagging image")
	}

	return newTag, nil
}

func getBuildID(op *cloudbuild.Operation) (string, error) {
	if op.Metadata == nil {
		return "", errors.New("missing Metadata in operation")
	}
	var buildMeta cloudbuild.BuildOperationMetadata
	if err := json.Unmarshal([]byte(op.Metadata), &buildMeta); err != nil {
		return "", err
	}
	if buildMeta.Build == nil {
		return "", errors.New("missing Build in operation metadata")
	}
	return buildMeta.Build.Id, nil
}

func getDigest(b *cloudbuild.Build) (string, error) {
	if b.Results == nil || len(b.Results.Images) == 0 {
		return "", errors.New("build failed")
	}
	return b.Results.Images[0].Digest, nil
}

func (b *Builder) getLogs(ctx context.Context, offset int64, bucket, objectName string) (io.ReadCloser, error) {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	r, err := c.Bucket(bucket).Object(objectName).NewRangeReader(ctx, offset, -1)
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			switch gerr.Code {
			// case http.
			case 404, 416, 429, 503:
				logrus.Debugf("Status Code: %d, %s", gerr.Code, gerr.Body)
				return nil, nil
			}
		}
		if err == cstorage.ErrObjectNotExist {
			logrus.Debugf("Logs for %s %s not uploaded yet...", bucket, objectName)
			return nil, nil
		}
		return nil, errors.Wrap(err, "unknown error")
	}
	return r, nil
}

func (b *Builder) checkBucketProjectCorrect(ctx context.Context, projectID, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	it := c.Buckets(ctx, projectID)
	// Set the prefix to the bucket we're looking for to only return that bucket and buckets with that prefix
	// that we'll filter further later on
	it.Prefix = bucket
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			return errors.Wrap(err, "bucket not found")
		}
		if err != nil {
			return errors.Wrap(err, "iterating over buckets")
		}
		// Since we can't filter on bucket name specifically, only prefix, we need to check equality here and not just prefix
		if attrs.Name == bucket {
			return nil
		}
	}
}

func (b *Builder) createBucketIfNotExists(ctx context.Context, projectID, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	_, err = c.Bucket(bucket).Attrs(ctx)

	if err == nil {
		// Bucket exists
		return nil
	}

	if err != cstorage.ErrBucketNotExist {
		return errors.Wrapf(err, "getting bucket %s", bucket)
	}

	err = c.Bucket(bucket).Create(ctx, projectID, &cstorage.BucketAttrs{
		Name: bucket,
	})
	if e, ok := err.(*googleapi.Error); ok {
		if e.Code == http.StatusConflict {
			// 409 errors are ok, there could have been a race condition or eventual consistency.
			logrus.Debugf("Not creating bucket, got a 409 error indicating it already exists.")
			return nil
		}
	}

	if err != nil {
		return err
	}
	logrus.Debugf("Created bucket %s in %s", bucket, projectID)
	return nil
}
<|endoftext|>


<commit_before>package iptables

import (
	"fmt"
	"os/exec"

	"github.com/tmrts/flamingo/pkg/sys"
	"github.com/tmrts/flamingo/pkg/sys/firewall"
	"github.com/tmrts/flamingo/pkg/util"
)

type Target string

const (
	AcceptTarget Target = "ACCEPT"
	ReturnTarget Target = "RETURN"
	DropTarget   Target = "DROP"
)

type Operation string

const (
	Append Operation = "append"
	Delete Operation = "delete"
	Insert Operation = "insert"
	Check  Operation = "check"
)

type Implementation struct {
	sys.Executor
}

type Rule struct {
	Source      []string `flag:"source"`
	Destination []string `flag:"destination"`

	FromInterface string `flag:"in-interface"`
	ToInterface   string `flag:"out-interface"`

	IsSyncPackage bool `flag:"syn"`

	Protocol string `flag:"protocol"`

	Target Target `flag:"jump"`
}

func (r *Rule) FlagForm() []string {
	return util.GetFlagFormOfStruct(*r)
}

func (impl *Implementation) CheckDependencies() error {
	_, err := exec.LookPath("iptables")
	if err != nil {
		return fmt.Errorf("iptables: 'iptables' executable could not be found")
	}

	return nil
}

func (impl *Implementation) Perform(op Operation, c firewall.Chain, r *Rule) error {
	if r.Target == "" {
		return fmt.Errorf("iptables: rule doesnt have a target")
	}

	table := fmt.Sprintf("--table=%v", c.Table)
	action := fmt.Sprintf("--%v=%v", op, c.Name)

	args := []string{"--wait", table, action}

	args = append(args, r.FlagForm()...)

	_, err := impl.Execute("iptables", args...)

	return err
}
<commit_msg>Implement Stringer for iptables.Rule<commit_after>package iptables

import (
	"fmt"
	"os/exec"
	"strings"

	"github.com/tmrts/flamingo/pkg/sys"
	"github.com/tmrts/flamingo/pkg/sys/firewall"
	"github.com/tmrts/flamingo/pkg/util"
)

type Target string

const (
	AcceptTarget Target = "ACCEPT"
	ReturnTarget Target = "RETURN"
	DropTarget   Target = "DROP"
)

type Operation string

const (
	Append Operation = "append"
	Delete Operation = "delete"
	Insert Operation = "insert"
	Check  Operation = "check"
)

type Implementation struct {
	sys.Executor
}

type Rule struct {
	Source      []string `flag:"source"`
	Destination []string `flag:"destination"`

	FromInterface string `flag:"in-interface"`
	ToInterface   string `flag:"out-interface"`

	IsSyncPackage bool `flag:"syn"`

	Protocol string `flag:"protocol"`

	Target Target `flag:"jump"`
}

func (r *Rule) FlagForm() []string {
	return util.GetFlagFormOfStruct(*r)
}

func (r *Rule) String() string {
	return strings.Join(r.FlagForm(), " ")
}

func (impl *Implementation) CheckDependencies() error {
	_, err := exec.LookPath("iptables")
	if err != nil {
		return fmt.Errorf("iptables: 'iptables' executable could not be found")
	}

	return nil
}

func (impl *Implementation) Perform(op Operation, c firewall.Chain, r *Rule) error {
	if r.Target == "" {
		return fmt.Errorf("iptables: rule doesnt have a target")
	}

	table := fmt.Sprintf("--table=%v", c.Table)
	action := fmt.Sprintf("--%v=%v", op, c.Name)

	args := []string{"--wait", table, action}

	args = append(args, r.FlagForm()...)

	_, err := impl.Execute("iptables", args...)

	return err
}
<|endoftext|>


<commit_before>/*
 * This file is part of the kubevirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package rest

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"net/url"

	"github.com/emicklei/go-restful"
	"github.com/golang/mock/gomock"
	"github.com/gorilla/websocket"
	"github.com/libvirt/libvirt-go"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/client-go/tools/record"

	"kubevirt.io/kubevirt/pkg/logging"
	"kubevirt.io/kubevirt/pkg/virt-handler/virtwrap"
)

var _ = Describe("Console", func() {
	var handler http.Handler

	var mockConn *virtwrap.MockConnection
	var mockDomain *virtwrap.MockVirDomain
	var mockStream *virtwrap.MockStream
	var ctrl *gomock.Controller
	var mockRecorder *record.FakeRecorder
	var server *httptest.Server
	var wsUrl *url.URL

	logging.DefaultLogger().SetIOWriter(GinkgoWriter)

	dial := func(vm string, console string) *websocket.Conn {
		wsUrl.Scheme = "ws"
		wsUrl.Path = "/api/v1/console/" + vm
		wsUrl.RawQuery = "console=" + console
		c, _, err := websocket.DefaultDialer.Dial(wsUrl.String(), nil)
		Expect(err).ToNot(HaveOccurred())
		return c
	}

	get := func(vm string) (*http.Response, error) {
		wsUrl.Scheme = "http"
		wsUrl.Path = "/api/v1/console/" + vm
		return http.DefaultClient.Get(wsUrl.String())
	}

	BeforeEach(func() {
		var err error
		// Set up mocks
		ctrl = gomock.NewController(GinkgoT())
		mockConn = virtwrap.NewMockConnection(ctrl)
		mockDomain = virtwrap.NewMockVirDomain(ctrl)
		mockStream = virtwrap.NewMockStream(ctrl)
		mockRecorder = record.NewFakeRecorder(10)

		// Set up web service
		ws := new(restful.WebService)
		handler = http.Handler(restful.NewContainer().Add(ws))

		ws.Route(ws.GET("/api/v1/console/{name}").To(NewConsoleResource(mockConn).Console))
		server = httptest.NewServer(handler)
		wsUrl, err = url.Parse(server.URL)
		Expect(err).ToNot(HaveOccurred())
	})

	It("should return 404 if VM does not exist", func() {
		mockConn.EXPECT().LookupDomainByName("testvm").Return(nil, libvirt.Error{Code: libvirt.ERR_NO_DOMAIN})
		r, err := get("testvm")
		Expect(err).ToNot(HaveOccurred())
		Expect(r.StatusCode).To(Equal(http.StatusNotFound))
	})

	It("should return 500 if domain can't be looked up", func() {
		mockConn.EXPECT().LookupDomainByName("testvm").Return(nil, libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
		r, err := get("testvm")
		Expect(err).ToNot(HaveOccurred())
		Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
	})

	Context("with existing domain", func() {
		BeforeEach(func() {
			// Make sure that we always free the domain after use
			mockDomain.EXPECT().Free()
		})

		It("should return 500 if uid of domain can't be looked up", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return("", libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 500 if creating a stream fails", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(nil, libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 500 if opening a console connection fails", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(mockStream, nil)
			stream := &libvirt.Stream{}
			mockStream.EXPECT().UnderlyingStream().Return(stream)
			mockStream.EXPECT().Close()
			mockDomain.EXPECT().OpenConsole("", stream, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 400 if ws upgrade does not work", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(mockStream, nil)
			stream := &libvirt.Stream{}
			mockStream.EXPECT().UnderlyingStream().Return(stream)
			mockStream.EXPECT().Close()
			mockDomain.EXPECT().OpenConsole("", stream, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(nil)
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusBadRequest))
		})
		It("should proxy websocket traffic", func() {

			var in []byte
			var out []byte
			inBuf := bytes.NewBuffer(in)
			outBuf := bytes.NewBuffer(out)
			outBuf.WriteString("hello client!")
			stream := &fakeStream{in: inBuf, out: outBuf, s: &libvirt.Stream{}}

			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(stream, nil)
			mockDomain.EXPECT().OpenConsole("console0", stream.s, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(nil)

			con := dial("testvm", "console0")
			defer con.Close()
			err := con.WriteMessage(websocket.TextMessage, []byte("hello console!"))
			Expect(err).ToNot(HaveOccurred())

			// FIXME, there is somewhere a buffer or timeout which delays the actual send
			// Eventually(inBuf.String).Should(Equal("hello console!"))

			t, body, err := con.ReadMessage()
			Expect(t).To(Equal(websocket.TextMessage))
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(Equal([]byte("hello client!")))
		})

	})
	AfterEach(func() {
		ctrl.Finish()
		server.Close()
	})
})

type fakeStream struct {
	in  *bytes.Buffer
	out *bytes.Buffer
	s   *libvirt.Stream
}

func (s *fakeStream) Write(p []byte) (n int, err error) {
	return s.in.Write(p)
}

func (s *fakeStream) Read(p []byte) (n int, err error) {
	return s.out.Read(p)
}

func (s *fakeStream) Close() (e error) {
	return nil
}

func (s *fakeStream) UnderlyingStream() *libvirt.Stream {
	return s.s
}
<commit_msg>tests: Wait until console handler is done before checking mocks<commit_after>/*
 * This file is part of the kubevirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package rest

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"net/url"

	"github.com/emicklei/go-restful"
	"github.com/golang/mock/gomock"
	"github.com/gorilla/websocket"
	"github.com/libvirt/libvirt-go"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/client-go/tools/record"

	"kubevirt.io/kubevirt/pkg/logging"
	"kubevirt.io/kubevirt/pkg/virt-handler/virtwrap"
)

var _ = Describe("Console", func() {
	var handler http.Handler

	var mockConn *virtwrap.MockConnection
	var mockDomain *virtwrap.MockVirDomain
	var mockStream *virtwrap.MockStream
	var ctrl *gomock.Controller
	var mockRecorder *record.FakeRecorder
	var server *httptest.Server
	var wsUrl *url.URL
	var serverDone chan bool

	logging.DefaultLogger().SetIOWriter(GinkgoWriter)

	dial := func(vm string, console string) *websocket.Conn {
		wsUrl.Scheme = "ws"
		wsUrl.Path = "/api/v1/console/" + vm
		wsUrl.RawQuery = "console=" + console
		c, _, err := websocket.DefaultDialer.Dial(wsUrl.String(), nil)
		Expect(err).ToNot(HaveOccurred())
		return c
	}

	get := func(vm string) (*http.Response, error) {
		wsUrl.Scheme = "http"
		wsUrl.Path = "/api/v1/console/" + vm
		return http.DefaultClient.Get(wsUrl.String())
	}

	BeforeEach(func() {
		var err error
		// Set up mocks
		ctrl = gomock.NewController(GinkgoT())
		mockConn = virtwrap.NewMockConnection(ctrl)
		mockDomain = virtwrap.NewMockVirDomain(ctrl)
		mockStream = virtwrap.NewMockStream(ctrl)
		mockRecorder = record.NewFakeRecorder(10)

		// Set up web service
		ws := new(restful.WebService)
		handler = http.Handler(restful.NewContainer().Add(ws))

		// Give us a chance to detect when the request is done. Otherwise we
		// don't know when to check mock invokations
		serverDone = make(chan bool)
		waiter := func(request *restful.Request, response *restful.Response) {
			NewConsoleResource(mockConn).Console(request, response)
			close(serverDone)
		}
		ws.Route(ws.GET("/api/v1/console/{name}").To(waiter))
		server = httptest.NewServer(handler)
		wsUrl, err = url.Parse(server.URL)
		Expect(err).ToNot(HaveOccurred())
	})

	It("should return 404 if VM does not exist", func() {
		mockConn.EXPECT().LookupDomainByName("testvm").Return(nil, libvirt.Error{Code: libvirt.ERR_NO_DOMAIN})
		r, err := get("testvm")
		Expect(err).ToNot(HaveOccurred())
		Expect(r.StatusCode).To(Equal(http.StatusNotFound))
	})

	It("should return 500 if domain can't be looked up", func() {
		mockConn.EXPECT().LookupDomainByName("testvm").Return(nil, libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
		r, err := get("testvm")
		Expect(err).ToNot(HaveOccurred())
		Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
	})

	Context("with existing domain", func() {
		BeforeEach(func() {
			// Make sure that we always free the domain after use
			mockDomain.EXPECT().Free()
		})

		It("should return 500 if uid of domain can't be looked up", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return("", libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 500 if creating a stream fails", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(nil, libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 500 if opening a console connection fails", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(mockStream, nil)
			stream := &libvirt.Stream{}
			mockStream.EXPECT().UnderlyingStream().Return(stream)
			mockStream.EXPECT().Close()
			mockDomain.EXPECT().OpenConsole("", stream, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(libvirt.Error{Code: libvirt.ERR_INVALID_CONN})
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusInternalServerError))
		})
		It("should return 400 if ws upgrade does not work", func() {
			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(mockStream, nil)
			stream := &libvirt.Stream{}
			mockStream.EXPECT().UnderlyingStream().Return(stream)
			mockStream.EXPECT().Close()
			mockDomain.EXPECT().OpenConsole("", stream, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(nil)
			r, err := get("testvm")
			Expect(err).ToNot(HaveOccurred())
			Expect(r.StatusCode).To(Equal(http.StatusBadRequest))
		})
		It("should proxy websocket traffic", func() {

			var in []byte
			var out []byte
			inBuf := bytes.NewBuffer(in)
			outBuf := bytes.NewBuffer(out)
			outBuf.WriteString("hello client!")
			stream := &fakeStream{in: inBuf, out: outBuf, s: &libvirt.Stream{}}

			mockConn.EXPECT().LookupDomainByName("testvm").Return(mockDomain, nil)
			mockDomain.EXPECT().GetUUIDString().Return(string(uuid.NewUUID()), nil)
			mockConn.EXPECT().NewStream(libvirt.StreamFlags(0)).Return(stream, nil)
			mockDomain.EXPECT().OpenConsole("console0", stream.s, libvirt.DomainConsoleFlags(libvirt.DOMAIN_CONSOLE_FORCE)).Return(nil)

			con := dial("testvm", "console0")
			defer con.Close()
			err := con.WriteMessage(websocket.TextMessage, []byte("hello console!"))
			Expect(err).ToNot(HaveOccurred())

			// FIXME, there is somewhere a buffer or timeout which delays the actual send
			// Eventually(inBuf.String).Should(Equal("hello console!"))

			t, body, err := con.ReadMessage()
			Expect(t).To(Equal(websocket.TextMessage))
			Expect(err).ToNot(HaveOccurred())
			Expect(body).To(Equal([]byte("hello client!")))
		})

	})
	AfterEach(func() {
		server.Close()
		<-serverDone
		ctrl.Finish()
	})
})

type fakeStream struct {
	in  *bytes.Buffer
	out *bytes.Buffer
	s   *libvirt.Stream
}

func (s *fakeStream) Write(p []byte) (n int, err error) {
	return s.in.Write(p)
}

func (s *fakeStream) Read(p []byte) (n int, err error) {
	return s.out.Read(p)
}

func (s *fakeStream) Close() (e error) {
	return nil
}

func (s *fakeStream) UnderlyingStream() *libvirt.Stream {
	return s.s
}
<|endoftext|>


<commit_before>package s3manager_test

import (
	"context"
	"os"
	"testing"

	"github.com/admpub/nging/v3/application/dbschema"
	"github.com/admpub/nging/v3/application/library/s3manager/s3client"
	"github.com/stretchr/testify/assert"
)

func TestStat(t *testing.T) {
	return
	cfg := &dbschema.NgingCloudStorage{
		Key:      os.Getenv(`S3_KEY`),
		Secret:   os.Getenv(`S3_SECRET`),
		Secure:   `Y`,
		Region:   os.Getenv(`S3_REGION`),
		Bucket:   os.Getenv(`S3_BUCKET`),
		Endpoint: os.Getenv(`S3_ENDPOINT`),
		Baseurl:  os.Getenv(`S3_BASEURL`),
	}
	mgr, err := s3client.New(cfg, 1024000)
	if err != nil {
		panic(err)
	}
	exists, err := mgr.Exists(context.Background(), ``)
	if err != nil {
		panic(err)
	}
	assert.Equal(t, false, exists)
}
<commit_msg>update<commit_after>package s3manager_test

import (
	"context"
	"os"
	"path/filepath"
	"testing"

	"github.com/admpub/godotenv"
	"github.com/admpub/nging/v3/application/dbschema"
	"github.com/admpub/nging/v3/application/library/s3manager/s3client"
	"github.com/stretchr/testify/assert"
	"github.com/webx-top/echo"
)

func TestStat(t *testing.T) {
	return
	projectDir := filepath.Join(echo.Wd(), `../../../`)
	envFile := filepath.Join(projectDir, `.env`)
	err := godotenv.Overload(envFile)
	if err != nil {
		panic(err)
	}
	cfg := &dbschema.NgingCloudStorage{
		Key:      os.Getenv(`S3_KEY`),
		Secret:   os.Getenv(`S3_SECRET`),
		Secure:   `Y`,
		Region:   os.Getenv(`S3_REGION`),
		Bucket:   os.Getenv(`S3_BUCKET`),
		Endpoint: os.Getenv(`S3_ENDPOINT`),
		Baseurl:  os.Getenv(`S3_BASEURL`),
	}
	mgr, err := s3client.New(cfg, 1024000)
	if err != nil {
		panic(err)
	}
	exists, err := mgr.Exists(context.Background(), ``)
	if err != nil {
		panic(err)
	}
	assert.Equal(t, false, exists)
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"testing"
)

func TestIsURLHTTP(t *testing.T) {
	url := "http://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLHTTPS(t *testing.T) {
	url := "https://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLWrong(t *testing.T) {
	url := "somehost.com"

	b := isURL(url)
	expect(t, b, false)
}

func TestIsURLWrongTLD(t *testing.T) {
	url := "http://somehost."

	b := isURL(url)
	expect(t, b, false)
}
<commit_msg>test file exists<commit_after>package hoverfly

import (
	"testing"
)

func TestIsURLHTTP(t *testing.T) {
	url := "http://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLHTTPS(t *testing.T) {
	url := "https://somehost.com"

	b := isURL(url)
	expect(t, b, true)
}

func TestIsURLWrong(t *testing.T) {
	url := "somehost.com"

	b := isURL(url)
	expect(t, b, false)
}

func TestIsURLWrongTLD(t *testing.T) {
	url := "http://somehost."

	b := isURL(url)
	expect(t, b, false)
}

func TestFileExists(t *testing.T) {
	fp := "examples/exports/readthedocs.json"

	ex, err := exists(fp)
	expect(t, ex, true)
	expect(t, err, nil)
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v\n", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, file := range buildPkg.GoFiles {
		f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
		if err != nil {
			fmt.Fprintf(os.Stderr, "could not parse %q: %v\n", file, err)
			continue
		}
		for name := range f.Scope.Objects {
			if ast.IsExported(name) {
				exports[name] = true
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<commit_msg>go.tools/imports: recognize exports in cgo files<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package imports

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"

	"code.google.com/p/go.tools/astutil"
)

// importToGroup is a list of functions which map from an import path to
// a group number.
var importToGroup = []func(importPath string) (num int, ok bool){
	func(importPath string) (num int, ok bool) {
		if strings.HasPrefix(importPath, "appengine") {
			return 2, true
		}
		return
	},
	func(importPath string) (num int, ok bool) {
		if strings.Contains(importPath, ".") {
			return 1, true
		}
		return
	},
}

func importGroup(importPath string) int {
	for _, fn := range importToGroup {
		if n, ok := fn(importPath); ok {
			return n
		}
	}
	return 0
}

func fixImports(fset *token.FileSet, f *ast.File) (added []string, err error) {
	// refs are a set of possible package references currently unsatisfied by imports.
	// first key: either base package (e.g. "fmt") or renamed package
	// second key: referenced package symbol (e.g. "Println")
	refs := make(map[string]map[string]bool)

	// decls are the current package imports. key is base package or renamed package.
	decls := make(map[string]*ast.ImportSpec)

	// collect potential uses of packages.
	var visitor visitFn
	visitor = visitFn(func(node ast.Node) ast.Visitor {
		if node == nil {
			return visitor
		}
		switch v := node.(type) {
		case *ast.ImportSpec:
			if v.Name != nil {
				decls[v.Name.Name] = v
			} else {
				local := importPathToName(strings.Trim(v.Path.Value, `\"`))
				decls[local] = v
			}
		case *ast.SelectorExpr:
			xident, ok := v.X.(*ast.Ident)
			if !ok {
				break
			}
			if xident.Obj != nil {
				// if the parser can resolve it, it's not a package ref
				break
			}
			pkgName := xident.Name
			if refs[pkgName] == nil {
				refs[pkgName] = make(map[string]bool)
			}
			if decls[pkgName] == nil {
				refs[pkgName][v.Sel.Name] = true
			}
		}
		return visitor
	})
	ast.Walk(visitor, f)

	// Search for imports matching potential package references.
	searches := 0
	type result struct {
		ipath string
		name  string
		err   error
	}
	results := make(chan result)
	for pkgName, symbols := range refs {
		if len(symbols) == 0 {
			continue // skip over packages already imported
		}
		go func(pkgName string, symbols map[string]bool) {
			ipath, rename, err := findImport(pkgName, symbols)
			r := result{ipath: ipath, err: err}
			if rename {
				r.name = pkgName
			}
			results <- r
		}(pkgName, symbols)
		searches++
	}
	for i := 0; i < searches; i++ {
		result := <-results
		if result.err != nil {
			return nil, result.err
		}
		if result.ipath != "" {
			if result.name != "" {
				astutil.AddNamedImport(fset, f, result.name, result.ipath)
			} else {
				astutil.AddImport(fset, f, result.ipath)
			}
			added = append(added, result.ipath)
		}
	}

	// Nil out any unused ImportSpecs, to be removed in following passes
	unusedImport := map[string]bool{}
	for pkg, is := range decls {
		if refs[pkg] == nil && pkg != "_" && pkg != "." {
			unusedImport[strings.Trim(is.Path.Value, `"`)] = true
		}
	}
	for ipath := range unusedImport {
		if ipath == "C" {
			// Don't remove cgo stuff.
			continue
		}
		astutil.DeleteImport(fset, f, ipath)
	}

	return added, nil
}

// importPathToName returns the package name for the given import path.
var importPathToName = importPathToNameGoPath

// importPathToNameBasic assumes the package name is the base of import path.
func importPathToNameBasic(importPath string) (packageName string) {
	return path.Base(importPath)
}

// importPathToNameGoPath finds out the actual package name, as declared in its .go files.
// If there's a problem, it falls back to using importPathToNameBasic.
func importPathToNameGoPath(importPath string) (packageName string) {
	if buildPkg, err := build.Import(importPath, "", 0); err == nil {
		return buildPkg.Name
	} else {
		return importPathToNameBasic(importPath)
	}
}

type pkg struct {
	importpath string // full pkg import path, e.g. "net/http"
	dir        string // absolute file path to pkg directory e.g. "/usr/lib/go/src/fmt"
}

var pkgIndexOnce sync.Once

var pkgIndex struct {
	sync.Mutex
	m map[string][]pkg // shortname => []pkg, e.g "http" => "net/http"
}

// gate is a semaphore for limiting concurrency.
type gate chan struct{}

func (g gate) enter() { g <- struct{}{} }
func (g gate) leave() { <-g }

// fsgate protects the OS & filesystem from too much concurrency.
// Too much disk I/O -> too many threads -> swapping and bad scheduling.
var fsgate = make(gate, 8)

func loadPkgIndex() {
	pkgIndex.Lock()
	pkgIndex.m = make(map[string][]pkg)
	pkgIndex.Unlock()

	var wg sync.WaitGroup
	for _, path := range build.Default.SrcDirs() {
		fsgate.enter()
		f, err := os.Open(path)
		if err != nil {
			fsgate.leave()
			fmt.Fprint(os.Stderr, err)
			continue
		}
		children, err := f.Readdir(-1)
		f.Close()
		fsgate.leave()
		if err != nil {
			fmt.Fprint(os.Stderr, err)
			continue
		}
		for _, child := range children {
			if child.IsDir() {
				wg.Add(1)
				go func(path, name string) {
					defer wg.Done()
					loadPkg(&wg, path, name)
				}(path, child.Name())
			}
		}
	}
	wg.Wait()
}

func loadPkg(wg *sync.WaitGroup, root, pkgrelpath string) {
	importpath := filepath.ToSlash(pkgrelpath)
	dir := filepath.Join(root, importpath)

	fsgate.enter()
	defer fsgate.leave()
	pkgDir, err := os.Open(dir)
	if err != nil {
		return
	}
	children, err := pkgDir.Readdir(-1)
	pkgDir.Close()
	if err != nil {
		return
	}
	// hasGo tracks whether a directory actually appears to be a
	// Go source code directory. If $GOPATH == $HOME, and
	// $HOME/src has lots of other large non-Go projects in it,
	// then the calls to importPathToName below can be expensive.
	hasGo := false
	for _, child := range children {
		name := child.Name()
		if name == "" {
			continue
		}
		if c := name[0]; c == '.' || ('0' <= c && c <= '9') {
			continue
		}
		if strings.HasSuffix(name, ".go") {
			hasGo = true
		}
		if child.IsDir() {
			wg.Add(1)
			go func(root, name string) {
				defer wg.Done()
				loadPkg(wg, root, name)
			}(root, filepath.Join(importpath, name))
		}
	}
	if hasGo {
		shortName := importPathToName(importpath)
		pkgIndex.Lock()
		pkgIndex.m[shortName] = append(pkgIndex.m[shortName], pkg{
			importpath: importpath,
			dir:        dir,
		})
		pkgIndex.Unlock()
	}

}

// loadExports returns a list exports for a package.
var loadExports = loadExportsGoPath

func loadExportsGoPath(dir string) map[string]bool {
	exports := make(map[string]bool)
	buildPkg, err := build.ImportDir(dir, 0)
	if err != nil {
		if strings.Contains(err.Error(), "no buildable Go source files in") {
			return nil
		}
		fmt.Fprintf(os.Stderr, "could not import %q: %v\n", dir, err)
		return nil
	}
	fset := token.NewFileSet()
	for _, files := range [...][]string{buildPkg.GoFiles, buildPkg.CgoFiles} {
		for _, file := range files {
			f, err := parser.ParseFile(fset, filepath.Join(dir, file), nil, 0)
			if err != nil {
				fmt.Fprintf(os.Stderr, "could not parse %q: %v\n", file, err)
				continue
			}
			for name := range f.Scope.Objects {
				if ast.IsExported(name) {
					exports[name] = true
				}
			}
		}
	}
	return exports
}

// findImport searches for a package with the given symbols.
// If no package is found, findImport returns "".
// Declared as a variable rather than a function so goimports can be easily
// extended by adding a file with an init function.
var findImport = findImportGoPath

func findImportGoPath(pkgName string, symbols map[string]bool) (string, bool, error) {
	// Fast path for the standard library.
	// In the common case we hopefully never have to scan the GOPATH, which can
	// be slow with moving disks.
	if pkg, rename, ok := findImportStdlib(pkgName, symbols); ok {
		return pkg, rename, nil
	}

	// TODO(sameer): look at the import lines for other Go files in the
	// local directory, since the user is likely to import the same packages
	// in the current Go file.  Return rename=true when the other Go files
	// use a renamed package that's also used in the current file.

	pkgIndexOnce.Do(loadPkgIndex)

	// Collect exports for packages with matching names.
	var wg sync.WaitGroup
	var pkgsMu sync.Mutex // guards pkgs
	// full importpath => exported symbol => True
	// e.g. "net/http" => "Client" => True
	pkgs := make(map[string]map[string]bool)
	pkgIndex.Lock()
	for _, pkg := range pkgIndex.m[pkgName] {
		wg.Add(1)
		go func(importpath, dir string) {
			defer wg.Done()
			exports := loadExports(dir)
			if exports != nil {
				pkgsMu.Lock()
				pkgs[importpath] = exports
				pkgsMu.Unlock()
			}
		}(pkg.importpath, pkg.dir)
	}
	pkgIndex.Unlock()
	wg.Wait()

	// Filter out packages missing required exported symbols.
	for symbol := range symbols {
		for importpath, exports := range pkgs {
			if !exports[symbol] {
				delete(pkgs, importpath)
			}
		}
	}
	if len(pkgs) == 0 {
		return "", false, nil
	}

	// If there are multiple candidate packages, the shortest one wins.
	// This is a heuristic to prefer the standard library (e.g. "bytes")
	// over e.g. "github.com/foo/bar/bytes".
	shortest := ""
	for importPath := range pkgs {
		if shortest == "" || len(importPath) < len(shortest) {
			shortest = importPath
		}
	}
	return shortest, false, nil
}

type visitFn func(node ast.Node) ast.Visitor

func (fn visitFn) Visit(node ast.Node) ast.Visitor {
	return fn(node)
}

func findImportStdlib(shortPkg string, symbols map[string]bool) (importPath string, rename, ok bool) {
	for symbol := range symbols {
		path := stdlib[shortPkg+"."+symbol]
		if path == "" {
			return "", false, false
		}
		if importPath != "" && importPath != path {
			// Ambiguous. Symbols pointed to different things.
			return "", false, false
		}
		importPath = path
	}
	return importPath, false, importPath != ""
}
<|endoftext|>


<commit_before>package joystick

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/veandco/go-sdl2/sdl"
	"github.com/hybridgroup/gobot"
)

var _ gobot.Driver = (*JoystickDriver)(nil)

// JoystickDriver represents a joystick
type JoystickDriver struct {
	name       string
	interval   time.Duration
	connection gobot.Connection
	configPath string
	config     joystickConfig
	poll       func() sdl.Event
	halt       chan bool
	gobot.Eventer
}

// pair is a JSON representation of name and id
type pair struct {
	Name string `json:"name"`
	ID   int    `json:"id"`
}

// hat is a JSON representation of hat, name and id
type hat struct {
	Hat  int    `json:"hat"`
	Name string `json:"name"`
	ID   int    `json:"id"`
}

// joystickConfig is a JSON representation of configuration values
type joystickConfig struct {
	Name    string `json:"name"`
	GUID    string `json:"guid"`
	Axis    []pair `json:"axis"`
	Buttons []pair `json:"buttons"`
	Hats    []hat  `json:"Hats"`
}

// NewJoystickDriver returns a new JoystickDriver with a polling interval of
// 10 Milliseconds given a JoystickAdaptor, name and json button configuration
// file location.
//
// Optinally accepts:
//  time.Duration: Interval at which the JoystickDriver is polled for new information
func NewJoystickDriver(a *JoystickAdaptor, name string, config string, v ...time.Duration) *JoystickDriver {
	d := &JoystickDriver{
		name:       name,
		connection: a,
		Eventer:    gobot.NewEventer(),
		configPath: config,
		poll: func() sdl.Event {
			return sdl.PollEvent()
		},
		interval: 10 * time.Millisecond,
		halt:     make(chan bool, 0),
	}

	if len(v) > 0 {
		d.interval = v[0]
	}

	d.AddEvent("error")
	return d
}

// Name returns the JoystickDrivers name
func (j *JoystickDriver) Name() string { return j.name }

// Connection returns the JoystickDrivers connection
func (j *JoystickDriver) Connection() gobot.Connection { return j.connection }

// adaptor returns joystick adaptor
func (j *JoystickDriver) adaptor() *JoystickAdaptor {
	return j.Connection().(*JoystickAdaptor)
}

// Start and polls the state of the joystick at the given interval.
//
// Emits the Events:
//	Error error - On button error
//	Events defined in the json button configuration file.
//	They will have the format:
//		[button]_press
//		[button]_release
//		[axis]
func (j *JoystickDriver) Start() (errs []error) {
	file, err := ioutil.ReadFile(j.configPath)
	if err != nil {
		return []error{err}
	}

	var jsontype joystickConfig
	json.Unmarshal(file, &jsontype)
	j.config = jsontype

	for _, value := range j.config.Buttons {
		j.AddEvent(fmt.Sprintf("%s_press", value.Name))
		j.AddEvent(fmt.Sprintf("%s_release", value.Name))
	}
	for _, value := range j.config.Axis {
		j.AddEvent(value.Name)
	}
	for _, value := range j.config.Hats {
		j.AddEvent(value.Name)
	}

	go func() {
		for {
			event := j.poll()
			if event != nil {
				if err = j.handleEvent(event); err != nil {
					gobot.Publish(j.Event("error"), err)
				}
			}
			select {
			case <-time.After(j.interval):
			case <-j.halt:
				return
			}
		}
	}()
	return
}

// Halt stops joystick driver
func (j *JoystickDriver) Halt() (errs []error) {
	j.halt <- true
	return
}

// HandleEvent publishes an specific event according to data received
func (j *JoystickDriver) handleEvent(event sdl.Event) error {
	switch data := event.(type) {
	case *sdl.JoyAxisEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			axis := j.findName(data.Axis, j.config.Axis)
			if axis == "" {
				return fmt.Errorf("Unknown Axis: %v", data.Axis)
			}
			gobot.Publish(j.Event(axis), data.Value)
		}
	case *sdl.JoyButtonEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			button := j.findName(data.Button, j.config.Buttons)
			if button == "" {
				return fmt.Errorf("Unknown Button: %v", data.Button)
			}
			if data.State == 1 {
				gobot.Publish(j.Event(fmt.Sprintf("%s_press", button)), nil)
			}
			gobot.Publish(j.Event(fmt.Sprintf("%s_release", button)), nil)
		}
	case *sdl.JoyHatEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			hat := j.findHatName(data.Value, data.Hat, j.config.Hats)
			if hat == "" {
				return fmt.Errorf("Unknown Hat: %v %v", data.Hat, data.Value)
			}
			gobot.Publish(j.Event(hat), true)
		}
	}
	return nil
}

func (j *JoystickDriver) findName(id uint8, list []pair) string {
	for _, value := range list {
		if int(id) == value.ID {
			return value.Name
		}
	}
	return ""
}

// findHatName returns name from hat found by id in provided list
func (j *JoystickDriver) findHatName(id uint8, hat uint8, list []hat) string {
	for _, lHat := range list {
		if int(id) == lHat.ID && int(hat) == lHat.Hat {
			return lHat.Name
		}
	}
	return ""
}
<commit_msg>Fix for analog (quick changes lag)<commit_after>package joystick

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/veandco/go-sdl2/sdl"
	"github.com/hybridgroup/gobot"
)

var _ gobot.Driver = (*JoystickDriver)(nil)

// JoystickDriver represents a joystick
type JoystickDriver struct {
	name       string
	interval   time.Duration
	connection gobot.Connection
	configPath string
	config     joystickConfig
	poll       func() sdl.Event
	halt       chan bool
	gobot.Eventer
}

// pair is a JSON representation of name and id
type pair struct {
	Name string `json:"name"`
	ID   int    `json:"id"`
}

// hat is a JSON representation of hat, name and id
type hat struct {
	Hat  int    `json:"hat"`
	Name string `json:"name"`
	ID   int    `json:"id"`
}

// joystickConfig is a JSON representation of configuration values
type joystickConfig struct {
	Name    string `json:"name"`
	GUID    string `json:"guid"`
	Axis    []pair `json:"axis"`
	Buttons []pair `json:"buttons"`
	Hats    []hat  `json:"Hats"`
}

// NewJoystickDriver returns a new JoystickDriver with a polling interval of
// 10 Milliseconds given a JoystickAdaptor, name and json button configuration
// file location.
//
// Optinally accepts:
//  time.Duration: Interval at which the JoystickDriver is polled for new information
func NewJoystickDriver(a *JoystickAdaptor, name string, config string, v ...time.Duration) *JoystickDriver {
	d := &JoystickDriver{
		name:       name,
		connection: a,
		Eventer:    gobot.NewEventer(),
		configPath: config,
		poll: func() sdl.Event {
			return sdl.PollEvent()
		},
		interval: 10 * time.Millisecond,
		halt:     make(chan bool, 0),
	}

	if len(v) > 0 {
		d.interval = v[0]
	}

	d.AddEvent("error")
	return d
}

// Name returns the JoystickDrivers name
func (j *JoystickDriver) Name() string { return j.name }

// Connection returns the JoystickDrivers connection
func (j *JoystickDriver) Connection() gobot.Connection { return j.connection }

// adaptor returns joystick adaptor
func (j *JoystickDriver) adaptor() *JoystickAdaptor {
	return j.Connection().(*JoystickAdaptor)
}

// Start and polls the state of the joystick at the given interval.
//
// Emits the Events:
//	Error error - On button error
//	Events defined in the json button configuration file.
//	They will have the format:
//		[button]_press
//		[button]_release
//		[axis]
func (j *JoystickDriver) Start() (errs []error) {
	file, err := ioutil.ReadFile(j.configPath)
	if err != nil {
		return []error{err}
	}

	var jsontype joystickConfig
	json.Unmarshal(file, &jsontype)
	j.config = jsontype

	for _, value := range j.config.Buttons {
		j.AddEvent(fmt.Sprintf("%s_press", value.Name))
		j.AddEvent(fmt.Sprintf("%s_release", value.Name))
	}
	for _, value := range j.config.Axis {
		j.AddEvent(value.Name)
	}
	for _, value := range j.config.Hats {
		j.AddEvent(value.Name)
	}

	go func() {
		for {
			for event := j.poll(); event != nil; event = j.poll() {
				if err = j.handleEvent(event); err != nil {
					gobot.Publish(j.Event("error"), err)
				}
			}
			select {
			case <-time.After(j.interval):
			case <-j.halt:
				return
			}
		}
	}()
	return
}

// Halt stops joystick driver
func (j *JoystickDriver) Halt() (errs []error) {
	j.halt <- true
	return
}

// HandleEvent publishes an specific event according to data received
func (j *JoystickDriver) handleEvent(event sdl.Event) error {
	switch data := event.(type) {
	case *sdl.JoyAxisEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			axis := j.findName(data.Axis, j.config.Axis)
			if axis == "" {
				return fmt.Errorf("Unknown Axis: %v", data.Axis)
			}
			gobot.Publish(j.Event(axis), data.Value)
		}
	case *sdl.JoyButtonEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			button := j.findName(data.Button, j.config.Buttons)
			if button == "" {
				return fmt.Errorf("Unknown Button: %v", data.Button)
			}
			if data.State == 1 {
				gobot.Publish(j.Event(fmt.Sprintf("%s_press", button)), nil)
			}
			gobot.Publish(j.Event(fmt.Sprintf("%s_release", button)), nil)
		}
	case *sdl.JoyHatEvent:
		if data.Which == j.adaptor().joystick.InstanceID() {
			hat := j.findHatName(data.Value, data.Hat, j.config.Hats)
			if hat == "" {
				return fmt.Errorf("Unknown Hat: %v %v", data.Hat, data.Value)
			}
			gobot.Publish(j.Event(hat), true)
		}
	}
	return nil
}

func (j *JoystickDriver) findName(id uint8, list []pair) string {
	for _, value := range list {
		if int(id) == value.ID {
			return value.Name
		}
	}
	return ""
}

// findHatName returns name from hat found by id in provided list
func (j *JoystickDriver) findHatName(id uint8, hat uint8, list []hat) string {
	for _, lHat := range list {
		if int(id) == lHat.ID && int(hat) == lHat.Hat {
			return lHat.Name
		}
	}
	return ""
}
<|endoftext|>