<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 watch

import (
	"strings"
	"time"

	"github.com/jeevatkm/go-model"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"
	"k8s.io/client-go/util/workqueue"

	kubev1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/controller"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/logging"
	"kubevirt.io/kubevirt/pkg/virt-controller/services"
)

func NewVMController(restClient *rest.RESTClient, vmService services.VMService, queue workqueue.RateLimitingInterface, vmCache cache.Store, vmInformer cache.SharedIndexInformer, podInformer cache.SharedIndexInformer, recorder record.EventRecorder, clientset kubecli.KubevirtClient) *VMController {
	return &VMController{
		restClient:  restClient,
		vmService:   vmService,
		queue:       queue,
		store:       vmCache,
		vmInformer:  vmInformer,
		podInformer: podInformer,
		recorder:    recorder,
		clientset:   clientset,
	}
}

type VMController struct {
	restClient  *rest.RESTClient
	vmService   services.VMService
	clientset   kubecli.KubevirtClient
	queue       workqueue.RateLimitingInterface
	store       cache.Store
	vmInformer  cache.SharedIndexInformer
	podInformer cache.SharedIndexInformer
	recorder    record.EventRecorder
}

func (c *VMController) Run(threadiness int, stopCh chan struct{}) {
	defer controller.HandlePanic()
	defer c.queue.ShutDown()
	logging.DefaultLogger().Info().Msg("Starting controller.")

	// Wait for cache sync before we start the pod controller
	cache.WaitForCacheSync(stopCh, c.vmInformer.HasSynced, c.podInformer.HasSynced)

	// Start the actual work
	for i := 0; i < threadiness; i++ {
		go wait.Until(c.runWorker, time.Second, stopCh)
	}

	<-stopCh
	logging.DefaultLogger().Info().Msg("Stopping controller.")
}

func (c *VMController) runWorker() {
	for c.Execute() {
	}
}

func (c *VMController) Execute() bool {
	key, quit := c.queue.Get()
	if quit {
		return false
	}
	defer c.queue.Done(key)
	if err := c.execute(key.(string)); err != nil {
		logging.DefaultLogger().Info().Reason(err).Msgf("reenqueuing VM %v", key)
		c.queue.AddRateLimited(key)
	} else {
		logging.DefaultLogger().Info().V(4).Msgf("processed VM %v", key)
		c.queue.Forget(key)
	}
	return true
}

func (c *VMController) execute(key string) error {

	// Fetch the latest Vm state from cache
	obj, exists, err := c.store.GetByKey(key)

	if err != nil {
		return err
	}

	// Retrieve the VM
	var vm *kubev1.VirtualMachine
	if !exists {
		namespace, name, err := cache.SplitMetaNamespaceKey(key)
		if err != nil {
			return err
		}
		vm = kubev1.NewVMReferenceFromNameWithNS(namespace, name)
	} else {
		vm = obj.(*kubev1.VirtualMachine)
	}
	logger := logging.DefaultLogger().Object(vm)

	if !exists {
		// Delete VM Pods
		err := c.vmService.DeleteVMPod(vm)
		if err != nil {
			logger.Error().Reason(err).Msg("Deleting VM target Pod failed.")
			return err
		}
		logger.Info().Msg("Deleting VM target Pod succeeded.")
		return nil
	}

	switch vm.Status.Phase {
	case kubev1.VmPhaseUnset, kubev1.Pending:
		// Schedule the VM

		// Deep copy the object, so that we can safely manipulate it
		vmCopy := kubev1.VirtualMachine{}
		model.Copy(&vmCopy, vm)
		logger = logging.DefaultLogger().Object(&vmCopy)

		// Check if there are already outdated VM Pods
		pods, err := c.vmService.GetRunningVMPods(&vmCopy)
		if err != nil {
			logger.Error().Reason(err).Msg("Fetching VM pods failed.")
			return err
		}

		// If there are already pods, delete them before continuing ...
		if len(pods.Items) > 0 {
			logger.Error().Msg("VM Pods do already exist, will clean up before continuing.")
			if err := c.vmService.DeleteVMPod(&vmCopy); err != nil {
				logger.Error().Reason(err).Msg("Deleting VM pods failed.")
				return err
			}
		}

		// Defaulting and setting constants
		// TODO move defaulting to virt-api
		// TODO move constants to virt-handler and remove from the spec
		if vmCopy.Spec.Domain == nil {
			spec := kubev1.NewMinimalDomainSpec()
			vmCopy.Spec.Domain = spec
		}

		// TODO when we move this to virt-api, we have to block that they are set on POST or changed on PUT
		graphics := vmCopy.Spec.Domain.Devices.Graphics
		for i, _ := range graphics {
			if strings.ToLower(graphics[i].Type) == "spice" {
				graphics[i].Port = int32(-1)
				graphics[i].Listen = kubev1.Listen{
					Address: "0.0.0.0",
					Type:    "address",
				}

			}
		}

		// Create a Pod which will be the VM destination
		if err := c.vmService.StartVMPod(&vmCopy); err != nil {
			logger.Error().Reason(err).Msg("Defining a target pod for the VM failed.")
			return err
		}

		// Mark the VM as "initialized". After the created Pod above is scheduled by
		// kubernetes, virt-handler can take over.
		vmCopy.Status.Phase = kubev1.Scheduling
		if err := c.restClient.Put().Resource("virtualmachines").Body(&vmCopy).Name(vmCopy.ObjectMeta.Name).Namespace(vmCopy.ObjectMeta.Namespace).Do().Error(); err != nil {
			logger.Error().Reason(err).Msg("Updating the VM state to 'Scheduling' failed.")
			return err
		}
		logger.Info().Msg("Handing over the VM to the scheduler succeeded.")
	case kubev1.Scheduling:
		// Target Pod for the VM was already created, check if it is  running and update the VM to Scheduled

		// Deep copy the object, so that we can safely manipulate it
		vmCopy := kubev1.VirtualMachine{}
		model.Copy(&vmCopy, vm)
		logger = logging.DefaultLogger().Object(&vmCopy)

		pods, err := c.vmService.GetRunningVMPods(&vmCopy)
		if err != nil {
			logger.Error().Reason(err).Msg("Fetching VM pods failed.")
			return err
		}

		//TODO, we can improve the pod checks here, for now they are as good as before the refactoring
		// So far, no running Pod found, we will sooner or later get a started event.
		// If not, something is wrong and the VM, stay stuck in the Scheduling phase
		if len(pods.Items) == 0 {
			logger.Info().V(3).Msg("No VM target pod in running state found.")
			return nil
		}

		// Whatever is going on here, I don't know what to do, don't reprocess this
		if len(pods.Items) > 1 {
			logger.Error().V(3).Msg("More than one VM target pods found.")
			return nil
		}

		// Pod is not yet running
		if pods.Items[0].Status.Phase != k8sv1.PodRunning {
			return nil
		}

		if verifyReadiness(&pods.Items[0]) == false {
			logger.Info().V(2).Msg("Waiting on all virt-launcher containers to be marked ready")
			return nil
		}

		// VM got scheduled
		vmCopy.Status.Phase = kubev1.Scheduled

		// FIXME we store this in the metadata since field selectors are currently not working for TPRs
		if vmCopy.GetObjectMeta().GetLabels() == nil {
			vmCopy.ObjectMeta.Labels = map[string]string{}
		}
		vmCopy.ObjectMeta.Labels[kubev1.NodeNameLabel] = pods.Items[0].Spec.NodeName
		vmCopy.Status.NodeName = pods.Items[0].Spec.NodeName
		if _, err := c.vmService.PutVm(&vmCopy); err != nil {
			logger.Error().Reason(err).Msg("Updating the VM state to 'Scheduled' failed.")
			return err
		}
		logger.Info().Msgf("VM successfully scheduled to %s.", vmCopy.Status.NodeName)
	}
	return nil
}

func verifyReadiness(pod *k8sv1.Pod) bool {
	for _, containerStatus := range pod.Status.ContainerStatuses {
		if containerStatus.Ready == false {
			return false
		}
	}

	return true
}

func vmLabelHandler(vmQueue workqueue.RateLimitingInterface) func(obj interface{}) {
	return func(obj interface{}) {
		phase := obj.(*k8sv1.Pod).Status.Phase
		namespace := obj.(*k8sv1.Pod).ObjectMeta.Namespace
		appLabel, hasAppLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.AppLabel]
		domainLabel, hasDomainLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.DomainLabel]
		_, hasMigrationLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.MigrationLabel]

		if phase != k8sv1.PodRunning {
			// Filter out non running pods from Queue
			return
		} else if hasMigrationLabel {
			// Filter out migration target pods
			return
		} else if hasDomainLabel == false || hasAppLabel == false {
			// missing required labels
			return
		} else if appLabel != "virt-launcher" {
			// ensure we're looking just for virt-launcher pods
			return
		}
		vmQueue.Add(namespace + "/" + domainLabel)
	}
}
<commit_msg>Don't create VM pod until previous one is not running<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 watch

import (
	"strings"
	"time"

	"github.com/jeevatkm/go-model"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/tools/record"
	"k8s.io/client-go/util/workqueue"

	kubev1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/controller"
	"kubevirt.io/kubevirt/pkg/kubecli"
	"kubevirt.io/kubevirt/pkg/logging"
	"kubevirt.io/kubevirt/pkg/virt-controller/services"
)

func NewVMController(restClient *rest.RESTClient, vmService services.VMService, queue workqueue.RateLimitingInterface, vmCache cache.Store, vmInformer cache.SharedIndexInformer, podInformer cache.SharedIndexInformer, recorder record.EventRecorder, clientset kubecli.KubevirtClient) *VMController {
	return &VMController{
		restClient:  restClient,
		vmService:   vmService,
		queue:       queue,
		store:       vmCache,
		vmInformer:  vmInformer,
		podInformer: podInformer,
		recorder:    recorder,
		clientset:   clientset,
	}
}

type VMController struct {
	restClient  *rest.RESTClient
	vmService   services.VMService
	clientset   kubecli.KubevirtClient
	queue       workqueue.RateLimitingInterface
	store       cache.Store
	vmInformer  cache.SharedIndexInformer
	podInformer cache.SharedIndexInformer
	recorder    record.EventRecorder
}

func (c *VMController) Run(threadiness int, stopCh chan struct{}) {
	defer controller.HandlePanic()
	defer c.queue.ShutDown()
	logging.DefaultLogger().Info().Msg("Starting controller.")

	// Wait for cache sync before we start the pod controller
	cache.WaitForCacheSync(stopCh, c.vmInformer.HasSynced, c.podInformer.HasSynced)

	// Start the actual work
	for i := 0; i < threadiness; i++ {
		go wait.Until(c.runWorker, time.Second, stopCh)
	}

	<-stopCh
	logging.DefaultLogger().Info().Msg("Stopping controller.")
}

func (c *VMController) runWorker() {
	for c.Execute() {
	}
}

func (c *VMController) Execute() bool {
	key, quit := c.queue.Get()
	if quit {
		return false
	}
	defer c.queue.Done(key)
	isPending, err := c.execute(key.(string))

	if err != nil {
		logging.DefaultLogger().Info().Reason(err).Msgf("reenqueuing VM %v", key)
		c.queue.AddRateLimited(key)
	} else if isPending {
		logging.DefaultLogger().Info().V(4).Msgf("reenqueuing pending VM %v", key)
		c.queue.AddAfter(key, 1*time.Second)
		c.queue.Forget(key)
	} else {
		logging.DefaultLogger().Info().V(4).Msgf("processed VM %v", key)
		c.queue.Forget(key)
	}
	return true
}

func (c *VMController) execute(key string) (bool, error) {

	// Fetch the latest Vm state from cache
	obj, exists, err := c.store.GetByKey(key)

	if err != nil {
		return false, err
	}

	// Retrieve the VM
	var vm *kubev1.VirtualMachine
	if !exists {
		namespace, name, err := cache.SplitMetaNamespaceKey(key)
		if err != nil {
			return false, err
		}
		vm = kubev1.NewVMReferenceFromNameWithNS(namespace, name)
	} else {
		vm = obj.(*kubev1.VirtualMachine)
	}
	logger := logging.DefaultLogger().Object(vm)

	if !exists {
		// Delete VM Pods
		err := c.vmService.DeleteVMPod(vm)
		if err != nil {
			logger.Error().Reason(err).Msg("Deleting VM target Pod failed.")
			return false, err
		}
		logger.Info().Msg("Deleting VM target Pod succeeded.")
		return false, nil
	}

	switch vm.Status.Phase {
	case kubev1.VmPhaseUnset, kubev1.Pending:
		// Schedule the VM

		// Deep copy the object, so that we can safely manipulate it
		vmCopy := kubev1.VirtualMachine{}
		model.Copy(&vmCopy, vm)
		logger = logging.DefaultLogger().Object(&vmCopy)

		// Check if there are already outdated VM Pods
		pods, err := c.vmService.GetRunningVMPods(&vmCopy)
		if err != nil {
			logger.Error().Reason(err).Msg("Fetching VM pods failed.")
			return false, err
		}

		// If there are already pods, delete them before continuing ...
		if len(pods.Items) > 0 {
			logger.Error().Msg("VM Pods do already exist, will clean up before continuing.")
			if err := c.vmService.DeleteVMPod(&vmCopy); err != nil {
				logger.Error().Reason(err).Msg("Deleting VM pods failed.")
				return false, err
			}
			return true, nil
		}

		// Defaulting and setting constants
		// TODO move defaulting to virt-api
		// TODO move constants to virt-handler and remove from the spec
		if vmCopy.Spec.Domain == nil {
			spec := kubev1.NewMinimalDomainSpec()
			vmCopy.Spec.Domain = spec
		}

		// TODO when we move this to virt-api, we have to block that they are set on POST or changed on PUT
		graphics := vmCopy.Spec.Domain.Devices.Graphics
		for i, _ := range graphics {
			if strings.ToLower(graphics[i].Type) == "spice" {
				graphics[i].Port = int32(-1)
				graphics[i].Listen = kubev1.Listen{
					Address: "0.0.0.0",
					Type:    "address",
				}

			}
		}

		// Create a Pod which will be the VM destination
		if err := c.vmService.StartVMPod(&vmCopy); err != nil {
			logger.Error().Reason(err).Msg("Defining a target pod for the VM failed.")
			return false, err
		}

		// Mark the VM as "initialized". After the created Pod above is scheduled by
		// kubernetes, virt-handler can take over.
		vmCopy.Status.Phase = kubev1.Scheduling
		if err := c.restClient.Put().Resource("virtualmachines").Body(&vmCopy).Name(vmCopy.ObjectMeta.Name).Namespace(vmCopy.ObjectMeta.Namespace).Do().Error(); err != nil {
			logger.Error().Reason(err).Msg("Updating the VM state to 'Scheduling' failed.")
			return false, err
		}
		logger.Info().Msg("Handing over the VM to the scheduler succeeded.")
	case kubev1.Scheduling:
		// Target Pod for the VM was already created, check if it is  running and update the VM to Scheduled

		// Deep copy the object, so that we can safely manipulate it
		vmCopy := kubev1.VirtualMachine{}
		model.Copy(&vmCopy, vm)
		logger = logging.DefaultLogger().Object(&vmCopy)

		pods, err := c.vmService.GetRunningVMPods(&vmCopy)
		if err != nil {
			logger.Error().Reason(err).Msg("Fetching VM pods failed.")
			return false, err
		}

		//TODO, we can improve the pod checks here, for now they are as good as before the refactoring
		// So far, no running Pod found, we will sooner or later get a started event.
		// If not, something is wrong and the VM, stay stuck in the Scheduling phase
		if len(pods.Items) == 0 {
			logger.Info().V(3).Msg("No VM target pod in running state found.")
			return false, nil
		}

		// Whatever is going on here, I don't know what to do, don't reprocess this
		if len(pods.Items) > 1 {
			logger.Error().V(3).Msg("More than one VM target pods found.")
			return false, nil
		}

		// Pod is not yet running
		if pods.Items[0].Status.Phase != k8sv1.PodRunning {
			return false, nil
		}

		if verifyReadiness(&pods.Items[0]) == false {
			logger.Info().V(2).Msg("Waiting on all virt-launcher containers to be marked ready")
			return false, nil
		}

		// VM got scheduled
		vmCopy.Status.Phase = kubev1.Scheduled

		// FIXME we store this in the metadata since field selectors are currently not working for TPRs
		if vmCopy.GetObjectMeta().GetLabels() == nil {
			vmCopy.ObjectMeta.Labels = map[string]string{}
		}
		vmCopy.ObjectMeta.Labels[kubev1.NodeNameLabel] = pods.Items[0].Spec.NodeName
		vmCopy.Status.NodeName = pods.Items[0].Spec.NodeName
		if _, err := c.vmService.PutVm(&vmCopy); err != nil {
			logger.Error().Reason(err).Msg("Updating the VM state to 'Scheduled' failed.")
			return false, err
		}
		logger.Info().Msgf("VM successfully scheduled to %s.", vmCopy.Status.NodeName)
	}
	return false, nil
}

func verifyReadiness(pod *k8sv1.Pod) bool {
	for _, containerStatus := range pod.Status.ContainerStatuses {
		if containerStatus.Ready == false {
			return false
		}
	}

	return true
}

func vmLabelHandler(vmQueue workqueue.RateLimitingInterface) func(obj interface{}) {
	return func(obj interface{}) {
		phase := obj.(*k8sv1.Pod).Status.Phase
		namespace := obj.(*k8sv1.Pod).ObjectMeta.Namespace
		appLabel, hasAppLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.AppLabel]
		domainLabel, hasDomainLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.DomainLabel]
		_, hasMigrationLabel := obj.(*k8sv1.Pod).ObjectMeta.Labels[kubev1.MigrationLabel]

		if phase != k8sv1.PodRunning {
			// Filter out non running pods from Queue
			return
		} else if hasMigrationLabel {
			// Filter out migration target pods
			return
		} else if hasDomainLabel == false || hasAppLabel == false {
			// missing required labels
			return
		} else if appLabel != "virt-launcher" {
			// ensure we're looking just for virt-launcher pods
			return
		}
		vmQueue.Add(namespace + "/" + domainLabel)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/miekg/dns"
	"gloon/mem_rs"
	"gloon/record_set"
	"gloon/redis_rs"
	"log"
	"regexp"
)

type Server struct {
	*dns.Server
	*record_set.RecordSet
	resolver *Resolver
	settings *Settings
}

var split_rex = regexp.MustCompile("[:= ]")

func NewServer(addr string, settings *Settings) (s *Server, err error) {
	s = &Server{}
	s.settings = settings

	// Set up the record store
	var store record_set.RecordStore
	switch settings.Store {
	case "redis":
		store, err = redis_rs.Create(settings.StoreOpts)
		if err != nil {
			log.Fatalf("Unable to create redis record set: %s", err.Error())
		}
	case "memory":
		store = mem_rs.Create()
	default:
		log.Fatalf("Unknown dns record store type %s specified", settings.Store)
	}
	s.RecordSet = record_set.Create(store)
	s.Server = &dns.Server{Addr: addr, Net: "udp"}
	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		s.handleDnsRequest(w, r)
	})
	s.resolver, err = NewResolver(settings)
	for _, v := range settings.Hostnames {
		parts := split_rex.Split(v, -1)
		if len(parts) == 2 {
			s.RecordSet.Put(dns.TypeA, parts[0], parts[1])
		}
	}
	return
}

func (s *Server) handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	defer func() {
		if r := recover(); r != nil {
			handlePanic(r)
		}
	}()
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false
	switch r.Opcode {
	case dns.OpcodeQuery:
		if s.processQuery(m) {
			w.WriteMsg(m)
			return
		}
	}
	if s.settings.DisableForwarding {
		log.Printf("WARNING: name not found and forwarding disabled")
		m.Rcode = dns.RcodeNameError
		w.WriteMsg(m)
		return
	}
	resp, err := s.resolver.Lookup(r)
	if err == nil {
		w.WriteMsg(resp)
	} else {
		log.Printf("Resolver err: %s", err.Error())
	}
}

func (s *Server) processQuery(m *dns.Msg) bool {
	answers := 0

	for _, q := range m.Question {
		var rr dns.RR
		switch q.Qtype {
		case dns.TypeA:
			ip := s.Get(dns.TypeA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s %d A %s", q.Name, s.settings.Ttl, ip))
			}
		case dns.TypePTR:
			host := s.Get(dns.TypePTR, q.Name)
			if host != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s %d PTR %s", q.Name, s.settings.Ttl, host))
			}
		case dns.TypeAAAA:
			ip := s.Get(dns.TypeAAAA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s %d AAAA %s", q.Name, s.settings.Ttl, ip))
			}
		default:
			return false // If we get a question we can't answer, bail
		}
		if rr != nil {
			m.Answer = append(m.Answer, rr)
			if s.settings.Debug {
				log.Printf("Resolved request. RR: %s", rr.String())
			}
			answers++
		}
	}
	if answers > 0 {
		return true
	}
	return false
}
<commit_msg>return NXDOMAIN when an AAAA record is requested for a host with an A record<commit_after>package main

import (
	"fmt"
	"github.com/miekg/dns"
	"gloon/mem_rs"
	"gloon/record_set"
	"gloon/redis_rs"
	"log"
	"regexp"
)

type Server struct {
	*dns.Server
	*record_set.RecordSet
	resolver *Resolver
	settings *Settings
}

var split_rex = regexp.MustCompile("[:= ]")

func NewServer(addr string, settings *Settings) (s *Server, err error) {
	s = &Server{}
	s.settings = settings

	// Set up the record store
	var store record_set.RecordStore
	switch settings.Store {
	case "redis":
		store, err = redis_rs.Create(settings.StoreOpts)
		if err != nil {
			log.Fatalf("Unable to create redis record set: %s", err.Error())
		}
	case "memory":
		store = mem_rs.Create()
	default:
		log.Fatalf("Unknown dns record store type %s specified", settings.Store)
	}
	s.RecordSet = record_set.Create(store)
	s.Server = &dns.Server{Addr: addr, Net: "udp"}
	dns.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		s.handleDnsRequest(w, r)
	})
	s.resolver, err = NewResolver(settings)
	for _, v := range settings.Hostnames {
		parts := split_rex.Split(v, -1)
		if len(parts) == 2 {
			s.RecordSet.Put(dns.TypeA, parts[0], parts[1])
		}
	}
	return
}

func (s *Server) handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	defer func() {
		if r := recover(); r != nil {
			handlePanic(r)
		}
	}()
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false
	switch r.Opcode {
	case dns.OpcodeQuery:
		if s.processQuery(m) {
			w.WriteMsg(m)
			return
		}
	}
	if s.settings.DisableForwarding {
		log.Printf("WARNING: name not found and forwarding disabled")
		m.Rcode = dns.RcodeNameError
		w.WriteMsg(m)
		return
	}
	resp, err := s.resolver.Lookup(r)
	if err == nil {
		w.WriteMsg(resp)
	} else {
		log.Printf("Resolver err: %s", err.Error())
	}
}

func (s *Server) processQuery(m *dns.Msg) bool {
	answers := 0

	for _, q := range m.Question {
		var rr dns.RR
		switch q.Qtype {
		case dns.TypeA:
			ip := s.Get(dns.TypeA, q.Name)
			if ip != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s %d A %s", q.Name, s.settings.Ttl, ip))
			} else if s.settings.Debug {
				log.Printf("Missed A record for %s", q.Name)
			}
		case dns.TypePTR:
			host := s.Get(dns.TypePTR, q.Name)
			if host != "" {
				rr, _ = dns.NewRR(fmt.Sprintf("%s %d PTR %s", q.Name, s.settings.Ttl, host))
			}
		case dns.TypeAAAA: // Bail for now if we have an ipv4
			ip := s.Get(dns.TypeA, q.Name)
			if ip != "" {
				m.Rcode = dns.RcodeNameError
				return true
			}
		default:
			return false // If we get a question we can't answer, bail
		}
		if rr != nil {
			m.Answer = append(m.Answer, rr)
			if s.settings.Debug {
				log.Printf("Resolved request. RR: %s", rr.String())
			}
			answers++
		}
	}
	if answers > 0 {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Pani Networks 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 listener's nodes.go contains a subset of listener functionality
// where it connects to kubernetes using kubernetes client-go and watches
// node creation/deletion events and then adds/deletes the nodes
// to/from romana cluster appropriately.
package listener

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/elgs/gojq"
	log "github.com/romana/rlog"

	romanaApi "github.com/romana/core/common/api"
	romanaErrors "github.com/romana/core/common/api/errors"

	"github.com/romana/core/common/log/trace"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
)

func (l *KubeListener) kubeClientInit() error {
	var err error

	// Try generating config for kubernetes client-go from
	// in-cluster variables like KUBERNETES_SERVICE_HOST, etc
	// so that we can connect to kubernetes using them.
	kConfig, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("error: %s", err)
	}

	// Get a set of REST clients which connect to kubernetes services
	// from the config generated above.
	l.kubeClientSet, err = kubernetes.NewForConfig(kConfig)
	if err != nil {
		return fmt.Errorf("error while connecting to kubernetes: %s", err)
	}

	return nil
}

// nodeToHost converts a K8S node object to Romana host object.
func (l *KubeListener) nodeToHost(n interface{}) (romanaApi.Host, error) {
	var node *v1.Node
	var ok bool
	host := romanaApi.Host{}

	if node, ok = n.(*v1.Node); !ok {
		return host, fmt.Errorf("Expected node object, got %s: %T", n, n)
	}
	if node.Name == "" || node.Status.Addresses == nil || len(node.Status.Addresses) == 0 {
		return host, fmt.Errorf("Received invalid node name or IP Address: (%#v)", node)
	}
	host.Name = node.Name
	hostIP := net.ParseIP(node.Status.Addresses[0].Address)
	if hostIP == nil {
		return host, fmt.Errorf("Cannot parse address of node %s: %s", node.Name, node.Status.Addresses[0].Address)
	}
	host.IP = hostIP
	host.Tags = node.GetLabels()
	if l.nodeAttributes != nil && len(l.nodeAttributes) > 0 {
		host.K8SInfo = make(map[string]interface{})
		json, err := json.Marshal(node)
		if err != nil {
			return host, err
		}
		parser, err := gojq.NewStringQuery(string(json))
		if err != nil {
			return host, err
		}
		for _, nodeAttr := range l.nodeAttributes {
			val, err := parser.Query(nodeAttr)
			if err != nil {
				if strings.HasSuffix(err.Error(), "does not exist.") {
					// log.Tracef(trace.Inside, "Cannot find %s in %s", nodeAttr, string(json))
					continue
				}
				log.Tracef(trace.Inside, "Error querying %s: %s (%T)", nodeAttr, err, err)
				return host, err
			}
			host.K8SInfo[nodeAttr] = val
		}
		if host.K8SInfo != nil && len(host.K8SInfo) > 0 {
			log.Tracef(trace.Inside, "Set k8s_info of host %s", host)
		}
	}
	return host, nil
}

// syncNodes checks what nodes are defined in K8S cluster vs
// hosts defined in Romana and synchronizes them.
// In case syncNodes() is called multiple
// times, it would lock here, check the status of syncNodesRunning flag
// and bail out if there is another instance of syncNodes() running.
// Note that that can only happen when the syncNodes() routine was taking
// so long that the next scheduled one has started, which is unlikely,
// but it doesn't hurt to ensure that they don't step on each other.
// Note also that individual handlers (on node add/delete/update) will not
// synchronize -- in the worst case, they will fail and be picked up by next
// synchronization.
func (l *KubeListener) syncNodes() {
	log.Trace(trace.Inside, "Entering syncNodes()\n")

	var err error
	l.syncNodesMutex.Lock()
	if l.syncNodesRunning {
		// This really can only happen if syncNodes was taking too long
		// and the next scheduled syncNodes started. Very unlikely, but
		// doesn'thurt to check.
		log.Infof("syncNodes() already running, will exit.")
		l.syncNodesMutex.Unlock()
		return
	}
	l.syncNodesRunning = true
	l.syncNodesMutex.Unlock()

	k8sNodesList := l.nodeStore.List()

	romanaHostList := l.client.IPAM.ListHosts()
	log.Debugf("Comparing Romana host list %d vs K8S node list %d", len(k8sNodesList), len(romanaHostList.Hosts))

	var nodeInRomana bool
	var romanaHost romanaApi.Host
	var hostToAdd romanaApi.Host

	for _, n := range k8sNodesList {
		if hostToAdd, err = l.nodeToHost(n); err != nil {
			log.Error(err)
			continue
		}

		log.Tracef(trace.Inside, "Checking if node %s is in Romana", hostToAdd)

		nodeInRomana = false
		// This is not super-efficient but as we don't have that many hosts for now
		// to deal with, it can do.
		for _, romanaHost = range romanaHostList.Hosts {
			if romanaHost.IP.String() == hostToAdd.IP.String() {
				nodeInRomana = true
				break
			}
		}
		if !nodeInRomana {
			log.Tracef(trace.Inside, "Trying to add host %s to Romana", hostToAdd)
			if err = l.romanaHostAdd(hostToAdd); err != nil {
				log.Error(err)
			}
		}
	}

	var hostToRemove romanaApi.Host
	var hostInK8S bool

	var node *v1.Node
	for _, romanaHost = range romanaHostList.Hosts {
		hostInK8S = false
		for _, n := range k8sNodesList {
			node = n.(*v1.Node)
			if hostToRemove, err = l.nodeToHost(node); err != nil {
				log.Error(err)
				continue
			}
			if hostToRemove.IP.String() == romanaHost.IP.String() {
				hostInK8S = true
				break
			}
		}
		log.Tracef(trace.Inside, "Checking if host %s is in K8S: %t", romanaHost, hostInK8S)

		if !hostInK8S {
			if err = l.romanaHostRemove(hostToRemove); err != nil {
				log.Error(err)
			}
		}
	}
	l.syncNodesMutex.Lock()
	l.syncNodesRunning = false
	l.syncNodesMutex.Unlock()
}

// ProcessNodeEvents processes kubernetes node events, there by
// adding/deleting nodes to/from romana cluster automatically
// when they are added/removed to/from kubernetes cluster.
func (l *KubeListener) ProcessNodeEvents(done <-chan struct{}) {
	log.Debug("In ProcessNodeEvents()")

	// nodeWatcher is a new ListWatch object created from the specified
	// kubeClientSet which k8s.io/client-go exports for watching node events.
	nodeWatcher := cache.NewListWatchFromClient(
		l.kubeClientSet.CoreV1Client.RESTClient(),
		"nodes",
		api.NamespaceAll,
		fields.Everything())

	// Setup a notifications for specific events using NewInformer.
	l.nodeStore, l.nodeInformer = cache.NewInformer(
		nodeWatcher,
		&v1.Node{},
		// TODO this can also be configurable
		time.Minute,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    l.kubernetesAddNodeEventHandler,
			UpdateFunc: l.kubernetesUpdateNodeEventHandler,
			DeleteFunc: l.kubernetesDeleteNodeEventHandler,
		},
	)

	log.Infof("Starting receving node events.")
	go l.nodeInformer.Run(done)

	log.Infof("Starting sync ticker with %s", l.syncNodesInterval)
	l.syncNodesTicker = time.NewTicker(l.syncNodesInterval)

	log.Infof("Starting initial sync ticker with %s", initialSyncInterval)
	initialSyncTicker := time.NewTicker(initialSyncInterval)
	log.Infof("Setting timeout for initial sync ticker to %s from now", initialSyncDuration)
	initialSyncTimeout := time.After(initialSyncDuration)

	// The label to break out of the loop must be placed before the loop, per
	// https://golang.org/ref/spec#Break_statements
	// "If there is a label, it must be that of an enclosing "for", "switch", or
	// "select" statement, and that is the one whose execution terminates.
INITIAL_SYNC:
	for {
		select {
		case <-initialSyncTicker.C:
			if l.nodeInformer.HasSynced() {
				log.Info("Initial synchronization completed\n")
				l.initialNodesSyncDone = true
				initialSyncTicker.Stop()
				l.syncNodes()
				// This actually causes the loop to exit.
				break INITIAL_SYNC
			}
		case <-initialSyncTimeout:
			log.Errorf("Timeout after %s while synchronizing nodes\n", initialSyncDuration)
			os.Exit(1)

		}
	}

	go func() {
		for {
			select {
			case t := <-l.syncNodesTicker.C:
				if !l.initialNodesSyncDone {
					continue
				}
				log.Tracef(trace.Inside, "Entering timed syncNodes() at %s\n", t)
				l.syncNodes()
			}
		}
	}()

}

// kubernetesAddNodeEventHandler is called when Kubernetes reports an
// add node event.
func (l *KubeListener) kubernetesAddNodeEventHandler(n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring add event")
		return
	}
	if hostToAdd, err := l.nodeToHost(n); err != nil {
		log.Errorf("Error handling node add event: %s", err)
	} else if err = l.romanaHostAdd(hostToAdd); err != nil {
		log.Errorf("Error handling node add event: %s", err)
	}
}

// kubernetesDeleteNodeEventHandler is called when Kubernetes reports a
// delete node event.
func (l *KubeListener) kubernetesDeleteNodeEventHandler(n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring delete event")
		return
	}
	if hostToRemove, err := l.nodeToHost(n); err != nil {
		log.Errorf("Error handling node remove event: %s", err)
	} else if err = l.romanaHostRemove(hostToRemove); err != nil {
		log.Errorf("Error handling node remove event: %s", err)
	}
}

// kubernetesUpdateNodeEventHandler is called when Kubernetes reports an
// update node event. It calls syncNodes to sync romana/kubernetes
// host list.
func (l *KubeListener) kubernetesUpdateNodeEventHandler(o, n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring update	 event")
		return
	}
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Expected Node object, received (%T: %s)", n, n)
		return
	}

	host, err := l.nodeToHost(node)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
		return
	}

	err = l.client.IPAM.UpdateHostLabels(host)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
	}
	err = l.client.IPAM.UpdateHostK8SInfo(host)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
	}
}

// romanaHostAdd connects to romana API and adds a node to
// the romana cluster.
func (l *KubeListener) romanaHostAdd(host romanaApi.Host) error {
	var ok bool
	err := l.client.IPAM.AddHost(host)
	if _, ok = err.(romanaErrors.RomanaExistsError); ok {
		log.Infof("Host %s already exists, ignoring addition.", host)
		return nil
	} else if err == nil {
		log.Infof("Host (%s) successfully added to Romana cluster.", host)
		return nil
	}
	return err
}

// romanaHostRemove connects to romana API and removes a node from
// the romana cluster.
func (l *KubeListener) romanaHostRemove(host romanaApi.Host) error {
	err := l.client.IPAM.RemoveHost(host)
	if _, ok := err.(romanaErrors.RomanaNotFoundError); ok {
		log.Infof("Host %s is not found, ignoring removal", host)
		return nil
	} else if err == nil {
		log.Infof("Host %s successfully removed from Romana cluster", host)
		return nil
	}
	return err
}
<commit_msg>Ensure the correct node is deleted. Fixes variable names and scope<commit_after>// Copyright (c) 2016-2017 Pani Networks 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 listener's nodes.go contains a subset of listener functionality
// where it connects to kubernetes using kubernetes client-go and watches
// node creation/deletion events and then adds/deletes the nodes
// to/from romana cluster appropriately.
package listener

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/elgs/gojq"
	log "github.com/romana/rlog"

	romanaApi "github.com/romana/core/common/api"
	romanaErrors "github.com/romana/core/common/api/errors"

	"github.com/romana/core/common/log/trace"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/fields"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
)

func (l *KubeListener) kubeClientInit() error {
	var err error

	// Try generating config for kubernetes client-go from
	// in-cluster variables like KUBERNETES_SERVICE_HOST, etc
	// so that we can connect to kubernetes using them.
	kConfig, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("error: %s", err)
	}

	// Get a set of REST clients which connect to kubernetes services
	// from the config generated above.
	l.kubeClientSet, err = kubernetes.NewForConfig(kConfig)
	if err != nil {
		return fmt.Errorf("error while connecting to kubernetes: %s", err)
	}

	return nil
}

// nodeToHost converts a K8S node object to Romana host object.
func (l *KubeListener) nodeToHost(n interface{}) (romanaApi.Host, error) {
	var node *v1.Node
	var ok bool
	host := romanaApi.Host{}

	if node, ok = n.(*v1.Node); !ok {
		return host, fmt.Errorf("Expected node object, got %s: %T", n, n)
	}
	if node.Name == "" || node.Status.Addresses == nil || len(node.Status.Addresses) == 0 {
		return host, fmt.Errorf("Received invalid node name or IP Address: (%#v)", node)
	}
	host.Name = node.Name
	hostIP := net.ParseIP(node.Status.Addresses[0].Address)
	if hostIP == nil {
		return host, fmt.Errorf("Cannot parse address of node %s: %s", node.Name, node.Status.Addresses[0].Address)
	}
	host.IP = hostIP
	host.Tags = node.GetLabels()
	if l.nodeAttributes != nil && len(l.nodeAttributes) > 0 {
		host.K8SInfo = make(map[string]interface{})
		json, err := json.Marshal(node)
		if err != nil {
			return host, err
		}
		parser, err := gojq.NewStringQuery(string(json))
		if err != nil {
			return host, err
		}
		for _, nodeAttr := range l.nodeAttributes {
			val, err := parser.Query(nodeAttr)
			if err != nil {
				if strings.HasSuffix(err.Error(), "does not exist.") {
					// log.Tracef(trace.Inside, "Cannot find %s in %s", nodeAttr, string(json))
					continue
				}
				log.Tracef(trace.Inside, "Error querying %s: %s (%T)", nodeAttr, err, err)
				return host, err
			}
			host.K8SInfo[nodeAttr] = val
		}
		if host.K8SInfo != nil && len(host.K8SInfo) > 0 {
			log.Tracef(trace.Inside, "Set k8s_info of host %s", host)
		}
	}
	return host, nil
}

// syncNodes checks what nodes are defined in K8S cluster vs
// hosts defined in Romana and synchronizes them.
// In case syncNodes() is called multiple
// times, it would lock here, check the status of syncNodesRunning flag
// and bail out if there is another instance of syncNodes() running.
// Note that that can only happen when the syncNodes() routine was taking
// so long that the next scheduled one has started, which is unlikely,
// but it doesn't hurt to ensure that they don't step on each other.
// Note also that individual handlers (on node add/delete/update) will not
// synchronize -- in the worst case, they will fail and be picked up by next
// synchronization.
func (l *KubeListener) syncNodes() {
	log.Trace(trace.Inside, "Entering syncNodes()\n")

	var err error
	l.syncNodesMutex.Lock()
	if l.syncNodesRunning {
		// This really can only happen if syncNodes was taking too long
		// and the next scheduled syncNodes started. Very unlikely, but
		// doesn'thurt to check.
		log.Infof("syncNodes() already running, will exit.")
		l.syncNodesMutex.Unlock()
		return
	}
	l.syncNodesRunning = true
	l.syncNodesMutex.Unlock()

	k8sNodesList := l.nodeStore.List()
	romanaHostList := l.client.IPAM.ListHosts()

	log.Debugf("Comparing Romana host list %d vs K8S node list %d", len(k8sNodesList), len(romanaHostList.Hosts))

	// Check for nodes that exist in kubernetes but not registered as romana hosts.
	// Add hosts that are missing
	for _, n := range k8sNodesList {
		host, err := l.nodeToHost(n)
		if err != nil {
			log.Error(err)
			continue
		}

		log.Tracef(trace.Inside, "Checking if node %s is in Romana", host)

		nodeInRomana := false
		// This is not super-efficient but as we don't have that many hosts for now
		// to deal with, it can do.
		for _, romanaHost := range romanaHostList.Hosts {
			if romanaHost.IP.String() == host.IP.String() {
				nodeInRomana = true
				break
			}
		}
		if !nodeInRomana {
			log.Tracef(trace.Inside, "Trying to add host %s to Romana", host)
			if err = l.romanaHostAdd(host); err != nil {
				log.Error(err)
			}
		}
	}

	// Check for hosts that are registered with romana but don't exist as kubernetes nodes.
	// Remove hosts that are missing in kubernetes
	for _, romanaHost := range romanaHostList.Hosts {
		hostInK8S := false
		for _, n := range k8sNodesList {
			node := n.(*v1.Node)
			host, err := l.nodeToHost(node)
			if err != nil {
				log.Error(err)
				continue
			}
			if host.IP.String() == romanaHost.IP.String() {
				hostInK8S = true
				break
			}
		}
		log.Tracef(trace.Inside, "Checking if host %s is in K8S: %t", romanaHost, hostInK8S)

		if !hostInK8S {
			if err = l.romanaHostRemove(romanaHost); err != nil {
				log.Error(err)
			}
		}
	}
	l.syncNodesMutex.Lock()
	l.syncNodesRunning = false
	l.syncNodesMutex.Unlock()
}

// ProcessNodeEvents processes kubernetes node events, there by
// adding/deleting nodes to/from romana cluster automatically
// when they are added/removed to/from kubernetes cluster.
func (l *KubeListener) ProcessNodeEvents(done <-chan struct{}) {
	log.Debug("In ProcessNodeEvents()")

	// nodeWatcher is a new ListWatch object created from the specified
	// kubeClientSet which k8s.io/client-go exports for watching node events.
	nodeWatcher := cache.NewListWatchFromClient(
		l.kubeClientSet.CoreV1Client.RESTClient(),
		"nodes",
		api.NamespaceAll,
		fields.Everything())

	// Setup a notifications for specific events using NewInformer.
	l.nodeStore, l.nodeInformer = cache.NewInformer(
		nodeWatcher,
		&v1.Node{},
		// TODO this can also be configurable
		time.Minute,
		cache.ResourceEventHandlerFuncs{
			AddFunc:    l.kubernetesAddNodeEventHandler,
			UpdateFunc: l.kubernetesUpdateNodeEventHandler,
			DeleteFunc: l.kubernetesDeleteNodeEventHandler,
		},
	)

	log.Infof("Starting receving node events.")
	go l.nodeInformer.Run(done)

	log.Infof("Starting sync ticker with %s", l.syncNodesInterval)
	l.syncNodesTicker = time.NewTicker(l.syncNodesInterval)

	log.Infof("Starting initial sync ticker with %s", initialSyncInterval)
	initialSyncTicker := time.NewTicker(initialSyncInterval)
	log.Infof("Setting timeout for initial sync ticker to %s from now", initialSyncDuration)
	initialSyncTimeout := time.After(initialSyncDuration)

	// The label to break out of the loop must be placed before the loop, per
	// https://golang.org/ref/spec#Break_statements
	// "If there is a label, it must be that of an enclosing "for", "switch", or
	// "select" statement, and that is the one whose execution terminates.
INITIAL_SYNC:
	for {
		select {
		case <-initialSyncTicker.C:
			if l.nodeInformer.HasSynced() {
				log.Info("Initial synchronization completed\n")
				l.initialNodesSyncDone = true
				initialSyncTicker.Stop()
				l.syncNodes()
				// This actually causes the loop to exit.
				break INITIAL_SYNC
			}
		case <-initialSyncTimeout:
			log.Errorf("Timeout after %s while synchronizing nodes\n", initialSyncDuration)
			os.Exit(1)

		}
	}

	go func() {
		for {
			select {
			case t := <-l.syncNodesTicker.C:
				if !l.initialNodesSyncDone {
					continue
				}
				log.Tracef(trace.Inside, "Entering timed syncNodes() at %s\n", t)
				l.syncNodes()
			}
		}
	}()

}

// kubernetesAddNodeEventHandler is called when Kubernetes reports an
// add node event.
func (l *KubeListener) kubernetesAddNodeEventHandler(n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring add event")
		return
	}
	if hostToAdd, err := l.nodeToHost(n); err != nil {
		log.Errorf("Error handling node add event: %s", err)
	} else if err = l.romanaHostAdd(hostToAdd); err != nil {
		log.Errorf("Error handling node add event: %s", err)
	}
}

// kubernetesDeleteNodeEventHandler is called when Kubernetes reports a
// delete node event.
func (l *KubeListener) kubernetesDeleteNodeEventHandler(n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring delete event")
		return
	}
	if hostToRemove, err := l.nodeToHost(n); err != nil {
		log.Errorf("Error handling node remove event: %s", err)
	} else if err = l.romanaHostRemove(hostToRemove); err != nil {
		log.Errorf("Error handling node remove event: %s", err)
	}
}

// kubernetesUpdateNodeEventHandler is called when Kubernetes reports an
// update node event. It calls syncNodes to sync romana/kubernetes
// host list.
func (l *KubeListener) kubernetesUpdateNodeEventHandler(o, n interface{}) {
	if !l.initialNodesSyncDone {
		log.Debug("Initial synchronization not completed, ignoring update	 event")
		return
	}
	node, ok := n.(*v1.Node)
	if !ok {
		log.Errorf("Expected Node object, received (%T: %s)", n, n)
		return
	}

	host, err := l.nodeToHost(node)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
		return
	}

	err = l.client.IPAM.UpdateHostLabels(host)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
	}
	err = l.client.IPAM.UpdateHostK8SInfo(host)
	if err != nil {
		log.Errorf("Cannot update node %s: %s", node.Name, err)
	}
}

// romanaHostAdd connects to romana API and adds a node to
// the romana cluster.
func (l *KubeListener) romanaHostAdd(host romanaApi.Host) error {
	var ok bool
	err := l.client.IPAM.AddHost(host)
	if _, ok = err.(romanaErrors.RomanaExistsError); ok {
		log.Infof("Host %s already exists, ignoring addition.", host)
		return nil
	} else if err == nil {
		log.Infof("Host (%s) successfully added to Romana cluster.", host)
		return nil
	}
	return err
}

// romanaHostRemove connects to romana API and removes a node from
// the romana cluster.
func (l *KubeListener) romanaHostRemove(host romanaApi.Host) error {
	err := l.client.IPAM.RemoveHost(host)
	if _, ok := err.(romanaErrors.RomanaNotFoundError); ok {
		log.Infof("Host %s is not found, ignoring removal", host)
		return nil
	} else if err == nil {
		log.Infof("Host %s successfully removed from Romana cluster", host)
		return nil
	}
	return err
}
<|endoftext|>


<commit_before>// Package mailjet provides methods for interacting with the last version of the Mailjet API.
// The goal of this component is to simplify the usage of the MailJet API for GO developers.
//
// For more details, see the full API Documentation at http://dev.mailjet.com/
package mailjet

import (
	"bytes"
	"fmt"

	"encoding/json"

	"io"
	"log"
	"net/http"
	"net/textproto"
	"os"
	"strings"
)

// NewMailjetClient returns a new MailjetClient using an public apikey
// and an secret apikey to be used when authenticating to API.
func NewMailjetClient(apiKeyPublic, apiKeyPrivate string, baseURL ...string) *Client {
	httpClient := NewHTTPClient(apiKeyPublic, apiKeyPrivate)
	smtpClient := NewSMTPClient(apiKeyPublic, apiKeyPrivate)

	client := &Client{
		httpClient: httpClient,
		smtpClient: smtpClient,
		apiBase:    apiBase,
	}

	if len(baseURL) > 0 {
		client.apiBase = baseURL[0]
	}
	return client
}

// NewClient function
func NewClient(httpCl HTTPClientInterface, smtpCl SMTPClientInterface, baseURL ...string) *Client {

	client := &Client{
		httpClient: httpCl,
		smtpClient: smtpCl,
		apiBase:    apiBase,
	}

	if len(baseURL) > 0 {
		client.apiBase = baseURL[0]
	}
	return client
}

// SetBaseURL sets the base URL
func (c *Client) SetBaseURL(baseURL string) {
	c.apiBase = baseURL
}

// APIKeyPublic returns the public key.
func (c *Client) APIKeyPublic() string {
	return c.httpClient.APIKeyPublic()
}

// APIKeyPrivate returns the secret key.
func (c *Client) APIKeyPrivate() string {
	return c.httpClient.APIKeyPrivate()

}

// SetURL function to set the base url of the wrapper instance
func (c *Client) SetURL(baseURL string) {
	c.Lock()
	defer c.Unlock()
	c.apiBase = baseURL
}

// Client returns the underlying http client
func (c *Client) Client() *http.Client {
	return c.httpClient.Client()
}

// SetClient allows to customize http client.
func (c *Client) SetClient(client *http.Client) {
	c.Lock()
	defer c.Unlock()
	c.httpClient.SetClient(client)
}

// Filter applies a filter with the defined key and value.
func Filter(key, value string) RequestOptions {
	return func(req *http.Request) {
		q := req.URL.Query()
		q.Add(key, value)
		req.URL.RawQuery = strings.Replace(q.Encode(), "%2B", "+", 1)
	}
}

// SortOrder defines the order of the result.
type SortOrder int

// These are the two possible order.
const (
	SortDesc = SortOrder(iota)
	SortAsc
)

var debugOut io.Writer = os.Stderr

// SetDebugOutput sets the output destination for the debug.
func SetDebugOutput(w io.Writer) {
	debugOut = w
	log.SetOutput(w)
}

// Sort applies the Sort filter to the request.
func Sort(value string, order SortOrder) RequestOptions {
	if order == SortDesc {
		value = value + "+DESC"
	}
	return Filter("Sort", value)
}

// List issues a GET to list the specified resource
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
func (c *Client) List(resource string, resp interface{}, options ...RequestOptions) (count, total int, err error) {
	url := buildURL(c.apiBase, &Request{Resource: resource})
	req, err := createRequest("GET", url, nil, nil, options...)
	if err != nil {
		return count, total, err
	}

	c.Lock()
	defer c.Unlock()
	return c.httpClient.Send(req).Read(resp).Call()
}

// Get issues a GET to view a resource specifying an id
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
// Without an specified ID in MailjetRequest, it is the same as List.
func (c *Client) Get(mr *Request, resp interface{}, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, mr)
	req, err := createRequest("GET", url, nil, nil, options...)
	if err != nil {
		return err
	}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).Read(resp).Call()
	return err
}

// Post issues a POST to create a new resource
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
func (c *Client) Post(fmr *FullRequest, resp interface{}, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, fmr.Info)
	req, err := createRequest("POST", url, fmr.Payload, nil, options...)
	if err != nil {
		return err
	}

	headers := map[string]string{"Content-Type": "application/json"}
	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Read(resp).Call()
	return err
}

// Put is used to update a resource.
// Fields to be updated must be specified by the string array onlyFields.
// If onlyFields is nil, all fields except these with the tag read_only, are updated.
// Filters can be add via functional options.
func (c *Client) Put(fmr *FullRequest, onlyFields []string, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, fmr.Info)
	req, err := createRequest("PUT", url, fmr.Payload, onlyFields, options...)
	if err != nil {
		return err
	}

	headers := map[string]string{"Content-Type": "application/json"}
	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Call()
	return err
}

// Delete is used to delete a resource.
func (c *Client) Delete(mr *Request) (err error) {
	url := buildURL(c.apiBase, mr)
	req, err := createRequest("DELETE", url, nil, nil)
	if err != nil {
		return err
	}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).Call()
	return err
}

// SendMail send mail via API.
func (c *Client) SendMail(data *InfoSendMail) (res *SentResult, err error) {
	url := c.apiBase + "/send/message"
	req, err := createRequest("POST", url, data, nil)
	if err != nil {
		return res, err
	}

	headers := map[string]string{"Content-Type": "application/json"}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Read(&res).Call()
	return res, err
}

// SendMailSMTP send mail via SMTP.
func (c *Client) SendMailSMTP(info *InfoSMTP) error {
	return c.smtpClient.SendMail(
		info.From,
		info.Recipients,
		buildMessage(info.Header, info.Content))
}

func buildMessage(header textproto.MIMEHeader, content []byte) []byte {
	buff := bytes.NewBuffer(nil)
	for key, values := range header {
		buff.WriteString(fmt.Sprintf("%s: %s\r\n", key, strings.Join(values, ", ")))
	}
	buff.WriteString("\r\n")
	buff.Write(content)

	return buff.Bytes()
}

// SendMailV31 sends a mail to the send API v3.1
func (c *Client) SendMailV31(data *MessagesV31) (*ResultsV31, error) {
	url := c.apiBase + ".1/send"
	req, err := createRequest("POST", url, data, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth(c.APIKeyPublic(), c.APIKeyPrivate())

	r, err := c.httpClient.SendMailV31(req)
	if err != nil {
		return nil, err
	}

	decoder := json.NewDecoder(r.Body)

	switch r.StatusCode {
	case http.StatusOK:

		var res ResultsV31
		if err := decoder.Decode(&res); err != nil {
			return nil, err
		}
		return &res, nil

	case http.StatusBadRequest, http.StatusForbidden:

		var apiFeedbackErr APIFeedbackErrorsV31
		if err := decoder.Decode(&apiFeedbackErr); err != nil {
			return nil, err
		}
		return nil, &apiFeedbackErr

	default:

		var errInfo ErrorInfoV31
		if err := decoder.Decode(&errInfo); err != nil {
			return nil, err
		}
		return nil, &errInfo
	}
}
<commit_msg>add request options to send mail methods<commit_after>// Package mailjet provides methods for interacting with the last version of the Mailjet API.
// The goal of this component is to simplify the usage of the MailJet API for GO developers.
//
// For more details, see the full API Documentation at http://dev.mailjet.com/
package mailjet

import (
	"bytes"
	"context"
	"fmt"

	"encoding/json"

	"io"
	"log"
	"net/http"
	"net/textproto"
	"os"
	"strings"
)

// NewMailjetClient returns a new MailjetClient using an public apikey
// and an secret apikey to be used when authenticating to API.
func NewMailjetClient(apiKeyPublic, apiKeyPrivate string, baseURL ...string) *Client {
	httpClient := NewHTTPClient(apiKeyPublic, apiKeyPrivate)
	smtpClient := NewSMTPClient(apiKeyPublic, apiKeyPrivate)

	client := &Client{
		httpClient: httpClient,
		smtpClient: smtpClient,
		apiBase:    apiBase,
	}

	if len(baseURL) > 0 {
		client.apiBase = baseURL[0]
	}
	return client
}

// NewClient function
func NewClient(httpCl HTTPClientInterface, smtpCl SMTPClientInterface, baseURL ...string) *Client {

	client := &Client{
		httpClient: httpCl,
		smtpClient: smtpCl,
		apiBase:    apiBase,
	}

	if len(baseURL) > 0 {
		client.apiBase = baseURL[0]
	}
	return client
}

// SetBaseURL sets the base URL
func (c *Client) SetBaseURL(baseURL string) {
	c.apiBase = baseURL
}

// APIKeyPublic returns the public key.
func (c *Client) APIKeyPublic() string {
	return c.httpClient.APIKeyPublic()
}

// APIKeyPrivate returns the secret key.
func (c *Client) APIKeyPrivate() string {
	return c.httpClient.APIKeyPrivate()

}

// SetURL function to set the base url of the wrapper instance
func (c *Client) SetURL(baseURL string) {
	c.Lock()
	defer c.Unlock()
	c.apiBase = baseURL
}

// Client returns the underlying http client
func (c *Client) Client() *http.Client {
	return c.httpClient.Client()
}

// SetClient allows to customize http client.
func (c *Client) SetClient(client *http.Client) {
	c.Lock()
	defer c.Unlock()
	c.httpClient.SetClient(client)
}

// Filter applies a filter with the defined key and value.
func Filter(key, value string) RequestOptions {
	return func(req *http.Request) {
		q := req.URL.Query()
		q.Add(key, value)
		req.URL.RawQuery = strings.Replace(q.Encode(), "%2B", "+", 1)
	}
}

// SortOrder defines the order of the result.
type SortOrder int

// These are the two possible order.
const (
	SortDesc = SortOrder(iota)
	SortAsc
)

var debugOut io.Writer = os.Stderr

// SetDebugOutput sets the output destination for the debug.
func SetDebugOutput(w io.Writer) {
	debugOut = w
	log.SetOutput(w)
}

// Sort applies the Sort filter to the request.
func Sort(value string, order SortOrder) RequestOptions {
	if order == SortDesc {
		value = value + "+DESC"
	}
	return Filter("Sort", value)
}

// List issues a GET to list the specified resource
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
func (c *Client) List(resource string, resp interface{}, options ...RequestOptions) (count, total int, err error) {
	url := buildURL(c.apiBase, &Request{Resource: resource})
	req, err := createRequest("GET", url, nil, nil, options...)
	if err != nil {
		return count, total, err
	}

	c.Lock()
	defer c.Unlock()
	return c.httpClient.Send(req).Read(resp).Call()
}

// Get issues a GET to view a resource specifying an id
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
// Without an specified ID in MailjetRequest, it is the same as List.
func (c *Client) Get(mr *Request, resp interface{}, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, mr)
	req, err := createRequest("GET", url, nil, nil, options...)
	if err != nil {
		return err
	}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).Read(resp).Call()
	return err
}

// Post issues a POST to create a new resource
// and stores the result in the value pointed to by res.
// Filters can be add via functional options.
func (c *Client) Post(fmr *FullRequest, resp interface{}, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, fmr.Info)
	req, err := createRequest("POST", url, fmr.Payload, nil, options...)
	if err != nil {
		return err
	}

	headers := map[string]string{"Content-Type": "application/json"}
	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Read(resp).Call()
	return err
}

// Put is used to update a resource.
// Fields to be updated must be specified by the string array onlyFields.
// If onlyFields is nil, all fields except these with the tag read_only, are updated.
// Filters can be add via functional options.
func (c *Client) Put(fmr *FullRequest, onlyFields []string, options ...RequestOptions) (err error) {
	url := buildURL(c.apiBase, fmr.Info)
	req, err := createRequest("PUT", url, fmr.Payload, onlyFields, options...)
	if err != nil {
		return err
	}

	headers := map[string]string{"Content-Type": "application/json"}
	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Call()
	return err
}

// Delete is used to delete a resource.
func (c *Client) Delete(mr *Request) (err error) {
	url := buildURL(c.apiBase, mr)
	req, err := createRequest("DELETE", url, nil, nil)
	if err != nil {
		return err
	}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).Call()
	return err
}

// SendMail send mail via API.
func (c *Client) SendMail(data *InfoSendMail, options ...RequestOptions) (res *SentResult, err error) {
	url := c.apiBase + "/send/message"
	req, err := createRequest("POST", url, data, nil, options...)
	if err != nil {
		return res, err
	}

	headers := map[string]string{"Content-Type": "application/json"}

	c.Lock()
	defer c.Unlock()
	_, _, err = c.httpClient.Send(req).With(headers).Read(&res).Call()
	return res, err
}

// SendMailSMTP send mail via SMTP.
func (c *Client) SendMailSMTP(info *InfoSMTP) error {
	return c.smtpClient.SendMail(
		info.From,
		info.Recipients,
		buildMessage(info.Header, info.Content))
}

func buildMessage(header textproto.MIMEHeader, content []byte) []byte {
	buff := bytes.NewBuffer(nil)
	for key, values := range header {
		buff.WriteString(fmt.Sprintf("%s: %s\r\n", key, strings.Join(values, ", ")))
	}
	buff.WriteString("\r\n")
	buff.Write(content)

	return buff.Bytes()
}

// SendMailV31 sends a mail to the send API v3.1
func (c *Client) SendMailV31(data *MessagesV31, options ...RequestOptions) (*ResultsV31, error) {
	url := c.apiBase + ".1/send"
	req, err := createRequest("POST", url, data, nil, options...)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth(c.APIKeyPublic(), c.APIKeyPrivate())

	r, err := c.httpClient.SendMailV31(req)
	if err != nil {
		return nil, err
	}

	decoder := json.NewDecoder(r.Body)

	switch r.StatusCode {
	case http.StatusOK:

		var res ResultsV31
		if err := decoder.Decode(&res); err != nil {
			return nil, err
		}
		return &res, nil

	case http.StatusBadRequest, http.StatusForbidden:

		var apiFeedbackErr APIFeedbackErrorsV31
		if err := decoder.Decode(&apiFeedbackErr); err != nil {
			return nil, err
		}
		return nil, &apiFeedbackErr

	default:

		var errInfo ErrorInfoV31
		if err := decoder.Decode(&errInfo); err != nil {
			return nil, err
		}
		return nil, &errInfo
	}
}
<|endoftext|>


<commit_before>package makeindex

import (
	//"github.com/indraniel/srasearch/sra"
	"github.com/indraniel/srasearch/utils"
	"bufio"
	//"encoding/json"
	"github.com/blevesearch/bleve"
	"io"
	"log"
	"os"
	"strings"
	"time"
)

func CreateSearchIndex(input, output string) {
	if _, err := os.Stat(output); err == nil {
		log.Fatalf("Index: '%s' %s %s",
			output,
			"already exists!",
			"Please move or delete it and try again.",
		)
	}

	log.Print("Starting to make a brand new index")

	batchSize := 100

	f, gzf := utils.OpenGZFile(input)
	defer utils.CloseGZFile(f, gzf)

	mapping := bleve.NewIndexMapping()
	index, err := bleve.New(output, mapping)
	if err != nil {
		log.Fatal("Trouble making a bleve index!")
	}

	reader := bufio.NewReader(gzf)
	line, err := reader.ReadString('\n')

	count := 0
	batch := index.NewBatch()
	startTime := time.Now()
	batchCount := 0
	for err == nil {
		elems := strings.SplitN(line, ",", 2)
		docId, json := elems[0], elems[1]
		json = strings.TrimRight(json, "\n")

		batch.Index(docId, []byte(json))
		batchCount++

		if batchCount >= batchSize {
			//log.Println("Submitting a batch to the index")
			err := index.Batch(batch)
			//log.Println("[DONE] Submitting a batch to the index")
			if err != nil {
				log.Fatal("Couldn't batch!")
				//				return err
			}
			batch = index.NewBatch()
			batchCount = 0
		}
		count++
		if count%1000 == 0 {
			indexDuration := time.Since(startTime)
			indexDurationSeconds := float64(indexDuration) / float64(time.Second)
			timePerDoc := float64(indexDuration) / float64(count)
			log.Printf("Indexed %d documents, in %.2fs (average %.2fms/doc)\n", count, indexDurationSeconds, timePerDoc/float64(time.Millisecond))
		}
		line, err = reader.ReadString('\n')
	}

	if err != io.EOF {
		log.Fatalln("[err] reading line ", count, "in", f.Name(), ":", err)
	}

	// flush the last batch
	if batchCount > 0 {
		err = index.Batch(batch)
		if err != nil {
			log.Fatal(err)
		}
	}

	indexDuration := time.Since(startTime)
	indexDurationSeconds := float64(indexDuration) / float64(time.Second)
	timePerDoc := float64(indexDuration) / float64(count)
	log.Printf("Indexed %d documents, in %.2fs (average %.2fms/doc)\n", count, indexDurationSeconds, timePerDoc/float64(time.Millisecond))
	log.Println("All Done")
}

//func buildIndexMapping() (*bleve.IndexMapping, error) {
//}

// v1
//		var si sra.SraItem
//		err := json.Unmarshal([]byte(line), &si)
//		if err != nil {
//			log.Fatal("Trouble unmarshal JSON to SraItem:", err)
//		}
//		//		log.Println("(", i, ")", "Indexing", si.Id)
//
//		err = index.Index(si.Id, si)
//		if err != nil {
//			log.Fatal("Trouble indexing sra item:", err)
//		}
//		i++
//		if i%25000 == 0 {
//			log.Println("Indexed", i, "entries")
//		}
<commit_msg>+ show a more explicit err message<commit_after>package makeindex

import (
	//"github.com/indraniel/srasearch/sra"
	"github.com/indraniel/srasearch/utils"
	"bufio"
	//"encoding/json"
	"github.com/blevesearch/bleve"
	"io"
	"log"
	"os"
	"strings"
	"time"
)

func CreateSearchIndex(input, output string) {
	if _, err := os.Stat(output); err == nil {
		log.Fatalf("Index: '%s' %s %s",
			output,
			"already exists!",
			"Please move or delete it and try again.",
		)
	}

	log.Print("Starting to make a brand new index")

	batchSize := 100

	f, gzf := utils.OpenGZFile(input)
	defer utils.CloseGZFile(f, gzf)

	mapping := bleve.NewIndexMapping()
	index, err := bleve.New(output, mapping)
	if err != nil {
		log.Fatalln("Trouble making a bleve index! :", err)
	}

	reader := bufio.NewReader(gzf)
	line, err := reader.ReadString('\n')

	count := 0
	batch := index.NewBatch()
	startTime := time.Now()
	batchCount := 0
	for err == nil {
		elems := strings.SplitN(line, ",", 2)
		docId, json := elems[0], elems[1]
		json = strings.TrimRight(json, "\n")

		batch.Index(docId, []byte(json))
		batchCount++

		if batchCount >= batchSize {
			//log.Println("Submitting a batch to the index")
			err := index.Batch(batch)
			//log.Println("[DONE] Submitting a batch to the index")
			if err != nil {
				log.Fatal("Couldn't batch!")
				//				return err
			}
			batch = index.NewBatch()
			batchCount = 0
		}
		count++
		if count%1000 == 0 {
			indexDuration := time.Since(startTime)
			indexDurationSeconds := float64(indexDuration) / float64(time.Second)
			timePerDoc := float64(indexDuration) / float64(count)
			log.Printf("Indexed %d documents, in %.2fs (average %.2fms/doc)\n", count, indexDurationSeconds, timePerDoc/float64(time.Millisecond))
		}
		line, err = reader.ReadString('\n')
	}

	if err != io.EOF {
		log.Fatalln("[err] reading line ", count, "in", f.Name(), ":", err)
	}

	// flush the last batch
	if batchCount > 0 {
		err = index.Batch(batch)
		if err != nil {
			log.Fatal(err)
		}
	}

	indexDuration := time.Since(startTime)
	indexDurationSeconds := float64(indexDuration) / float64(time.Second)
	timePerDoc := float64(indexDuration) / float64(count)
	log.Printf("Indexed %d documents, in %.2fs (average %.2fms/doc)\n", count, indexDurationSeconds, timePerDoc/float64(time.Millisecond))
	log.Println("All Done")
}

//func buildIndexMapping() (*bleve.IndexMapping, error) {
//}

// v1
//		var si sra.SraItem
//		err := json.Unmarshal([]byte(line), &si)
//		if err != nil {
//			log.Fatal("Trouble unmarshal JSON to SraItem:", err)
//		}
//		//		log.Println("(", i, ")", "Indexing", si.Id)
//
//		err = index.Index(si.Id, si)
//		if err != nil {
//			log.Fatal("Trouble indexing sra item:", err)
//		}
//		i++
//		if i%25000 == 0 {
//			log.Println("Indexed", i, "entries")
//		}
<|endoftext|>


<commit_before>package toml_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"github.com/pelletier/go-toml/v2"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

//nolint:funlen
func TestMarshal(t *testing.T) {
	t.Parallel()

	examples := []struct {
		desc     string
		v        interface{}
		expected string
		err      bool
	}{
		{
			desc: "simple map and string",
			v: map[string]string{
				"hello": "world",
			},
			expected: "hello = 'world'",
		},
		{
			desc: "map with new line in key",
			v: map[string]string{
				"hel\nlo": "world",
			},
			err: true,
		},
		{
			desc: `map with " in key`,
			v: map[string]string{
				`hel"lo`: "world",
			},
			expected: `'hel"lo' = 'world'`,
		},
		{
			desc: "map in map and string",
			v: map[string]map[string]string{
				"table": {
					"hello": "world",
				},
			},
			expected: `
[table]
hello = 'world'`,
		},
		{
			desc: "map in map in map and string",
			v: map[string]map[string]map[string]string{
				"this": {
					"is": {
						"a": "test",
					},
				},
			},
			expected: `
[this]
[this.is]
a = 'test'`,
		},
		{
			desc: "map in map in map and string with values",
			v: map[string]interface{}{
				"this": map[string]interface{}{
					"is": map[string]string{
						"a": "test",
					},
					"also": "that",
				},
			},
			expected: `
[this]
also = 'that'
[this.is]
a = 'test'`,
		},
		{
			desc: "simple string array",
			v: map[string][]string{
				"array": {"one", "two", "three"},
			},
			expected: `array = ['one', 'two', 'three']`,
		},
		{
			desc:     "empty string array",
			v:        map[string][]string{},
			expected: ``,
		},
		{
			desc:     "map",
			v:        map[string][]string{},
			expected: ``,
		},
		{
			desc: "nested string arrays",
			v: map[string][][]string{
				"array": {{"one", "two"}, {"three"}},
			},
			expected: `array = [['one', 'two'], ['three']]`,
		},
		{
			desc: "mixed strings and nested string arrays",
			v: map[string][]interface{}{
				"array": {"a string", []string{"one", "two"}, "last"},
			},
			expected: `array = ['a string', ['one', 'two'], 'last']`,
		},
		{
			desc: "array of maps",
			v: map[string][]map[string]string{
				"top": {
					{"map1.1": "v1.1"},
					{"map2.1": "v2.1"},
				},
			},
			expected: `
[[top]]
'map1.1' = 'v1.1'
[[top]]
'map2.1' = 'v2.1'
`,
		},
		{
			desc: "map with two keys",
			v: map[string]string{
				"key1": "value1",
				"key2": "value2",
			},
			expected: `
key1 = 'value1'
key2 = 'value2'`,
		},
		{
			desc: "simple struct",
			v: struct {
				A string
			}{
				A: "foo",
			},
			expected: `A = 'foo'`,
		},
		{
			desc: "one level of structs within structs",
			v: struct {
				A interface{}
			}{
				A: struct {
					K1 string
					K2 string
				}{
					K1: "v1",
					K2: "v2",
				},
			},
			expected: `
[A]
K1 = 'v1'
K2 = 'v2'
`,
		},
		{
			desc: "structs in array with interfaces",
			v: map[string]interface{}{
				"root": map[string]interface{}{
					"nested": []interface{}{
						map[string]interface{}{"name": "Bob"},
						map[string]interface{}{"name": "Alice"},
					},
				},
			},
			expected: `
[root]
[[root.nested]]
name = 'Bob'
[[root.nested]]
name = 'Alice'
`,
		},
		{
			desc: "string escapes",
			v: map[string]interface{}{
				"a": `'"\`,
			},
			expected: `a = "'\"\\"`,
		},
		{
			desc: "string utf8 low",
			v: map[string]interface{}{
				"a": "'Ę",
			},
			expected: `a = "'Ę"`,
		},
		{
			desc: "string utf8 low 2",
			v: map[string]interface{}{
				"a": "'\u10A85",
			},
			expected: "a = \"'\u10A85\"",
		},
		{
			desc: "string utf8 low 2",
			v: map[string]interface{}{
				"a": "'\u10A85",
			},
			expected: "a = \"'\u10A85\"",
		},
		{
			desc: "emoji",
			v: map[string]interface{}{
				"a": "'😀",
			},
			expected: "a = \"'😀\"",
		},
		{
			desc: "control char",
			v: map[string]interface{}{
				"a": "'\u001A",
			},
			expected: `a = "'\u001A"`,
		},
		{
			desc: "multi-line string",
			v: map[string]interface{}{
				"a": "hello\nworld",
			},
			expected: `a = "hello\nworld"`,
		},
		{
			desc: "multi-line forced",
			v: struct {
				A string `multiline:"true"`
			}{
				A: "hello\nworld",
			},
			expected: `A = """
hello
world"""`,
		},
		{
			desc: "inline field",
			v: struct {
				A map[string]string `inline:"true"`
				B map[string]string
			}{
				A: map[string]string{
					"isinline": "yes",
				},
				B: map[string]string{
					"isinline": "no",
				},
			},
			expected: `
A = {isinline = 'yes'}
[B]
isinline = 'no'
`,
		},
		{
			desc: "mutiline array int",
			v: struct {
				A []int `multiline:"true"`
				B []int
			}{
				A: []int{1, 2, 3, 4},
				B: []int{1, 2, 3, 4},
			},
			expected: `
A = [
  1,
  2,
  3,
  4
]
B = [1, 2, 3, 4]
`,
		},
		{
			desc: "mutiline array in array",
			v: struct {
				A [][]int `multiline:"true"`
			}{
				A: [][]int{{1, 2}, {3, 4}},
			},
			expected: `
A = [
  [1, 2],
  [3, 4]
]
`,
		},
	}

	for _, e := range examples {
		e := e
		t.Run(e.desc, func(t *testing.T) {
			t.Parallel()

			b, err := toml.Marshal(e.v)
			if e.err {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			equalStringsIgnoreNewlines(t, e.expected, string(b))

			// make sure the output is always valid TOML
			defaultMap := map[string]interface{}{}
			err = toml.Unmarshal(b, &defaultMap)
			require.NoError(t, err)

			testWithAllFlags(t, func(t *testing.T, flags int) {
				var buf bytes.Buffer
				enc := toml.NewEncoder(&buf)
				setFlags(enc, flags)

				err := enc.Encode(e.v)
				require.NoError(t, err)

				inlineMap := map[string]interface{}{}
				err = toml.Unmarshal(buf.Bytes(), &inlineMap)
				require.NoError(t, err)

				require.Equal(t, defaultMap, inlineMap)
			})
		})
	}
}

type flagsSetters []struct {
	name string
	f    func(enc *toml.Encoder, flag bool)
}

var allFlags = flagsSetters{
	{"arrays-multiline", (*toml.Encoder).SetArraysMultiline},
	{"tables-inline", (*toml.Encoder).SetTablesInline},
}

func setFlags(enc *toml.Encoder, flags int) {
	for i := 0; i < len(allFlags); i++ {
		enabled := flags&1 > 0
		allFlags[i].f(enc, enabled)
	}
}

func testWithAllFlags(t *testing.T, testfn func(t *testing.T, flags int)) {
	t.Helper()
	testWithFlags(t, 0, allFlags, testfn)
}

func testWithFlags(t *testing.T, flags int, setters flagsSetters, testfn func(t *testing.T, flags int)) {
	t.Helper()

	if len(setters) == 0 {
		testfn(t, flags)
		return
	}

	s := setters[0]

	for _, enabled := range []bool{false, true} {
		name := fmt.Sprintf("%s=%t", s.name, enabled)
		newFlags := flags << 1
		if enabled {
			newFlags++
		}
		t.Run(name, func(t *testing.T) {
			testWithFlags(t, newFlags, setters[1:], testfn)
		})
	}
}

func equalStringsIgnoreNewlines(t *testing.T, expected string, actual string) {
	t.Helper()
	cutset := "\n"
	assert.Equal(t, strings.Trim(expected, cutset), strings.Trim(actual, cutset))
}

func TestIssue436(t *testing.T) {
	t.Parallel()

	data := []byte(`{"a": [ { "b": { "c": "d" } } ]}`)

	var v interface{}
	err := json.Unmarshal(data, &v)
	require.NoError(t, err)

	var buf bytes.Buffer
	err = toml.NewEncoder(&buf).Encode(v)
	require.NoError(t, err)

	expected := `
[[a]]
[a.b]
c = 'd'
`
	equalStringsIgnoreNewlines(t, expected, buf.String())
}
<commit_msg>Add test for issue #424<commit_after>package toml_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"github.com/pelletier/go-toml/v2"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

//nolint:funlen
func TestMarshal(t *testing.T) {
	t.Parallel()

	examples := []struct {
		desc     string
		v        interface{}
		expected string
		err      bool
	}{
		{
			desc: "simple map and string",
			v: map[string]string{
				"hello": "world",
			},
			expected: "hello = 'world'",
		},
		{
			desc: "map with new line in key",
			v: map[string]string{
				"hel\nlo": "world",
			},
			err: true,
		},
		{
			desc: `map with " in key`,
			v: map[string]string{
				`hel"lo`: "world",
			},
			expected: `'hel"lo' = 'world'`,
		},
		{
			desc: "map in map and string",
			v: map[string]map[string]string{
				"table": {
					"hello": "world",
				},
			},
			expected: `
[table]
hello = 'world'`,
		},
		{
			desc: "map in map in map and string",
			v: map[string]map[string]map[string]string{
				"this": {
					"is": {
						"a": "test",
					},
				},
			},
			expected: `
[this]
[this.is]
a = 'test'`,
		},
		{
			desc: "map in map in map and string with values",
			v: map[string]interface{}{
				"this": map[string]interface{}{
					"is": map[string]string{
						"a": "test",
					},
					"also": "that",
				},
			},
			expected: `
[this]
also = 'that'
[this.is]
a = 'test'`,
		},
		{
			desc: "simple string array",
			v: map[string][]string{
				"array": {"one", "two", "three"},
			},
			expected: `array = ['one', 'two', 'three']`,
		},
		{
			desc:     "empty string array",
			v:        map[string][]string{},
			expected: ``,
		},
		{
			desc:     "map",
			v:        map[string][]string{},
			expected: ``,
		},
		{
			desc: "nested string arrays",
			v: map[string][][]string{
				"array": {{"one", "two"}, {"three"}},
			},
			expected: `array = [['one', 'two'], ['three']]`,
		},
		{
			desc: "mixed strings and nested string arrays",
			v: map[string][]interface{}{
				"array": {"a string", []string{"one", "two"}, "last"},
			},
			expected: `array = ['a string', ['one', 'two'], 'last']`,
		},
		{
			desc: "array of maps",
			v: map[string][]map[string]string{
				"top": {
					{"map1.1": "v1.1"},
					{"map2.1": "v2.1"},
				},
			},
			expected: `
[[top]]
'map1.1' = 'v1.1'
[[top]]
'map2.1' = 'v2.1'
`,
		},
		{
			desc: "map with two keys",
			v: map[string]string{
				"key1": "value1",
				"key2": "value2",
			},
			expected: `
key1 = 'value1'
key2 = 'value2'`,
		},
		{
			desc: "simple struct",
			v: struct {
				A string
			}{
				A: "foo",
			},
			expected: `A = 'foo'`,
		},
		{
			desc: "one level of structs within structs",
			v: struct {
				A interface{}
			}{
				A: struct {
					K1 string
					K2 string
				}{
					K1: "v1",
					K2: "v2",
				},
			},
			expected: `
[A]
K1 = 'v1'
K2 = 'v2'
`,
		},
		{
			desc: "structs in array with interfaces",
			v: map[string]interface{}{
				"root": map[string]interface{}{
					"nested": []interface{}{
						map[string]interface{}{"name": "Bob"},
						map[string]interface{}{"name": "Alice"},
					},
				},
			},
			expected: `
[root]
[[root.nested]]
name = 'Bob'
[[root.nested]]
name = 'Alice'
`,
		},
		{
			desc: "string escapes",
			v: map[string]interface{}{
				"a": `'"\`,
			},
			expected: `a = "'\"\\"`,
		},
		{
			desc: "string utf8 low",
			v: map[string]interface{}{
				"a": "'Ę",
			},
			expected: `a = "'Ę"`,
		},
		{
			desc: "string utf8 low 2",
			v: map[string]interface{}{
				"a": "'\u10A85",
			},
			expected: "a = \"'\u10A85\"",
		},
		{
			desc: "string utf8 low 2",
			v: map[string]interface{}{
				"a": "'\u10A85",
			},
			expected: "a = \"'\u10A85\"",
		},
		{
			desc: "emoji",
			v: map[string]interface{}{
				"a": "'😀",
			},
			expected: "a = \"'😀\"",
		},
		{
			desc: "control char",
			v: map[string]interface{}{
				"a": "'\u001A",
			},
			expected: `a = "'\u001A"`,
		},
		{
			desc: "multi-line string",
			v: map[string]interface{}{
				"a": "hello\nworld",
			},
			expected: `a = "hello\nworld"`,
		},
		{
			desc: "multi-line forced",
			v: struct {
				A string `multiline:"true"`
			}{
				A: "hello\nworld",
			},
			expected: `A = """
hello
world"""`,
		},
		{
			desc: "inline field",
			v: struct {
				A map[string]string `inline:"true"`
				B map[string]string
			}{
				A: map[string]string{
					"isinline": "yes",
				},
				B: map[string]string{
					"isinline": "no",
				},
			},
			expected: `
A = {isinline = 'yes'}
[B]
isinline = 'no'
`,
		},
		{
			desc: "mutiline array int",
			v: struct {
				A []int `multiline:"true"`
				B []int
			}{
				A: []int{1, 2, 3, 4},
				B: []int{1, 2, 3, 4},
			},
			expected: `
A = [
  1,
  2,
  3,
  4
]
B = [1, 2, 3, 4]
`,
		},
		{
			desc: "mutiline array in array",
			v: struct {
				A [][]int `multiline:"true"`
			}{
				A: [][]int{{1, 2}, {3, 4}},
			},
			expected: `
A = [
  [1, 2],
  [3, 4]
]
`,
		},
	}

	for _, e := range examples {
		e := e
		t.Run(e.desc, func(t *testing.T) {
			t.Parallel()

			b, err := toml.Marshal(e.v)
			if e.err {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			equalStringsIgnoreNewlines(t, e.expected, string(b))

			// make sure the output is always valid TOML
			defaultMap := map[string]interface{}{}
			err = toml.Unmarshal(b, &defaultMap)
			require.NoError(t, err)

			testWithAllFlags(t, func(t *testing.T, flags int) {
				var buf bytes.Buffer
				enc := toml.NewEncoder(&buf)
				setFlags(enc, flags)

				err := enc.Encode(e.v)
				require.NoError(t, err)

				inlineMap := map[string]interface{}{}
				err = toml.Unmarshal(buf.Bytes(), &inlineMap)
				require.NoError(t, err)

				require.Equal(t, defaultMap, inlineMap)
			})
		})
	}
}

type flagsSetters []struct {
	name string
	f    func(enc *toml.Encoder, flag bool)
}

var allFlags = flagsSetters{
	{"arrays-multiline", (*toml.Encoder).SetArraysMultiline},
	{"tables-inline", (*toml.Encoder).SetTablesInline},
}

func setFlags(enc *toml.Encoder, flags int) {
	for i := 0; i < len(allFlags); i++ {
		enabled := flags&1 > 0
		allFlags[i].f(enc, enabled)
	}
}

func testWithAllFlags(t *testing.T, testfn func(t *testing.T, flags int)) {
	t.Helper()
	testWithFlags(t, 0, allFlags, testfn)
}

func testWithFlags(t *testing.T, flags int, setters flagsSetters, testfn func(t *testing.T, flags int)) {
	t.Helper()

	if len(setters) == 0 {
		testfn(t, flags)
		return
	}

	s := setters[0]

	for _, enabled := range []bool{false, true} {
		name := fmt.Sprintf("%s=%t", s.name, enabled)
		newFlags := flags << 1
		if enabled {
			newFlags++
		}
		t.Run(name, func(t *testing.T) {
			testWithFlags(t, newFlags, setters[1:], testfn)
		})
	}
}

func equalStringsIgnoreNewlines(t *testing.T, expected string, actual string) {
	t.Helper()
	cutset := "\n"
	assert.Equal(t, strings.Trim(expected, cutset), strings.Trim(actual, cutset))
}

func TestIssue436(t *testing.T) {
	t.Parallel()

	data := []byte(`{"a": [ { "b": { "c": "d" } } ]}`)

	var v interface{}
	err := json.Unmarshal(data, &v)
	require.NoError(t, err)

	var buf bytes.Buffer
	err = toml.NewEncoder(&buf).Encode(v)
	require.NoError(t, err)

	expected := `
[[a]]
[a.b]
c = 'd'
`
	equalStringsIgnoreNewlines(t, expected, buf.String())
}

func TestIssue424(t *testing.T) {
	type Message1 struct {
		Text string
	}

	type Message2 struct {
		Text string `multiline:"true"`
	}

	msg1 := Message1{"Hello\\World"}
	msg2 := Message2{"Hello\\World"}

	toml1, err := toml.Marshal(msg1)
	require.NoError(t, err)

	toml2, err := toml.Marshal(msg2)
	require.NoError(t, err)

	msg1parsed := Message1{}
	err = toml.Unmarshal(toml1, &msg1parsed)
	require.NoError(t, err)
	require.Equal(t, msg1, msg1parsed)

	msg2parsed := Message2{}
	err = toml.Unmarshal(toml2, &msg2parsed)
	require.NoError(t, err)
	require.Equal(t, msg2, msg2parsed)
}
<|endoftext|>


<commit_before>package dcnet

import (
	"crypto/rand"
	"crypto/sha256"
	"github.com/dedis/prifi/prifi-lib/config"
	"gopkg.in/dedis/kyber.v2"
	"gopkg.in/dedis/kyber.v2/suites"
	"gopkg.in/dedis/onet.v2/log"
)

// Clients compute:
// kappa_i = k_i + h * SUM_j(q_ij), where q_ij = H(p_ij) in group
// c' = k_i + c
//
// Trustees compute:
// sigma_i = SUM_i(q_ij), where q_ij = H(s_ij) in group
//
// Relay compute:
// k_i = SUM_i(kappa_i) - h * (SUM_j(sigma_i))
//     = SUM_i(h * SUM_j(q_ij)) + k_i - h * SUM_j(SUM_i(q_ij))
// c = k_i + c'
//

// Equivocation holds the functions needed for equivocation protection
type EquivocationProtection struct {
	history    kyber.Scalar
	randomness kyber.XOF
	suite      suites.Suite
}

// NewEquivocation creates the structure that handle equivocation protection
func NewEquivocation() *EquivocationProtection {
	e := new(EquivocationProtection)
	e.suite = config.CryptoSuite
	e.history = e.suite.Scalar().One()

	randomKey := make([]byte, 100)
	rand.Read(randomKey)
	e.randomness = e.suite.XOF(randomKey)

	return e
}

func (e *EquivocationProtection) randomScalar() kyber.Scalar {
	return e.suite.Scalar().Pick(e.randomness)
}

func (e *EquivocationProtection) hashInGroup(data []byte) kyber.Scalar {
	return e.suite.Scalar().SetBytes(data)
}

// Update History adds those bits to the history hash chain
func (e *EquivocationProtection) UpdateHistory(data []byte) {
	historyB, err := e.history.MarshalBinary()
	if err != nil {
		log.Fatal("Could not unmarshall bytes", err)
	}
	toBeHashed := make([]byte, len(historyB)+len(data))
	newPayload := sha256.Sum256(toBeHashed)
	e.history.SetBytes(newPayload[:])
}

// a function that takes a payload x, encrypt it as x' = x + k, and returns x' and kappa = k + history * (sum of the (hashes of pads))
func (e *EquivocationProtection) ClientEncryptPayload(slotOwner bool, x []byte, p_j [][]byte) ([]byte, []byte) {

	// hash the pads p_i into q_i
	q_j := make([]kyber.Scalar, len(p_j))
	for trustee_j := range q_j {
		q_j[trustee_j] = e.hashInGroup(p_j[trustee_j])
	}

	// sum of q_i
	sum := e.suite.Scalar().Zero()
	for _, p := range q_j {
		sum = sum.Add(sum, p)
	}

	product := sum.Mul(sum, e.history)

	//we're not the slot owner
	if !slotOwner {
		kappa_i := product
		kappa_i_bytes, err := kappa_i.MarshalBinary()
		if err != nil {
			log.Fatal("Couldn't marshall", err)
		}
		return x, kappa_i_bytes
	}

	k_i := e.randomScalar()
	k_i_bytes, err := k_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}

	// encrypt payload
	for i := range x {
		x[i] ^= k_i_bytes[i%len(k_i_bytes)]
	}

	// compute kappa
	kappa_i := k_i.Add(k_i, product)
	kappa_i_bytes, err := kappa_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}
	return x, kappa_i_bytes
}

// a function that takes returns the byte[] version of sigma_j
func (e *EquivocationProtection) TrusteeGetContribution(s_i [][]byte) []byte {

	// hash the pads p_i into q_i
	q_i := make([]kyber.Scalar, len(s_i))
	for client_i := range q_i {
		q_i[client_i] = e.hashInGroup(s_i[client_i])
	}

	// sum of q_i
	sum := e.suite.Scalar().Zero()
	for _, p := range q_i {
		sum = sum.Add(sum, p)
	}

	kappa_j := sum

	kappa_j_bytes, err := kappa_j.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}
	return kappa_j_bytes
}

// given all contributions, decodes the payload
func (e *EquivocationProtection) RelayDecode(encryptedPayload []byte, trusteesContributions [][]byte, clientsContributions [][]byte) []byte {

	//reconstitute the abstract.Point values
	trustee_kappa_j := make([]kyber.Scalar, len(trusteesContributions))
	for k, v := range trusteesContributions {
		trustee_kappa_j[k] = e.suite.Scalar().SetBytes(v)
	}
	client_kappa_i := make([]kyber.Scalar, len(clientsContributions))
	for k, v := range clientsContributions {
		client_kappa_i[k] = e.suite.Scalar().SetBytes(v)
	}

	// compute sum of trustees contribs
	sumTrustees := e.suite.Scalar().Zero()
	for _, v := range trustee_kappa_j {
		sumTrustees = sumTrustees.Add(sumTrustees, v)
	}

	// compute sum of clients contribs
	sumClients := e.suite.Scalar().Zero()
	for _, v := range client_kappa_i {
		//log.Lvl1("Adding in", v, "value is now", sumClients)
		sumClients = sumClients.Add(sumClients, v)
	}

	prod := sumTrustees.Mul(sumTrustees, e.history)
	k_i := sumClients.Sub(sumClients, prod)

	//now use k to decrypt the payload
	k_bytes, err := k_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}

	if len(k_bytes) == 0 {
		log.Lvl1("Error: Equivocation couldn't recover the k_bytes value")
		log.Lvl1("encryptedPayload:", encryptedPayload)
		for k, v := range trusteesContributions {
			log.Lvl1("trusteesContributions:", v, "mapped to", trustee_kappa_j[k])
		}
		for k, v := range clientsContributions {
			log.Lvl1("clientsContributions:", v, "mapped to", client_kappa_i[k])
		}
		log.Lvl1("sumTrustees:", sumTrustees)
		log.Lvl1("sumClients:", sumClients)
		log.Lvl1("history:", e.history)
		log.Lvl1("prod:", prod)
		log.Lvl1("k_i:", k_i)
		return make([]byte, 0)
	}

	// decrypt the payload
	for i := range encryptedPayload {
		encryptedPayload[i] ^= k_bytes[i%len(k_bytes)]
	}

	return encryptedPayload
}
<commit_msg>Some placeholders in the Equiv<commit_after>package dcnet

import (
	"crypto/rand"
	"crypto/sha256"
	"github.com/dedis/prifi/prifi-lib/config"
	"gopkg.in/dedis/kyber.v2"
	"gopkg.in/dedis/kyber.v2/suites"
	"gopkg.in/dedis/onet.v2/log"
)

// Clients compute:
// kappa_i = k_i + h * SUM_j(q_ij), where q_ij = H(p_ij) in group
// c' = k_i + c
//
// Trustees compute:
// sigma_i = SUM_i(q_ij), where q_ij = H(s_ij) in group
//
// Relay compute:
// k_i = SUM_i(kappa_i) - h * (SUM_j(sigma_i))
//     = SUM_i(h * SUM_j(q_ij)) + k_i - h * SUM_j(SUM_i(q_ij))
// c = k_i + c'
//

// Equivocation holds the functions needed for equivocation protection
type EquivocationProtection struct {
	history    kyber.Scalar
	randomness kyber.XOF
	suite      suites.Suite
}

// NewEquivocation creates the structure that handle equivocation protection
func NewEquivocation() *EquivocationProtection {
	e := new(EquivocationProtection)
	e.suite = config.CryptoSuite
	e.history = e.suite.Scalar().One()

	randomKey := make([]byte, 32)
	rand.Read(randomKey)
	e.randomness = e.suite.XOF(randomKey)

	return e
}

func (e *EquivocationProtection) randomScalar() kyber.Scalar {
	return e.suite.Scalar().Pick(e.randomness)
}

func (e *EquivocationProtection) hashInGroup(data []byte) kyber.Scalar {
	return e.suite.Scalar().SetBytes(data)
}

// Update History adds those bits to the history hash chain
func (e *EquivocationProtection) UpdateHistory(data []byte) {
	historyB, err := e.history.MarshalBinary()
	if err != nil {
		log.Fatal("Could not unmarshall bytes", err)
	}
	toBeHashed := make([]byte, len(historyB)+len(data))
	newPayload := sha256.Sum256(toBeHashed)
	e.history.SetBytes(newPayload[:])
}

// a function that takes a payload x, encrypt it as x' = x + k, and returns x' and kappa = k + history * (sum of the (hashes of pads))
func (e *EquivocationProtection) ClientEncryptPayload(slotOwner bool, x []byte, p_j [][]byte) ([]byte, []byte) {

	// hash the pads p_i into q_i
	q_j := make([]kyber.Scalar, len(p_j))
	for trustee_j := range q_j {
		q_j[trustee_j] = e.hashInGroup(p_j[trustee_j])
	}

	// sum of q_i
	sum := e.suite.Scalar().Zero()
	for _, p := range q_j {
		sum = sum.Add(sum, p)
	}

	product := sum.Mul(sum, e.history)

	//we're not the slot owner
	if !slotOwner {
		kappa_i := product
		kappa_i_bytes, err := kappa_i.MarshalBinary()
		if err != nil {
			log.Fatal("Couldn't marshall", err)
		}
		return x, kappa_i_bytes
	}

	k_i := e.randomScalar()
	k_i_bytes, err := k_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}

	// encrypt payload
	// LB->CV: Replace by traditional encryption (AES-GCM or what not) instead of XOR
	for i := range x {
		x[i] ^= k_i_bytes[i%len(k_i_bytes)]
	}

	// compute kappa
	kappa_i := k_i.Add(k_i, product)
	kappa_i_bytes, err := kappa_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}
	return x, kappa_i_bytes
}

func (e *EquivocationProtection) ClientProve() []byte {

	return nil
}

// a function that takes returns the byte[] version of sigma_j
func (e *EquivocationProtection) TrusteeGetContribution(s_i [][]byte) []byte {

	// hash the pads p_i into q_i
	q_i := make([]kyber.Scalar, len(s_i))
	for client_i := range q_i {
		q_i[client_i] = e.hashInGroup(s_i[client_i])
	}

	// sum of q_i
	sum := e.suite.Scalar().Zero()
	for _, p := range q_i {
		sum = sum.Add(sum, p)
	}

	kappa_j := sum

	kappa_j_bytes, err := kappa_j.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}
	return kappa_j_bytes
}

func (e *EquivocationProtection) TrusteeProve() []byte {

	return nil
}

// given all contributions, decodes the payload
func (e *EquivocationProtection) RelayDecode(encryptedPayload []byte, trusteesContributions [][]byte, clientsContributions [][]byte) []byte {

	//reconstitute the abstract.Point values
	trustee_kappa_j := make([]kyber.Scalar, len(trusteesContributions))
	for k, v := range trusteesContributions {
		trustee_kappa_j[k] = e.suite.Scalar().SetBytes(v)
	}
	client_kappa_i := make([]kyber.Scalar, len(clientsContributions))
	for k, v := range clientsContributions {
		client_kappa_i[k] = e.suite.Scalar().SetBytes(v)
	}

	// compute sum of trustees contribs
	sumTrustees := e.suite.Scalar().Zero()
	for _, v := range trustee_kappa_j {
		sumTrustees = sumTrustees.Add(sumTrustees, v)
	}

	// compute sum of clients contribs
	sumClients := e.suite.Scalar().Zero()
	for _, v := range client_kappa_i {
		//log.Lvl1("Adding in", v, "value is now", sumClients)
		sumClients = sumClients.Add(sumClients, v)
	}

	prod := sumTrustees.Mul(sumTrustees, e.history)
	k_i := sumClients.Sub(sumClients, prod)

	//now use k to decrypt the payload
	k_bytes, err := k_i.MarshalBinary()
	if err != nil {
		log.Fatal("Couldn't marshall", err)
	}

	if len(k_bytes) == 0 {
		log.Lvl1("Error: Equivocation couldn't recover the k_bytes value")
		log.Lvl1("encryptedPayload:", encryptedPayload)
		for k, v := range trusteesContributions {
			log.Lvl1("trusteesContributions:", v, "mapped to", trustee_kappa_j[k])
		}
		for k, v := range clientsContributions {
			log.Lvl1("clientsContributions:", v, "mapped to", client_kappa_i[k])
		}
		log.Lvl1("sumTrustees:", sumTrustees)
		log.Lvl1("sumClients:", sumClients)
		log.Lvl1("history:", e.history)
		log.Lvl1("prod:", prod)
		log.Lvl1("k_i:", k_i)
		return make([]byte, 0)
	}

	// decrypt the payload
	for i := range encryptedPayload {
		encryptedPayload[i] ^= k_bytes[i%len(k_bytes)]
	}

	return encryptedPayload
}
<|endoftext|>