<commit_before>package zookeeper

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

// Zookeeper is a zookeeper plugin
type Zookeeper struct {
	Servers []string
}

var sampleConfig = `
  ## An array of address to gather stats about. Specify an ip or hostname
  ## with port. ie localhost:2181, 10.0.0.1:2181, etc.

  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 2181 is used
  servers = [":2181"]
`

var defaultTimeout = time.Second * time.Duration(5)

// SampleConfig returns sample configuration message
func (z *Zookeeper) SampleConfig() string {
	return sampleConfig
}

// Description returns description of Zookeeper plugin
func (z *Zookeeper) Description() string {
	return `Reads 'mntr' stats from one or many zookeeper servers`
}

// Gather reads stats from all configured servers accumulates stats
func (z *Zookeeper) Gather(acc telegraf.Accumulator) error {
	if len(z.Servers) == 0 {
		return nil
	}

	for _, serverAddress := range z.Servers {
		acc.AddError(z.gatherServer(serverAddress, acc))
	}
	return nil
}

func (z *Zookeeper) gatherServer(address string, acc telegraf.Accumulator) error {
	var zookeeper_state string
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		address = address + ":2181"
	}

	c, err := net.DialTimeout("tcp", address, defaultTimeout)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return err
	}
	defer c.Close()

	// Extend connection
	c.SetDeadline(time.Now().Add(defaultTimeout))

	fmt.Fprintf(c, "%s\n", "mntr")
	rdr := bufio.NewReader(c)
	scanner := bufio.NewScanner(rdr)

	service := strings.Split(address, ":")
	if len(service) != 2 {
		return fmt.Errorf("Invalid service address: %s", address)
	}

	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()

		re := regexp.MustCompile(`^zk_(\w+)\s+([\w\.\-]+)`)
		parts := re.FindStringSubmatch(string(line))

		if len(parts) != 3 {
			return fmt.Errorf("unexpected line in mntr response: %q", line)
		}

		measurement := strings.TrimPrefix(parts[1], "zk_")
		if measurement == "server_state" {
			zookeeper_state = parts[2]
		} else {
			sValue := string(parts[2])

			iVal, err := strconv.ParseInt(sValue, 10, 64)
			if err == nil {
				fields[measurement] = iVal
			} else {
				fields[measurement] = sValue
			}
		}
	}
	tags := map[string]string{
		"server": service[0],
		"port":   service[1],
		"state":  zookeeper_state,
	}
	acc.AddFields("zookeeper", fields, tags)

	return nil
}

func init() {
	inputs.Add("zookeeper", func() telegraf.Input {
		return &Zookeeper{}
	})
}
<commit_msg>Default to localhost if zookeeper has no servers set (#3056)<commit_after>package zookeeper

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

// Zookeeper is a zookeeper plugin
type Zookeeper struct {
	Servers []string
}

var sampleConfig = `
  ## An array of address to gather stats about. Specify an ip or hostname
  ## with port. ie localhost:2181, 10.0.0.1:2181, etc.

  ## If no servers are specified, then localhost is used as the host.
  ## If no port is specified, 2181 is used
  servers = [":2181"]
`

var defaultTimeout = time.Second * time.Duration(5)

// SampleConfig returns sample configuration message
func (z *Zookeeper) SampleConfig() string {
	return sampleConfig
}

// Description returns description of Zookeeper plugin
func (z *Zookeeper) Description() string {
	return `Reads 'mntr' stats from one or many zookeeper servers`
}

// Gather reads stats from all configured servers accumulates stats
func (z *Zookeeper) Gather(acc telegraf.Accumulator) error {
	if len(z.Servers) == 0 {
		z.Servers = []string{":2181"}
	}

	for _, serverAddress := range z.Servers {
		acc.AddError(z.gatherServer(serverAddress, acc))
	}
	return nil
}

func (z *Zookeeper) gatherServer(address string, acc telegraf.Accumulator) error {
	var zookeeper_state string
	_, _, err := net.SplitHostPort(address)
	if err != nil {
		address = address + ":2181"
	}

	c, err := net.DialTimeout("tcp", address, defaultTimeout)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return err
	}
	defer c.Close()

	// Extend connection
	c.SetDeadline(time.Now().Add(defaultTimeout))

	fmt.Fprintf(c, "%s\n", "mntr")
	rdr := bufio.NewReader(c)
	scanner := bufio.NewScanner(rdr)

	service := strings.Split(address, ":")
	if len(service) != 2 {
		return fmt.Errorf("Invalid service address: %s", address)
	}

	fields := make(map[string]interface{})
	for scanner.Scan() {
		line := scanner.Text()

		re := regexp.MustCompile(`^zk_(\w+)\s+([\w\.\-]+)`)
		parts := re.FindStringSubmatch(string(line))

		if len(parts) != 3 {
			return fmt.Errorf("unexpected line in mntr response: %q", line)
		}

		measurement := strings.TrimPrefix(parts[1], "zk_")
		if measurement == "server_state" {
			zookeeper_state = parts[2]
		} else {
			sValue := string(parts[2])

			iVal, err := strconv.ParseInt(sValue, 10, 64)
			if err == nil {
				fields[measurement] = iVal
			} else {
				fields[measurement] = sValue
			}
		}
	}
	tags := map[string]string{
		"server": service[0],
		"port":   service[1],
		"state":  zookeeper_state,
	}
	acc.AddFields("zookeeper", fields, tags)

	return nil
}

func init() {
	inputs.Add("zookeeper", func() telegraf.Input {
		return &Zookeeper{}
	})
}
<|endoftext|>


<commit_before>package nodetemplate

import (
	"fmt"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	"github.com/rancher/norman/types/values"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/ref"
	corev1 "github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/labels"
	"strings"
)

type Store struct {
	types.Store
	NodePoolLister        v3.NodePoolLister
	CloudCredentialLister corev1.SecretLister
}

func (s *Store) Delete(apiContext *types.APIContext, schema *types.Schema, id string) (map[string]interface{}, error) {
	pools, err := s.NodePoolLister.List("", labels.Everything())
	if err != nil {
		return nil, err
	}
	for _, pool := range pools {
		if pool.Spec.NodeTemplateName == id {
			return nil, httperror.NewAPIError(httperror.MethodNotAllowed, "Template is in use by a node pool.")
		}
	}
	return s.Store.Delete(apiContext, schema, id)
}

func (s *Store) Create(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}) (map[string]interface{}, error) {
	if err := s.replaceCloudCredFields(data); err != nil {
		return data, err
	}
	return s.Store.Create(apiContext, schema, data)
}

func (s *Store) Update(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}, id string) (map[string]interface{}, error) {
	if err := s.replaceCloudCredFields(data); err != nil {
		return data, err
	}
	return s.Store.Update(apiContext, schema, data, id)
}

func (s *Store) replaceCloudCredFields(data map[string]interface{}) error {
	credID := convert.ToString(values.GetValueN(data, "cloudCredentialId"))
	if credID == "" {
		return nil
	}
	ns, name := ref.Parse(credID)
	if ns == "" || name == "" {
		return fmt.Errorf("invalid credID %s", credID)
	}
	cred, err := s.CloudCredentialLister.Get(namespace.GlobalNamespace, name)
	if err != nil {
		return fmt.Errorf("error getting cloud cred %s: %v", credID, err)
	}
	if len(cred.Data) == 0 {
		return fmt.Errorf("empty credID data %s", credID)
	}
	configName, credConfigName := "", ""
	for key := range cred.Data {
		splitKey := strings.SplitN(key, "-", 2)
		if len(splitKey) == 2 && strings.HasSuffix(splitKey[0], "credentialConfig") {
			configName = strings.Replace(splitKey[0], "credential", "", 1)
			credConfigName = splitKey[0]
			break
		}
	}
	if configName == "" {
		return fmt.Errorf("empty configName for credID %s", configName)
	}
	toReplace := convert.ToMapInterface(values.GetValueN(data, configName))
	if len(toReplace) == 0 {
		return nil
	}
	var fields []string
	for key := range cred.Data {
		splitKey := strings.SplitN(key, "-", 2)
		if len(splitKey) == 2 && splitKey[0] == credConfigName {
			toReplace[splitKey[1]] = ""
			fields = append(fields, splitKey[1])
		}
	}
	logrus.Debugf("replaceCloudCredFields: %v for credID %s", fields, credID)
	return nil
}
<commit_msg>remove fields from config if credID is set<commit_after>package nodetemplate

import (
	"fmt"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/norman/types"
	"github.com/rancher/norman/types/convert"
	"github.com/rancher/norman/types/values"
	"github.com/rancher/rancher/pkg/namespace"
	"github.com/rancher/rancher/pkg/ref"
	corev1 "github.com/rancher/types/apis/core/v1"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/labels"
	"strings"
)

type Store struct {
	types.Store
	NodePoolLister        v3.NodePoolLister
	CloudCredentialLister corev1.SecretLister
}

func (s *Store) Delete(apiContext *types.APIContext, schema *types.Schema, id string) (map[string]interface{}, error) {
	pools, err := s.NodePoolLister.List("", labels.Everything())
	if err != nil {
		return nil, err
	}
	for _, pool := range pools {
		if pool.Spec.NodeTemplateName == id {
			return nil, httperror.NewAPIError(httperror.MethodNotAllowed, "Template is in use by a node pool.")
		}
	}
	return s.Store.Delete(apiContext, schema, id)
}

func (s *Store) Create(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}) (map[string]interface{}, error) {
	if err := s.replaceCloudCredFields(data); err != nil {
		return data, err
	}
	return s.Store.Create(apiContext, schema, data)
}

func (s *Store) Update(apiContext *types.APIContext, schema *types.Schema, data map[string]interface{}, id string) (map[string]interface{}, error) {
	if err := s.replaceCloudCredFields(data); err != nil {
		return data, err
	}
	return s.Store.Update(apiContext, schema, data, id)
}

func (s *Store) replaceCloudCredFields(data map[string]interface{}) error {
	credID := convert.ToString(values.GetValueN(data, "cloudCredentialId"))
	if credID == "" {
		return nil
	}
	ns, name := ref.Parse(credID)
	if ns == "" || name == "" {
		return fmt.Errorf("invalid credID %s", credID)
	}
	cred, err := s.CloudCredentialLister.Get(namespace.GlobalNamespace, name)
	if err != nil {
		return fmt.Errorf("error getting cloud cred %s: %v", credID, err)
	}
	if len(cred.Data) == 0 {
		return fmt.Errorf("empty credID data %s", credID)
	}
	configName, credConfigName := "", ""
	for key := range cred.Data {
		splitKey := strings.SplitN(key, "-", 2)
		if len(splitKey) == 2 && strings.HasSuffix(splitKey[0], "credentialConfig") {
			configName = strings.Replace(splitKey[0], "credential", "", 1)
			credConfigName = splitKey[0]
			break
		}
	}
	if configName == "" {
		return fmt.Errorf("empty configName for credID %s", configName)
	}
	toReplace := convert.ToMapInterface(values.GetValueN(data, configName))
	if len(toReplace) == 0 {
		return nil
	}
	var fields []string
	for key := range cred.Data {
		splitKey := strings.SplitN(key, "-", 2)
		if len(splitKey) == 2 && splitKey[0] == credConfigName {
			delete(toReplace, splitKey[1])
			fields = append(fields, splitKey[1])
		}
	}
	logrus.Debugf("replaceCloudCredFields: %v for credID %s", fields, credID)
	return nil
}
<|endoftext|>


<commit_before>package daemon

import (
	"context"
	"net"
	"strings"
	"time"

	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	"github.com/telepresenceio/telepresence/v2/pkg/client/daemon/dbus"
	"github.com/telepresenceio/telepresence/v2/pkg/client/daemon/dns"
	"github.com/telepresenceio/telepresence/v2/pkg/tun"
)

func (o *outbound) tryResolveD(c context.Context, dev *tun.Device) error {
	// Connect to ResolveD via DBUS.
	if !dbus.IsResolveDRunning(c) {
		dlog.Error(c, "systemd-resolved is not running")
		return errResolveDNotConfigured
	}

	o.setSearchPathFunc = func(c context.Context, paths []string) {
		// When using systemd-resolved, we provide resolution of NAME.NAMESPACE by adding each
		// namespace as a route (a search entry prefixed with ~)
		namespaces := make(map[string]struct{})
		search := make([]string, 0)
		for i, path := range paths {
			if strings.ContainsRune(path, '.') {
				search = append(search, path)
			} else {
				namespaces[path] = struct{}{}
				// Turn namespace into a route
				paths[i] = "~" + path
			}
		}
		for _, sfx := range o.dnsConfig.IncludeSuffixes {
			paths = append(paths, "~"+strings.TrimPrefix(sfx, "."))
		}
		paths = append(paths, kubernetesZone+".")
		namespaces[tel2SubDomain] = struct{}{}

		o.domainsLock.Lock()
		o.namespaces = namespaces
		o.search = search
		o.domainsLock.Unlock()
		if err := dbus.SetLinkDomains(c, int(dev.Index()), paths...); err != nil {
			dlog.Errorf(c, "failed to set link domains on %q: %v", dev.Name(), err)
		} else {
			dlog.Debugf(c, "Link domains on device %q set to [%s]", dev.Name(), strings.Join(paths, ","))
		}
	}

	g := dgroup.NewGroup(c, dgroup.GroupConfig{})

	// DNS resolver
	initDone := make(chan struct{})

	var dnsServer *dns.Server
	g.Go("Server", func(c context.Context) error {
		select {
		case <-c.Done():
			initDone <- struct{}{}
			return nil
		case dnsIP := <-o.kubeDNS:
			listeners, err := o.dnsListeners(c)
			if err != nil {
				dlog.Error(c, err)
				initDone <- struct{}{}
				return err
			}
			// Create a new local address that the DNS resolver can listen to.
			dnsResolverAddr, err := splitToUDPAddr(listeners[0].LocalAddr())
			if err != nil {
				return err
			}

			o.router.configureDNS(c, dnsIP, uint16(53), dnsResolverAddr)
			dlog.Infof(c, "Configuring DNS IP %s", dnsIP)
			if err = dbus.SetLinkDNS(c, int(dev.Index()), dnsIP); err != nil {
				dlog.Error(c, err)
				initDone <- struct{}{}
				return errResolveDNotConfigured
			}
			defer func() {
				for _, l := range listeners {
					_ = l.Close()
				}
			}()
			dnsServer = dns.NewServer(c, listeners, nil, o.resolveInCluster)
			o.router.configureDNS(c, dnsIP, uint16(53), dnsResolverAddr)
			close(initDone)
			return dnsServer.Run(c)
		}
	})
	g.Go("SanityCheck", func(c context.Context) error {
		if _, ok := <-initDone; ok {
			// initDonw was not closed, bail out.
			return errResolveDNotConfigured
		}

		// Check if an attempt to resolve a DNS address reaches our DNS resolver, One second should be plenty

		cmdC, cmdCancel := context.WithTimeout(c, time.Second)
		defer cmdCancel()
		for cmdC.Err() == nil {
			_, _ = net.DefaultResolver.LookupHost(cmdC, "jhfweoitnkgyeta."+tel2SubDomain)
			if dnsServer.RequestCount() > 0 {
				close(o.dnsConfigured)
				return nil
			}
			dtime.SleepWithContext(cmdC, 100*time.Millisecond)
		}
		dlog.Error(c, "resolver did not receive requests from systemd-resolved")
		return errResolveDNotConfigured
	})
	return g.Wait()
}
<commit_msg>Ensure that the systemd-resolved DNS service is properly reverted<commit_after>package daemon

import (
	"context"
	"net"
	"strings"
	"time"

	"github.com/datawire/dlib/dcontext"
	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/dlib/dtime"
	"github.com/telepresenceio/telepresence/v2/pkg/client/daemon/dbus"
	"github.com/telepresenceio/telepresence/v2/pkg/client/daemon/dns"
	"github.com/telepresenceio/telepresence/v2/pkg/tun"
)

func (o *outbound) tryResolveD(c context.Context, dev *tun.Device) error {
	// Connect to ResolveD via DBUS.
	if !dbus.IsResolveDRunning(c) {
		dlog.Error(c, "systemd-resolved is not running")
		return errResolveDNotConfigured
	}

	o.setSearchPathFunc = func(c context.Context, paths []string) {
		// When using systemd-resolved, we provide resolution of NAME.NAMESPACE by adding each
		// namespace as a route (a search entry prefixed with ~)
		namespaces := make(map[string]struct{})
		search := make([]string, 0)
		for i, path := range paths {
			if strings.ContainsRune(path, '.') {
				search = append(search, path)
			} else {
				namespaces[path] = struct{}{}
				// Turn namespace into a route
				paths[i] = "~" + path
			}
		}
		for _, sfx := range o.dnsConfig.IncludeSuffixes {
			paths = append(paths, "~"+strings.TrimPrefix(sfx, "."))
		}
		paths = append(paths, kubernetesZone+".")
		namespaces[tel2SubDomain] = struct{}{}

		o.domainsLock.Lock()
		o.namespaces = namespaces
		o.search = search
		o.domainsLock.Unlock()
		if err := dbus.SetLinkDomains(c, int(dev.Index()), paths...); err != nil {
			dlog.Errorf(c, "failed to set link domains on %q: %v", dev.Name(), err)
		} else {
			dlog.Debugf(c, "Link domains on device %q set to [%s]", dev.Name(), strings.Join(paths, ","))
		}
	}

	g := dgroup.NewGroup(c, dgroup.GroupConfig{})

	// DNS resolver
	initDone := make(chan struct{})

	var dnsServer *dns.Server
	g.Go("Server", func(c context.Context) error {
		select {
		case <-c.Done():
			initDone <- struct{}{}
			return nil
		case dnsIP := <-o.kubeDNS:
			listeners, err := o.dnsListeners(c)
			if err != nil {
				dlog.Error(c, err)
				initDone <- struct{}{}
				return err
			}
			// Create a new local address that the DNS resolver can listen to.
			dnsResolverAddr, err := splitToUDPAddr(listeners[0].LocalAddr())
			if err != nil {
				return err
			}

			o.router.configureDNS(c, dnsIP, uint16(53), dnsResolverAddr)
			dlog.Infof(c, "Configuring DNS IP %s", dnsIP)
			if err = dbus.SetLinkDNS(c, int(dev.Index()), dnsIP); err != nil {
				dlog.Error(c, err)
				initDone <- struct{}{}
				return errResolveDNotConfigured
			}
			defer func() {
				// It's very likely that the context is cancelled here. We use it
				// anyway, stripped from cancellation, to retain logging.
				c, cancel := context.WithTimeout(dcontext.WithoutCancel(c), time.Second)
				defer cancel()
				dlog.Debugf(c, "Reverting Link settings for %s", dev.Name())
				o.setSearchPathFunc = nil
				o.router.configureDNS(c, nil, 0, nil) // Don't route from TUN-device
				if err = dbus.RevertLink(c, int(dev.Index())); err != nil {
					dlog.Error(c, err)
				}

				// No need to close listeners here. They are closed by the dnsServer
			}()
			dnsServer = dns.NewServer(c, listeners, nil, o.resolveInCluster)
			close(initDone)
			return dnsServer.Run(c)
		}
	})
	g.Go("SanityCheck", func(c context.Context) error {
		if _, ok := <-initDone; ok {
			// initDonw was not closed, bail out.
			return errResolveDNotConfigured
		}

		// Check if an attempt to resolve a DNS address reaches our DNS resolver, One second should be plenty

		cmdC, cmdCancel := context.WithTimeout(c, time.Second)
		defer cmdCancel()
		for cmdC.Err() == nil {
			_, _ = net.DefaultResolver.LookupHost(cmdC, "jhfweoitnkgyeta."+tel2SubDomain)
			if dnsServer.RequestCount() > 0 {
				close(o.dnsConfigured)
				return nil
			}
			dtime.SleepWithContext(cmdC, 100*time.Millisecond)
		}
		dlog.Error(c, "resolver did not receive requests from systemd-resolved")
		return errResolveDNotConfigured
	})
	return g.Wait()
}
<|endoftext|>


<commit_before>// Copyright 2019 Authors of Cilium
//
// 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 probes

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/command/exec"
	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
)

var (
	log          = logging.DefaultLogger.WithField(logfields.LogSubsys, "probes")
	once         sync.Once
	probeManager *ProbeManager
)

// KernelParam is a type based on string which represents CONFIG_* kernel
// parameters which usually have values "y", "n" or "m".
type KernelParam string

// Enabled checks whether the kernel parameter is enabled.
func (kp KernelParam) Enabled() bool {
	return kp == "y"
}

// Module checks whether the kernel parameter is enabled as a module.
func (kp KernelParam) Module() bool {
	return kp == "m"
}

// SystemConfig contains kernel configuration and sysctl parameters related to
// BPF functionality.
type SystemConfig struct {
	UnprivilegedBpfDisabled      int         `json:"unprivileged_bpf_disabled"`
	BpfJitEnable                 int         `json:"bpf_jit_enable"`
	BpfJitHarden                 int         `json:"bpf_jit_harden"`
	BpfJitKallsyms               int         `json:"bpf_jit_kallsyms"`
	BpfJitLimit                  int         `json:"bpf_jit_limit"`
	ConfigBpf                    KernelParam `json:"CONFIG_BPF"`
	ConfigBpfSyscall             KernelParam `json:"CONFIG_BPF_SYSCALL"`
	ConfigHaveEbpfJit            KernelParam `json:"CONFIG_HAVE_EBPF_JIT"`
	ConfigBpfJit                 KernelParam `json:"CONFIG_BPF_JIT"`
	ConfigBpfJitAlwaysOn         KernelParam `json:"CONFIG_BPF_JIT_ALWAYS_ON"`
	ConfigCgroups                KernelParam `json:"CONFIG_CGROUPS"`
	ConfigCgroupBpf              KernelParam `json:"CONFIG_CGROUP_BPF"`
	ConfigCgroupNetClassID       KernelParam `json:"CONFIG_CGROUP_NET_CLASSID"`
	ConfigSockCgroupData         KernelParam `json:"CONFIG_SOCK_CGROUP_DATA"`
	ConfigBpfEvents              KernelParam `json:"CONFIG_BPF_EVENTS"`
	ConfigKprobeEvents           KernelParam `json:"CONFIG_KPROBE_EVENTS"`
	ConfigUprobeEvents           KernelParam `json:"CONFIG_UPROBE_EVENTS"`
	ConfigTracing                KernelParam `json:"CONFIG_TRACING"`
	ConfigFtraceSyscalls         KernelParam `json:"CONFIG_FTRACE_SYSCALLS"`
	ConfigFunctionErrorInjection KernelParam `json:"CONFIG_FUNCTION_ERROR_INJECTION"`
	ConfigBpfKprobeOverride      KernelParam `json:"CONFIG_BPF_KPROBE_OVERRIDE"`
	ConfigNet                    KernelParam `json:"CONFIG_NET"`
	ConfigXdpSockets             KernelParam `json:"CONFIG_XDP_SOCKETS"`
	ConfigLwtunnelBpf            KernelParam `json:"CONFIG_LWTUNNEL_BPF"`
	ConfigNetActBpf              KernelParam `json:"CONFIG_NET_ACT_BPF"`
	ConfigNetClsBpf              KernelParam `json:"CONFIG_NET_CLS_BPF"`
	ConfigNetClsAct              KernelParam `json:"CONFIG_NET_CLS_ACT"`
	ConfigNetSchIngress          KernelParam `json:"CONFIG_NET_SCH_INGRESS"`
	ConfigXfrm                   KernelParam `json:"CONFIG_XFRM"`
	ConfigIPRouteClassID         KernelParam `json:"CONFIG_IP_ROUTE_CLASSID"`
	ConfigIPv6Seg6Bpf            KernelParam `json:"CONFIG_IPV6_SEG6_BPF"`
	ConfigBpfLircMode2           KernelParam `json:"CONFIG_BPF_LIRC_MODE2"`
	ConfigBpfStreamParser        KernelParam `json:"CONFIG_BPF_STREAM_PARSER"`
	ConfigNetfilterXtMatchBpf    KernelParam `json:"CONFIG_NETFILTER_XT_MATCH_BPF"`
	ConfigBpfilter               KernelParam `json:"CONFIG_BPFILTER"`
	ConfigBpfilterUmh            KernelParam `json:"CONFIG_BPFILTER_UMH"`
	ConfigTestBpf                KernelParam `json:"CONFIG_TEST_BPF"`
}

// MapTypes contains bools indicating which types of BPF maps the currently
// running kernel supports.
type MapTypes struct {
	HaveHashMapType                bool `json:"have_hash_map_type"`
	HaveArrayMapType               bool `json:"have_array_map_type"`
	HaveProgArrayMapType           bool `json:"have_prog_array_map_type"`
	HavePerfEventArrayMapType      bool `json:"have_perf_event_array_map_type"`
	HavePercpuHashMapType          bool `json:"have_percpu_hash_map_type"`
	HavePercpuArrayMapType         bool `json:"have_percpu_array_map_type"`
	HaveStackTraceMapType          bool `json:"have_stack_trace_map_type"`
	HaveCgroupArrayMapType         bool `json:"have_cgroup_array_map_type"`
	HaveLruHashMapType             bool `json:"have_lru_hash_map_type"`
	HaveLruPercpuHashMapType       bool `json:"have_lru_percpu_hash_map_type"`
	HaveLpmTrieMapType             bool `json:"have_lpm_trie_map_type"`
	HaveArrayOfMapsMapType         bool `json:"have_array_of_maps_map_type"`
	HaveHashOfMapsMapType          bool `json:"have_hash_of_maps_map_type"`
	HaveDevmapMapType              bool `json:"have_devmap_map_type"`
	HaveSockmapMapType             bool `json:"have_sockmap_map_type"`
	HaveCpumapMapType              bool `json:"have_cpumap_map_type"`
	HaveXskmapMapType              bool `json:"have_xskmap_map_type"`
	HaveSockhashMapType            bool `json:"have_sockhash_map_type"`
	HaveCgroupStorageMapType       bool `json:"have_cgroup_storage_map_type"`
	HaveReuseportSockarrayMapType  bool `json:"have_reuseport_sockarray_map_type"`
	HavePercpuCgroupStorageMapType bool `json:"have_percpu_cgroup_storage_map_type"`
	HaveQueueMapType               bool `json:"have_queue_map_type"`
	HaveStackMapType               bool `json:"have_stack_map_type"`
}

// Features contains BPF feature checks returned by bpftool.
type Features struct {
	SystemConfig `json:"system_config"`
	MapTypes     `json:"map_types"`
	Helpers      map[string][]string `json:"helpers"`
}

// ProbeManager is a manager of BPF feature checks.
type ProbeManager struct {
	features Features
}

// NewProbeManager returns a new instance of ProbeManager - a manager of BPF
// feature checks.
func NewProbeManager() *ProbeManager {
	newProbeManager := func() {
		var features Features
		out, err := exec.WithTimeout(
			defaults.ExecTimeout, "bpftool", "-j", "feature",
			"probe", "filter_out",
			"\\(trace\\|write_user\\)",
		).CombinedOutput(log, true)
		if err != nil {
			log.WithError(err).Fatal("could not run bpftool")
		}
		if err := json.Unmarshal(out, &features); err != nil {
			log.WithError(err).Fatal("could not parse bpftool output")
		}
		probeManager = &ProbeManager{features: features}
	}
	once.Do(newProbeManager)
	return probeManager
}

// SystemConfigProbes performs a check of kernel configuration parameters. It
// returns an error when parameters required by Cilium are not enabled. It logs
// warnings when optional parameters are not enabled.
func (p *ProbeManager) SystemConfigProbes() error {
	config := p.features.SystemConfig
	// Required
	if !config.ConfigBpf.Enabled() {
		return fmt.Errorf("CONFIG_BPF kernel parameter is required")
	}
	if !config.ConfigBpfSyscall.Enabled() {
		return fmt.Errorf(
			"CONFIG_BPF_SYSCALL kernel parameter is required")
	}
	if !config.ConfigNetSchIngress.Enabled() && !config.ConfigNetSchIngress.Module() {
		return fmt.Errorf(
			"CONFIG_NET_SCH_INGRESS kernel parameter (or module) is required")
	}
	if !config.ConfigNetClsBpf.Enabled() && !config.ConfigNetClsBpf.Module() {
		return fmt.Errorf(
			"CONFIG_NET_CLS_BPF kernel parameter (or module) is required")
	}
	if !config.ConfigNetClsAct.Enabled() {
		return fmt.Errorf(
			"CONFIG_NET_CLS_ACT kernel parameter is required")
	}
	if !config.ConfigBpfJit.Enabled() {
		return fmt.Errorf(
			"CONFIG_BPF_JIT kernel parameter is required")
	}
	if !config.ConfigHaveEbpfJit.Enabled() {
		return fmt.Errorf(
			"CONFIG_HAVE_EBPF_JIT kernel parameter is required")
	}
	// Optional
	if !config.ConfigCgroupBpf.Enabled() {
		log.Warning(
			"CONFIG_CGROUP_BPF optional kernel parameter is not in kernel configuration")
	}
	if !config.ConfigLwtunnelBpf.Enabled() {
		log.Warning(
			"CONFIG_LWTUNNEL_BPF optional kernel parameter is not in kernel configuration")
	}
	if !config.ConfigBpfEvents.Enabled() {
		log.Warning(
			"CONFIG_BPF_EVENTS optional kernel parameter is not in kernel configuration")
	}
	return nil
}

// GetMapTypes returns information about supported BPF map types.
func (p *ProbeManager) GetMapTypes() *MapTypes {
	return &p.features.MapTypes
}

// GetHelpers returns information about available BPF helpers for the given
// program type.
// If program type is not found, returns nil.
func (p *ProbeManager) GetHelpers(prog string) map[string]struct{} {
	for p, helpers := range p.features.Helpers {
		if prog+"_available_helpers" == p {
			ret := map[string]struct{}{}
			for _, h := range helpers {
				ret[h] = struct{}{}
			}
			return ret
		}
	}
	return nil
}
<commit_msg>datapath: Switch to upstream bpftool, remove additional arguments<commit_after>// Copyright 2019 Authors of Cilium
//
// 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 probes

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/command/exec"
	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
)

var (
	log          = logging.DefaultLogger.WithField(logfields.LogSubsys, "probes")
	once         sync.Once
	probeManager *ProbeManager
)

// KernelParam is a type based on string which represents CONFIG_* kernel
// parameters which usually have values "y", "n" or "m".
type KernelParam string

// Enabled checks whether the kernel parameter is enabled.
func (kp KernelParam) Enabled() bool {
	return kp == "y"
}

// Module checks whether the kernel parameter is enabled as a module.
func (kp KernelParam) Module() bool {
	return kp == "m"
}

// SystemConfig contains kernel configuration and sysctl parameters related to
// BPF functionality.
type SystemConfig struct {
	UnprivilegedBpfDisabled      int         `json:"unprivileged_bpf_disabled"`
	BpfJitEnable                 int         `json:"bpf_jit_enable"`
	BpfJitHarden                 int         `json:"bpf_jit_harden"`
	BpfJitKallsyms               int         `json:"bpf_jit_kallsyms"`
	BpfJitLimit                  int         `json:"bpf_jit_limit"`
	ConfigBpf                    KernelParam `json:"CONFIG_BPF"`
	ConfigBpfSyscall             KernelParam `json:"CONFIG_BPF_SYSCALL"`
	ConfigHaveEbpfJit            KernelParam `json:"CONFIG_HAVE_EBPF_JIT"`
	ConfigBpfJit                 KernelParam `json:"CONFIG_BPF_JIT"`
	ConfigBpfJitAlwaysOn         KernelParam `json:"CONFIG_BPF_JIT_ALWAYS_ON"`
	ConfigCgroups                KernelParam `json:"CONFIG_CGROUPS"`
	ConfigCgroupBpf              KernelParam `json:"CONFIG_CGROUP_BPF"`
	ConfigCgroupNetClassID       KernelParam `json:"CONFIG_CGROUP_NET_CLASSID"`
	ConfigSockCgroupData         KernelParam `json:"CONFIG_SOCK_CGROUP_DATA"`
	ConfigBpfEvents              KernelParam `json:"CONFIG_BPF_EVENTS"`
	ConfigKprobeEvents           KernelParam `json:"CONFIG_KPROBE_EVENTS"`
	ConfigUprobeEvents           KernelParam `json:"CONFIG_UPROBE_EVENTS"`
	ConfigTracing                KernelParam `json:"CONFIG_TRACING"`
	ConfigFtraceSyscalls         KernelParam `json:"CONFIG_FTRACE_SYSCALLS"`
	ConfigFunctionErrorInjection KernelParam `json:"CONFIG_FUNCTION_ERROR_INJECTION"`
	ConfigBpfKprobeOverride      KernelParam `json:"CONFIG_BPF_KPROBE_OVERRIDE"`
	ConfigNet                    KernelParam `json:"CONFIG_NET"`
	ConfigXdpSockets             KernelParam `json:"CONFIG_XDP_SOCKETS"`
	ConfigLwtunnelBpf            KernelParam `json:"CONFIG_LWTUNNEL_BPF"`
	ConfigNetActBpf              KernelParam `json:"CONFIG_NET_ACT_BPF"`
	ConfigNetClsBpf              KernelParam `json:"CONFIG_NET_CLS_BPF"`
	ConfigNetClsAct              KernelParam `json:"CONFIG_NET_CLS_ACT"`
	ConfigNetSchIngress          KernelParam `json:"CONFIG_NET_SCH_INGRESS"`
	ConfigXfrm                   KernelParam `json:"CONFIG_XFRM"`
	ConfigIPRouteClassID         KernelParam `json:"CONFIG_IP_ROUTE_CLASSID"`
	ConfigIPv6Seg6Bpf            KernelParam `json:"CONFIG_IPV6_SEG6_BPF"`
	ConfigBpfLircMode2           KernelParam `json:"CONFIG_BPF_LIRC_MODE2"`
	ConfigBpfStreamParser        KernelParam `json:"CONFIG_BPF_STREAM_PARSER"`
	ConfigNetfilterXtMatchBpf    KernelParam `json:"CONFIG_NETFILTER_XT_MATCH_BPF"`
	ConfigBpfilter               KernelParam `json:"CONFIG_BPFILTER"`
	ConfigBpfilterUmh            KernelParam `json:"CONFIG_BPFILTER_UMH"`
	ConfigTestBpf                KernelParam `json:"CONFIG_TEST_BPF"`
}

// MapTypes contains bools indicating which types of BPF maps the currently
// running kernel supports.
type MapTypes struct {
	HaveHashMapType                bool `json:"have_hash_map_type"`
	HaveArrayMapType               bool `json:"have_array_map_type"`
	HaveProgArrayMapType           bool `json:"have_prog_array_map_type"`
	HavePerfEventArrayMapType      bool `json:"have_perf_event_array_map_type"`
	HavePercpuHashMapType          bool `json:"have_percpu_hash_map_type"`
	HavePercpuArrayMapType         bool `json:"have_percpu_array_map_type"`
	HaveStackTraceMapType          bool `json:"have_stack_trace_map_type"`
	HaveCgroupArrayMapType         bool `json:"have_cgroup_array_map_type"`
	HaveLruHashMapType             bool `json:"have_lru_hash_map_type"`
	HaveLruPercpuHashMapType       bool `json:"have_lru_percpu_hash_map_type"`
	HaveLpmTrieMapType             bool `json:"have_lpm_trie_map_type"`
	HaveArrayOfMapsMapType         bool `json:"have_array_of_maps_map_type"`
	HaveHashOfMapsMapType          bool `json:"have_hash_of_maps_map_type"`
	HaveDevmapMapType              bool `json:"have_devmap_map_type"`
	HaveSockmapMapType             bool `json:"have_sockmap_map_type"`
	HaveCpumapMapType              bool `json:"have_cpumap_map_type"`
	HaveXskmapMapType              bool `json:"have_xskmap_map_type"`
	HaveSockhashMapType            bool `json:"have_sockhash_map_type"`
	HaveCgroupStorageMapType       bool `json:"have_cgroup_storage_map_type"`
	HaveReuseportSockarrayMapType  bool `json:"have_reuseport_sockarray_map_type"`
	HavePercpuCgroupStorageMapType bool `json:"have_percpu_cgroup_storage_map_type"`
	HaveQueueMapType               bool `json:"have_queue_map_type"`
	HaveStackMapType               bool `json:"have_stack_map_type"`
}

// Features contains BPF feature checks returned by bpftool.
type Features struct {
	SystemConfig `json:"system_config"`
	MapTypes     `json:"map_types"`
	Helpers      map[string][]string `json:"helpers"`
}

// ProbeManager is a manager of BPF feature checks.
type ProbeManager struct {
	features Features
}

// NewProbeManager returns a new instance of ProbeManager - a manager of BPF
// feature checks.
func NewProbeManager() *ProbeManager {
	newProbeManager := func() {
		var features Features
		out, err := exec.WithTimeout(
			defaults.ExecTimeout,
			"bpftool", "-j", "feature", "probe",
		).CombinedOutput(log, true)
		if err != nil {
			log.WithError(err).Fatal("could not run bpftool")
		}
		if err := json.Unmarshal(out, &features); err != nil {
			log.WithError(err).Fatal("could not parse bpftool output")
		}
		probeManager = &ProbeManager{features: features}
	}
	once.Do(newProbeManager)
	return probeManager
}

// SystemConfigProbes performs a check of kernel configuration parameters. It
// returns an error when parameters required by Cilium are not enabled. It logs
// warnings when optional parameters are not enabled.
func (p *ProbeManager) SystemConfigProbes() error {
	config := p.features.SystemConfig
	// Required
	if !config.ConfigBpf.Enabled() {
		return fmt.Errorf("CONFIG_BPF kernel parameter is required")
	}
	if !config.ConfigBpfSyscall.Enabled() {
		return fmt.Errorf(
			"CONFIG_BPF_SYSCALL kernel parameter is required")
	}
	if !config.ConfigNetSchIngress.Enabled() && !config.ConfigNetSchIngress.Module() {
		return fmt.Errorf(
			"CONFIG_NET_SCH_INGRESS kernel parameter (or module) is required")
	}
	if !config.ConfigNetClsBpf.Enabled() && !config.ConfigNetClsBpf.Module() {
		return fmt.Errorf(
			"CONFIG_NET_CLS_BPF kernel parameter (or module) is required")
	}
	if !config.ConfigNetClsAct.Enabled() {
		return fmt.Errorf(
			"CONFIG_NET_CLS_ACT kernel parameter is required")
	}
	if !config.ConfigBpfJit.Enabled() {
		return fmt.Errorf(
			"CONFIG_BPF_JIT kernel parameter is required")
	}
	if !config.ConfigHaveEbpfJit.Enabled() {
		return fmt.Errorf(
			"CONFIG_HAVE_EBPF_JIT kernel parameter is required")
	}
	// Optional
	if !config.ConfigCgroupBpf.Enabled() {
		log.Warning(
			"CONFIG_CGROUP_BPF optional kernel parameter is not in kernel configuration")
	}
	if !config.ConfigLwtunnelBpf.Enabled() {
		log.Warning(
			"CONFIG_LWTUNNEL_BPF optional kernel parameter is not in kernel configuration")
	}
	if !config.ConfigBpfEvents.Enabled() {
		log.Warning(
			"CONFIG_BPF_EVENTS optional kernel parameter is not in kernel configuration")
	}
	return nil
}

// GetMapTypes returns information about supported BPF map types.
func (p *ProbeManager) GetMapTypes() *MapTypes {
	return &p.features.MapTypes
}

// GetHelpers returns information about available BPF helpers for the given
// program type.
// If program type is not found, returns nil.
func (p *ProbeManager) GetHelpers(prog string) map[string]struct{} {
	for p, helpers := range p.features.Helpers {
		if prog+"_available_helpers" == p {
			ret := map[string]struct{}{}
			for _, h := range helpers {
				ret[h] = struct{}{}
			}
			return ret
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package encryptedstore

import (
	"reflect"

	"github.com/rancher/types/apis/core/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	typedv1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	defaultNamespace = "cattle-system"
)

type GenericEncryptedStore struct {
	prefix    string
	namespace string
	secrets   typedv1.SecretsGetter
}

func NewGenericEncrypedStore(prefix, namespace string, namespaceInterface v1.NamespaceInterface, secrets typedv1.SecretsGetter) (*GenericEncryptedStore, error) {
	if namespace == "" {
		namespace = defaultNamespace
	}

	_, err := namespaceInterface.Get(namespace, metav1.GetOptions{})
	if errors.IsNotFound(err) {
		ns := &corev1.Namespace{}
		ns.Name = namespace
		if _, err := namespaceInterface.Create(ns); err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	return &GenericEncryptedStore{
		prefix:    prefix,
		namespace: namespace,
		secrets:   secrets,
	}, nil
}

func (g *GenericEncryptedStore) Get(name string) (map[string]string, error) {
	sec, err := g.secrets.Secrets(g.namespace).Get(g.getKey(name), metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	result := map[string]string{}
	for k, v := range sec.Data {
		result[k] = string(v)
	}

	return result, nil
}

func (g *GenericEncryptedStore) getKey(name string) string {
	return g.prefix + name
}

func (g *GenericEncryptedStore) Set(name string, data map[string]string) error {
	sec, err := g.secrets.Secrets(g.namespace).Get(g.getKey(name), metav1.GetOptions{})
	if errors.IsNotFound(err) {
		sec = &corev1.Secret{}
		sec.Name = g.getKey(name)
		sec.StringData = data
		_, err := g.secrets.Secrets(g.namespace).Create(sec)
		return err
	} else if err != nil {
		return err
	}

	orig := sec.DeepCopy()
	if sec.Data == nil {
		sec.Data = map[string][]byte{}
	}
	for k, v := range data {
		sec.Data[k] = []byte(v)
	}

	if !reflect.DeepEqual(orig, sec) {
		_, err = g.secrets.Secrets(g.namespace).Update(sec)
	}
	return err
}

func (g *GenericEncryptedStore) Remove(name string) error {
	err := g.secrets.Secrets(g.namespace).Delete(g.getKey(name), nil)
	if errors.IsNotFound(err) {
		return nil
	}
	return err
}
<commit_msg>Retry on secrets save<commit_after>package encryptedstore

import (
	"reflect"

	"github.com/rancher/types/apis/core/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	typedv1 "k8s.io/client-go/kubernetes/typed/core/v1"
)

const (
	defaultNamespace = "cattle-system"
)

type GenericEncryptedStore struct {
	prefix    string
	namespace string
	secrets   typedv1.SecretsGetter
}

func NewGenericEncrypedStore(prefix, namespace string, namespaceInterface v1.NamespaceInterface, secrets typedv1.SecretsGetter) (*GenericEncryptedStore, error) {
	if namespace == "" {
		namespace = defaultNamespace
	}

	_, err := namespaceInterface.Get(namespace, metav1.GetOptions{})
	if errors.IsNotFound(err) {
		ns := &corev1.Namespace{}
		ns.Name = namespace
		if _, err := namespaceInterface.Create(ns); err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	return &GenericEncryptedStore{
		prefix:    prefix,
		namespace: namespace,
		secrets:   secrets,
	}, nil
}

func (g *GenericEncryptedStore) Get(name string) (map[string]string, error) {
	sec, err := g.secrets.Secrets(g.namespace).Get(g.getKey(name), metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	result := map[string]string{}
	for k, v := range sec.Data {
		result[k] = string(v)
	}

	return result, nil
}

func (g *GenericEncryptedStore) getKey(name string) string {
	return g.prefix + name
}

func (g *GenericEncryptedStore) Set(name string, data map[string]string) error {
	return g.set(name, data, 0)
}

func (g *GenericEncryptedStore) set(name string, data map[string]string, try int) error {
	sec, err := g.secrets.Secrets(g.namespace).Get(g.getKey(name), metav1.GetOptions{})
	if errors.IsNotFound(err) {
		sec = &corev1.Secret{}
		sec.Name = g.getKey(name)
		sec.StringData = data
		_, err := g.secrets.Secrets(g.namespace).Create(sec)
		return err
	} else if err != nil {
		return err
	}

	orig := sec.DeepCopy()
	if sec.Data == nil {
		sec.Data = map[string][]byte{}
	}
	for k, v := range data {
		sec.Data[k] = []byte(v)
	}

	if !reflect.DeepEqual(orig, sec) {
		_, err = g.secrets.Secrets(g.namespace).Update(sec)
		if err != nil && try < 5 {
			return g.set(name, data, try+1)
		}
	}
	return err
}

func (g *GenericEncryptedStore) Remove(name string) error {
	err := g.secrets.Secrets(g.namespace).Delete(g.getKey(name), nil)
	if errors.IsNotFound(err) {
		return nil
	}
	return err
}
<|endoftext|>


<commit_before>package positions

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	yaml "gopkg.in/yaml.v2"
)

const positionFileMode = 0700

// Config describes where to get postition information from.
type Config struct {
	SyncPeriod    time.Duration `yaml:"sync_period"`
	PositionsFile string        `yaml:"filename"`
}

// RegisterFlags register flags.
func (cfg *Config) RegisterFlags(flags *flag.FlagSet) {
	flags.DurationVar(&cfg.SyncPeriod, "positions.sync-period", 10*time.Second, "Period with this to sync the position file.")
	flag.StringVar(&cfg.PositionsFile, "positions.file", "/var/log/positions.yaml", "Location to read/write positions from.")
}

// Positions tracks how far through each file we've read.
type Positions struct {
	logger    log.Logger
	cfg       Config
	mtx       sync.Mutex
	positions map[string]string
	quit      chan struct{}
	done      chan struct{}
}

// File format for the positions data.
type File struct {
	Positions map[string]string `yaml:"positions"`
}

// New makes a new Positions.
func New(logger log.Logger, cfg Config) (*Positions, error) {
	positions, err := readPositionsFile(cfg.PositionsFile)
	if err != nil {
		return nil, err
	}

	p := &Positions{
		logger:    logger,
		cfg:       cfg,
		positions: positions,
		quit:      make(chan struct{}),
		done:      make(chan struct{}),
	}

	go p.run()
	return p, nil
}

// Stop the Position tracker.
func (p *Positions) Stop() {
	close(p.quit)
	<-p.done
}

// PutString records (asynchronsouly) how far we've read through a file.
// Unlike Put, it records a string offset and is only useful for
// JournalTargets which doesn't have integer offsets.
func (p *Positions) PutString(path string, pos string) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	p.positions[path] = pos
}

// Put records (asynchronously) how far we've read through a file.
func (p *Positions) Put(path string, pos int64) {
	p.PutString(path, strconv.FormatInt(pos, 10))
}

// GetString returns how far we've through a file as a string.
// JournalTarget writes a journal cursor to the positions file, while
// FileTarget writes an integer offset. Use Get to read the integer
// offset.
func (p *Positions) GetString(path string) string {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	return p.positions[path]
}

// Get returns how far we've read through a file. Returns an error
// if the value stored for the file is not an integer.
func (p *Positions) Get(path string) (int64, error) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	pos, ok := p.positions[path]
	if !ok {
		return 0, nil
	}
	return strconv.ParseInt(pos, 10, 64)
}

// Remove removes the position tracking for a filepath
func (p *Positions) Remove(path string) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	p.remove(path)
}

func (p *Positions) remove(path string) {
	delete(p.positions, path)
}

// SyncPeriod returns how often the positions file gets resynced
func (p *Positions) SyncPeriod() time.Duration {
	return p.cfg.SyncPeriod
}

func (p *Positions) run() {
	defer func() {
		p.save()
		level.Debug(p.logger).Log("msg", "positions saved")
		close(p.done)
	}()

	ticker := time.NewTicker(p.cfg.SyncPeriod)
	for {
		select {
		case <-p.quit:
			return
		case <-ticker.C:
			p.save()
			p.cleanup()
		}
	}
}

func (p *Positions) save() {
	p.mtx.Lock()
	positions := make(map[string]string, len(p.positions))
	for k, v := range p.positions {
		positions[k] = v
	}
	p.mtx.Unlock()

	if err := writePositionFile(p.cfg.PositionsFile, positions); err != nil {
		level.Error(p.logger).Log("msg", "error writing positions file", "error", err)
	}
}

func (p *Positions) cleanup() {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	toRemove := []string{}
	for k := range p.positions {
		// If the position file is prefixed with journal, it's a
		// JournalTarget cursor and not a file on disk.
		if strings.HasPrefix(k, "journal-") {
			continue
		}

		if _, err := os.Stat(k); err != nil {
			if os.IsNotExist(err) {
				// File no longer exists.
				toRemove = append(toRemove, k)
			} else {
				// Can't determine if file exists or not, some other error.
				level.Warn(p.logger).Log("msg", "could not determine if log file "+
					"still exists while cleaning positions file", "error", err)
			}
		}
	}
	for _, tr := range toRemove {
		p.remove(tr)
	}
}

func readPositionsFile(filename string) (map[string]string, error) {
	buf, err := ioutil.ReadFile(filepath.Clean(filename))
	if err != nil {
		if os.IsNotExist(err) {
			return map[string]string{}, nil
		}
		return nil, err
	}

	var p File
	if err := yaml.UnmarshalStrict(buf, &p); err != nil {
		return nil, err
	}

	return p.Positions, nil
}

func writePositionFile(filename string, positions map[string]string) error {
	buf, err := yaml.Marshal(File{
		Positions: positions,
	})
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filepath.Clean(filename), buf, os.FileMode(positionFileMode))
}
<commit_msg>pkg/promtail/positions: remove executable bit from positions file (#893)<commit_after>package positions

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	yaml "gopkg.in/yaml.v2"
)

const positionFileMode = 0600

// Config describes where to get postition information from.
type Config struct {
	SyncPeriod    time.Duration `yaml:"sync_period"`
	PositionsFile string        `yaml:"filename"`
}

// RegisterFlags register flags.
func (cfg *Config) RegisterFlags(flags *flag.FlagSet) {
	flags.DurationVar(&cfg.SyncPeriod, "positions.sync-period", 10*time.Second, "Period with this to sync the position file.")
	flag.StringVar(&cfg.PositionsFile, "positions.file", "/var/log/positions.yaml", "Location to read/write positions from.")
}

// Positions tracks how far through each file we've read.
type Positions struct {
	logger    log.Logger
	cfg       Config
	mtx       sync.Mutex
	positions map[string]string
	quit      chan struct{}
	done      chan struct{}
}

// File format for the positions data.
type File struct {
	Positions map[string]string `yaml:"positions"`
}

// New makes a new Positions.
func New(logger log.Logger, cfg Config) (*Positions, error) {
	positions, err := readPositionsFile(cfg.PositionsFile)
	if err != nil {
		return nil, err
	}

	p := &Positions{
		logger:    logger,
		cfg:       cfg,
		positions: positions,
		quit:      make(chan struct{}),
		done:      make(chan struct{}),
	}

	go p.run()
	return p, nil
}

// Stop the Position tracker.
func (p *Positions) Stop() {
	close(p.quit)
	<-p.done
}

// PutString records (asynchronsouly) how far we've read through a file.
// Unlike Put, it records a string offset and is only useful for
// JournalTargets which doesn't have integer offsets.
func (p *Positions) PutString(path string, pos string) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	p.positions[path] = pos
}

// Put records (asynchronously) how far we've read through a file.
func (p *Positions) Put(path string, pos int64) {
	p.PutString(path, strconv.FormatInt(pos, 10))
}

// GetString returns how far we've through a file as a string.
// JournalTarget writes a journal cursor to the positions file, while
// FileTarget writes an integer offset. Use Get to read the integer
// offset.
func (p *Positions) GetString(path string) string {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	return p.positions[path]
}

// Get returns how far we've read through a file. Returns an error
// if the value stored for the file is not an integer.
func (p *Positions) Get(path string) (int64, error) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	pos, ok := p.positions[path]
	if !ok {
		return 0, nil
	}
	return strconv.ParseInt(pos, 10, 64)
}

// Remove removes the position tracking for a filepath
func (p *Positions) Remove(path string) {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	p.remove(path)
}

func (p *Positions) remove(path string) {
	delete(p.positions, path)
}

// SyncPeriod returns how often the positions file gets resynced
func (p *Positions) SyncPeriod() time.Duration {
	return p.cfg.SyncPeriod
}

func (p *Positions) run() {
	defer func() {
		p.save()
		level.Debug(p.logger).Log("msg", "positions saved")
		close(p.done)
	}()

	ticker := time.NewTicker(p.cfg.SyncPeriod)
	for {
		select {
		case <-p.quit:
			return
		case <-ticker.C:
			p.save()
			p.cleanup()
		}
	}
}

func (p *Positions) save() {
	p.mtx.Lock()
	positions := make(map[string]string, len(p.positions))
	for k, v := range p.positions {
		positions[k] = v
	}
	p.mtx.Unlock()

	if err := writePositionFile(p.cfg.PositionsFile, positions); err != nil {
		level.Error(p.logger).Log("msg", "error writing positions file", "error", err)
	}
}

func (p *Positions) cleanup() {
	p.mtx.Lock()
	defer p.mtx.Unlock()
	toRemove := []string{}
	for k := range p.positions {
		// If the position file is prefixed with journal, it's a
		// JournalTarget cursor and not a file on disk.
		if strings.HasPrefix(k, "journal-") {
			continue
		}

		if _, err := os.Stat(k); err != nil {
			if os.IsNotExist(err) {
				// File no longer exists.
				toRemove = append(toRemove, k)
			} else {
				// Can't determine if file exists or not, some other error.
				level.Warn(p.logger).Log("msg", "could not determine if log file "+
					"still exists while cleaning positions file", "error", err)
			}
		}
	}
	for _, tr := range toRemove {
		p.remove(tr)
	}
}

func readPositionsFile(filename string) (map[string]string, error) {
	buf, err := ioutil.ReadFile(filepath.Clean(filename))
	if err != nil {
		if os.IsNotExist(err) {
			return map[string]string{}, nil
		}
		return nil, err
	}

	var p File
	if err := yaml.UnmarshalStrict(buf, &p); err != nil {
		return nil, err
	}

	return p.Positions, nil
}

func writePositionFile(filename string, positions map[string]string) error {
	buf, err := yaml.Marshal(File{
		Positions: positions,
	})
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filepath.Clean(filename), buf, os.FileMode(positionFileMode))
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"context"
	"testing"

	"github.com/grafana/grafana/pkg/models"
	. "github.com/smartystreets/goconvey/convey"
)

func TestStatsDataAccess(t *testing.T) {
	Convey("Testing Stats Data Access", t, func() {
		InitTestDB(t)

		Convey("Get system stats should not results in error", func() {
			query := models.GetSystemStatsQuery{}
			err := GetSystemStats(&query)
			So(err, ShouldBeNil)
		})

		Convey("Get system user count stats should not results in error", func() {
			query := models.GetSystemUserCountStatsQuery{}
			err := GetSystemUserCountStats(context.Background(), &query)
			So(err, ShouldBeNil)
		})

		Convey("Get datasource stats should not results in error", func() {
			query := models.GetDataSourceStatsQuery{}
			err := GetDataSourceStats(&query)
			So(err, ShouldBeNil)
		})

		Convey("Get datasource access stats should not results in error", func() {
			query := models.GetDataSourceAccessStatsQuery{}
			err := GetDataSourceAccessStats(&query)
			So(err, ShouldBeNil)
		})

		Convey("Get alert notifier stats should not results in error", func() {
			query := models.GetAlertNotifierUsageStatsQuery{}
			err := GetAlertNotifiersUsageStats(context.Background(), &query)
			So(err, ShouldBeNil)
		})

		Convey("Get admin stats should not result in error", func() {
			query := models.GetAdminStatsQuery{}
			err := GetAdminStats(&query)
			So(err, ShouldBeNil)
		})
	})
}
<commit_msg>SQLStore: Test admins/editors/viewers stats validity (#20751)<commit_after>package sqlstore

import (
	"context"
	"fmt"
	"testing"

	"github.com/grafana/grafana/pkg/models"
	"github.com/stretchr/testify/assert"
)

func TestStatsDataAccess(t *testing.T) {
	t.Run("Testing Stats Data Access", func(t *testing.T) {
		InitTestDB(t)

		t.Run("Get system stats should not results in error", func(t *testing.T) {
			populateDB(t)

			query := models.GetSystemStatsQuery{}
			err := GetSystemStats(&query)
			assert.Nil(t, err)
			assert.Equal(t, query.Result.Users, int64(3))
			assert.Equal(t, query.Result.Editors, 1)
			assert.Equal(t, query.Result.Viewers, 1)
			assert.Equal(t, query.Result.Admins, 3)
		})

		t.Run("Get system user count stats should not results in error", func(t *testing.T) {
			query := models.GetSystemUserCountStatsQuery{}
			err := GetSystemUserCountStats(context.Background(), &query)
			assert.Nil(t, err)
		})

		t.Run("Get datasource stats should not results in error", func(t *testing.T) {
			query := models.GetDataSourceStatsQuery{}
			err := GetDataSourceStats(&query)
			assert.Nil(t, err)
		})

		t.Run("Get datasource access stats should not results in error", func(t *testing.T) {
			query := models.GetDataSourceAccessStatsQuery{}
			err := GetDataSourceAccessStats(&query)
			assert.Nil(t, err)
		})

		t.Run("Get alert notifier stats should not results in error", func(t *testing.T) {
			query := models.GetAlertNotifierUsageStatsQuery{}
			err := GetAlertNotifiersUsageStats(context.Background(), &query)
			assert.Nil(t, err)
		})

		t.Run("Get admin stats should not result in error", func(t *testing.T) {
			query := models.GetAdminStatsQuery{}
			err := GetAdminStats(&query)
			assert.Nil(t, err)
		})
	})
}

func populateDB(t *testing.T) {
	users := make([]models.User, 3)
	for i := range users {
		cmd := &models.CreateUserCommand{
			Email:   fmt.Sprintf("usertest%v@test.com", i),
			Name:    fmt.Sprintf("user name %v", i),
			Login:   fmt.Sprintf("user_test_%v_login", i),
			OrgName: fmt.Sprintf("Org #%v", i),
		}
		err := CreateUser(context.Background(), cmd)
		assert.Nil(t, err)
		users[i] = cmd.Result
	}

	// get 1st user's organisation
	getOrgByIdQuery := &models.GetOrgByIdQuery{Id: users[0].OrgId}
	err := GetOrgById(getOrgByIdQuery)
	assert.Nil(t, err)
	org := getOrgByIdQuery.Result

	// add 2nd user as editor
	cmd := &models.AddOrgUserCommand{
		OrgId:  org.Id,
		UserId: users[1].Id,
		Role:   models.ROLE_EDITOR,
	}
	err = AddOrgUser(cmd)
	assert.Nil(t, err)

	// add 3rd user as viewer
	cmd = &models.AddOrgUserCommand{
		OrgId:  org.Id,
		UserId: users[2].Id,
		Role:   models.ROLE_VIEWER,
	}
	err = AddOrgUser(cmd)
	assert.Nil(t, err)

	// get 2nd user's organisation
	getOrgByIdQuery = &models.GetOrgByIdQuery{Id: users[1].OrgId}
	err = GetOrgById(getOrgByIdQuery)
	assert.Nil(t, err)
	org = getOrgByIdQuery.Result

	// add 1st user as admin
	cmd = &models.AddOrgUserCommand{
		OrgId:  org.Id,
		UserId: users[0].Id,
		Role:   models.ROLE_ADMIN,
	}
	err = AddOrgUser(cmd)
	assert.Nil(t, err)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 docker

import (
	"archive/tar"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/GoogleContainerTools/skaffold/testutil"
)

func TestDockerContext(t *testing.T) {
	tmpDir, cleanup := testutil.TempDir(t)
	defer cleanup()

	os.Mkdir(filepath.Join(tmpDir, "files"), 0750)
	ioutil.WriteFile(filepath.Join(tmpDir, "files", "ignored.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "files", "included.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, ".dockerignore"), []byte("**/ignored.txt\nalsoignored.txt"), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "Dockerfile"), []byte("FROM alpine\nCOPY ./files /files"), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "ignored.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "alsoignored.txt"), []byte(""), 0644)

	reader, writer := io.Pipe()
	go func() {
		err := CreateDockerTarContext(writer, "Dockerfile", tmpDir)
		if err != nil {
			writer.CloseWithError(err)
		} else {
			writer.Close()
		}
	}()

	files := make(map[string]bool)
	tr := tar.NewReader(reader)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		files[header.Name] = true
	}

	if files["ignored.txt"] {
		t.Error("File ignored.txt should have been excluded, but was not")
	}
	if files["alsoignored.txt"] {
		t.Error("File alsoignored.txt should have been excluded, but was not")
	}
	if files["files/ignored.txt"] {
		t.Error("File files/ignored.txt should have been excluded, but was not")
	}
	if !files["files/included.txt"] {
		t.Error("File files/included.txt should have been included, but was not")
	}
	if !files["Dockerfile"] {
		t.Error("File Dockerfile should have been included, but was not")
	}
}
<commit_msg>Fix flaky test<commit_after>/*
Copyright 2018 Google LLC

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 docker

import (
	"archive/tar"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/GoogleContainerTools/skaffold/testutil"
)

func TestDockerContext(t *testing.T) {
	tmpDir, cleanup := testutil.TempDir(t)
	defer cleanup()

	RetrieveImage = mockRetrieveImage
	defer func() {
		RetrieveImage = retrieveImage
	}()

	os.Mkdir(filepath.Join(tmpDir, "files"), 0750)
	ioutil.WriteFile(filepath.Join(tmpDir, "files", "ignored.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "files", "included.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, ".dockerignore"), []byte("**/ignored.txt\nalsoignored.txt"), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "Dockerfile"), []byte("FROM alpine\nCOPY ./files /files"), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "ignored.txt"), []byte(""), 0644)
	ioutil.WriteFile(filepath.Join(tmpDir, "alsoignored.txt"), []byte(""), 0644)

	reader, writer := io.Pipe()
	go func() {
		err := CreateDockerTarContext(writer, "Dockerfile", tmpDir)
		if err != nil {
			writer.CloseWithError(err)
		} else {
			writer.Close()
		}
	}()

	files := make(map[string]bool)
	tr := tar.NewReader(reader)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		files[header.Name] = true
	}

	if files["ignored.txt"] {
		t.Error("File ignored.txt should have been excluded, but was not")
	}
	if files["alsoignored.txt"] {
		t.Error("File alsoignored.txt should have been excluded, but was not")
	}
	if files["files/ignored.txt"] {
		t.Error("File files/ignored.txt should have been excluded, but was not")
	}
	if !files["files/included.txt"] {
		t.Error("File files/included.txt should have been included, but was not")
	}
	if !files["Dockerfile"] {
		t.Error("File Dockerfile should have been included, but was not")
	}
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright 2018-2019 Authors of Cilium

package testidentity

import (
	"context"
	"net"

	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/labels"
)

type IdentityAllocatorOwnerMock struct{}

func (i *IdentityAllocatorOwnerMock) UpdateIdentities(added, deleted cache.IdentityCache) {}

func (i *IdentityAllocatorOwnerMock) GetNodeSuffix() string {
	return "foo"
}

// MockIdentityAllocator is used as a mock identity allocator for unit tests.
type MockIdentityAllocator struct {
	cache.IdentityCache

	currentID        int // Regular identities
	localID          int // CIDR identities
	idRefCount       map[int]int
	ipToIdentity     map[string]int
	idToIdentity     map[int]*identity.Identity
	labelsToIdentity map[string]int // labels are sorted as a key
}

// NewMockIdentityAllocator returns a new mock identity allocator to be used
// for unit testing purposes. It can be used as a drop-in for "real" identity
// allocation in a testing context.
func NewMockIdentityAllocator(c cache.IdentityCache) *MockIdentityAllocator {
	if c == nil {
		c = cache.IdentityCache{}
	}
	return &MockIdentityAllocator{
		IdentityCache: c,

		currentID:        1000,
		localID:          int(identity.LocalIdentityFlag),
		ipToIdentity:     make(map[string]int),
		idRefCount:       make(map[int]int),
		idToIdentity:     make(map[int]*identity.Identity),
		labelsToIdentity: make(map[string]int),
	}
}

// WaitForInitialGlobalIdentities does nothing.
func (f *MockIdentityAllocator) WaitForInitialGlobalIdentities(context.Context) error {
	return nil
}

// IsLocalIdentityAllocatorInitialized returns true.
func (f *MockIdentityAllocator) IsLocalIdentityAllocatorInitialized() bool {
	return true
}

// GetIdentities returns the identities from the identity cache.
func (f *MockIdentityAllocator) GetIdentities() cache.IdentitiesModel {
	result := cache.IdentitiesModel{}
	return result.FromIdentityCache(f.IdentityCache)
}

// AllocateIdentity allocates a fake identity. It is meant to generally mock
// the canonical identity allocator logic.
func (f *MockIdentityAllocator) AllocateIdentity(_ context.Context, lbls labels.Labels, _ bool) (*identity.Identity, bool, error) {
	if reservedIdentity := identity.LookupReservedIdentityByLabels(lbls); reservedIdentity != nil {
		return reservedIdentity, false, nil
	}

	var id int
	if identity.RequiresGlobalIdentity(lbls) {
		id = f.currentID
		f.currentID++
	} else {
		id = f.localID
		f.localID++
	}

	f.IdentityCache[identity.NumericIdentity(id)] = lbls.LabelArray()
	f.labelsToIdentity[lbls.String()] = id

	realID := &identity.Identity{
		ID:     identity.NumericIdentity(id),
		Labels: lbls,
	}
	f.idToIdentity[id] = realID
	f.idRefCount[id]++

	return realID, true, nil
}

// Release releases a fake identity. It is meant to generally mock the
// canonical identity release logic.
func (f *MockIdentityAllocator) Release(_ context.Context, id *identity.Identity, _ bool) (released bool, err error) {
	count, ok := f.idRefCount[int(id.ID)]
	if !ok {
		return false, nil
	}
	if count == 1 {
		delete(f.idRefCount, int(id.ID))
	} else {
		count--
		f.idRefCount[int(id.ID)] = count
		return false, nil
	}
	delete(f.idToIdentity, int(id.ID))
	delete(f.IdentityCache, id.ID)
	return true, nil
}

// ReleaseSlice wraps Release for slices.
func (f *MockIdentityAllocator) ReleaseSlice(ctx context.Context, _ cache.IdentityAllocatorOwner, identities []*identity.Identity) error {
	for _, id := range identities {
		if _, err := f.Release(ctx, id, false); err != nil {
			return err
		}
	}
	return nil
}

// LookupIdentity looks up the labels in the mock identity store.
func (f *MockIdentityAllocator) LookupIdentity(ctx context.Context, lbls labels.Labels) *identity.Identity {
	if reservedIdentity := identity.LookupReservedIdentityByLabels(lbls); reservedIdentity != nil {
		return reservedIdentity
	}
	return f.idToIdentity[f.labelsToIdentity[lbls.String()]]
}

// LookupIdentityByID returns the identity corresponding to the id if the
// identity is a reserved identity. Otherwise, returns nil.
func (f *MockIdentityAllocator) LookupIdentityByID(ctx context.Context, id identity.NumericIdentity) *identity.Identity {
	if identity := identity.LookupReservedIdentity(id); identity != nil {
		return identity
	}
	return f.idToIdentity[int(id)]
}

// AllocateCIDRsForIPs allocates CIDR identities for the given IPs. It is meant
// to generally mock the CIDR identity allocator logic.
func (f *MockIdentityAllocator) AllocateCIDRsForIPs(IPs []net.IP, _ map[string]*identity.Identity) ([]*identity.Identity, error) {
	result := make([]*identity.Identity, 0, len(IPs))
	for _, ip := range IPs {
		id, ok := f.ipToIdentity[ip.String()]
		if !ok {
			id = f.localID
			f.ipToIdentity[ip.String()] = id
			f.idRefCount[id]++
			f.localID = id + 1
		}
		cidrLabels := append([]string{}, ip.String())
		result = append(result, &identity.Identity{
			ID:        identity.NumericIdentity(id),
			CIDRLabel: labels.NewLabelsFromModel(cidrLabels),
		})
	}
	return result, nil
}

func (f *MockIdentityAllocator) ReleaseCIDRIdentitiesByID(ctx context.Context, identities []identity.NumericIdentity) {
}

// GetIdentityCache returns the identity cache.
func (f *MockIdentityAllocator) GetIdentityCache() cache.IdentityCache {
	return f.IdentityCache
}
<commit_msg>testutils/identity: Implement proper refcounting for mock<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright 2018-2019 Authors of Cilium

package testidentity

import (
	"context"
	"net"

	"github.com/cilium/cilium/pkg/identity"
	"github.com/cilium/cilium/pkg/identity/cache"
	"github.com/cilium/cilium/pkg/labels"
)

type IdentityAllocatorOwnerMock struct{}

func (i *IdentityAllocatorOwnerMock) UpdateIdentities(added, deleted cache.IdentityCache) {}

func (i *IdentityAllocatorOwnerMock) GetNodeSuffix() string {
	return "foo"
}

// MockIdentityAllocator is used as a mock identity allocator for unit tests.
type MockIdentityAllocator struct {
	cache.IdentityCache

	currentID        int // Regular identities
	localID          int // CIDR identities
	ipToIdentity     map[string]int
	idToIdentity     map[int]*identity.Identity
	labelsToIdentity map[string]int // labels are sorted as a key
}

// NewMockIdentityAllocator returns a new mock identity allocator to be used
// for unit testing purposes. It can be used as a drop-in for "real" identity
// allocation in a testing context.
func NewMockIdentityAllocator(c cache.IdentityCache) *MockIdentityAllocator {
	if c == nil {
		c = cache.IdentityCache{}
	}
	return &MockIdentityAllocator{
		IdentityCache: c,

		currentID:        1000,
		localID:          int(identity.LocalIdentityFlag),
		ipToIdentity:     make(map[string]int),
		idToIdentity:     make(map[int]*identity.Identity),
		labelsToIdentity: make(map[string]int),
	}
}

// WaitForInitialGlobalIdentities does nothing.
func (f *MockIdentityAllocator) WaitForInitialGlobalIdentities(context.Context) error {
	return nil
}

// IsLocalIdentityAllocatorInitialized returns true.
func (f *MockIdentityAllocator) IsLocalIdentityAllocatorInitialized() bool {
	return true
}

// GetIdentities returns the identities from the identity cache.
func (f *MockIdentityAllocator) GetIdentities() cache.IdentitiesModel {
	result := cache.IdentitiesModel{}
	return result.FromIdentityCache(f.IdentityCache)
}

// AllocateIdentity allocates a fake identity. It is meant to generally mock
// the canonical identity allocator logic.
func (f *MockIdentityAllocator) AllocateIdentity(_ context.Context, lbls labels.Labels, _ bool) (*identity.Identity, bool, error) {
	if reservedIdentity := identity.LookupReservedIdentityByLabels(lbls); reservedIdentity != nil {
		return reservedIdentity, false, nil
	}

	requiresGlobal := identity.RequiresGlobalIdentity(lbls)

	if numID, ok := f.labelsToIdentity[lbls.String()]; ok && !requiresGlobal {
		id := f.idToIdentity[numID]
		id.ReferenceCount++
		return id, false, nil
	}

	var id int
	if requiresGlobal {
		id = f.currentID
		f.currentID++
	} else {
		id = f.localID
		f.localID++
	}

	f.IdentityCache[identity.NumericIdentity(id)] = lbls.LabelArray()
	f.labelsToIdentity[lbls.String()] = id

	realID := &identity.Identity{
		ID:             identity.NumericIdentity(id),
		Labels:         lbls,
		ReferenceCount: 1,
	}
	f.idToIdentity[id] = realID

	return realID, true, nil
}

// Release releases a fake identity. It is meant to generally mock the
// canonical identity release logic.
func (f *MockIdentityAllocator) Release(_ context.Context, id *identity.Identity, _ bool) (released bool, err error) {
	realID, ok := f.idToIdentity[int(id.ID)]
	if !ok {
		return false, nil
	}
	if realID.ReferenceCount == 1 {
		delete(f.idToIdentity, int(id.ID))
		delete(f.IdentityCache, id.ID)
	} else {
		realID.ReferenceCount--
		return false, nil
	}
	return true, nil
}

// ReleaseSlice wraps Release for slices.
func (f *MockIdentityAllocator) ReleaseSlice(ctx context.Context, _ cache.IdentityAllocatorOwner, identities []*identity.Identity) error {
	for _, id := range identities {
		if _, err := f.Release(ctx, id, false); err != nil {
			return err
		}
	}
	return nil
}

// LookupIdentity looks up the labels in the mock identity store.
func (f *MockIdentityAllocator) LookupIdentity(ctx context.Context, lbls labels.Labels) *identity.Identity {
	if reservedIdentity := identity.LookupReservedIdentityByLabels(lbls); reservedIdentity != nil {
		return reservedIdentity
	}
	return f.idToIdentity[f.labelsToIdentity[lbls.String()]]
}

// LookupIdentityByID returns the identity corresponding to the id if the
// identity is a reserved identity. Otherwise, returns nil.
func (f *MockIdentityAllocator) LookupIdentityByID(ctx context.Context, id identity.NumericIdentity) *identity.Identity {
	if identity := identity.LookupReservedIdentity(id); identity != nil {
		return identity
	}
	return f.idToIdentity[int(id)]
}

// AllocateCIDRsForIPs allocates CIDR identities for the given IPs. It is meant
// to generally mock the CIDR identity allocator logic.
func (f *MockIdentityAllocator) AllocateCIDRsForIPs(IPs []net.IP, _ map[string]*identity.Identity) ([]*identity.Identity, error) {
	result := make([]*identity.Identity, 0, len(IPs))
	for _, ip := range IPs {
		id, ok := f.ipToIdentity[ip.String()]
		if !ok {
			id = f.localID
			f.ipToIdentity[ip.String()] = id
			f.localID = id + 1
		}
		cidrLabels := append([]string{}, ip.String())
		result = append(result, &identity.Identity{
			ID:        identity.NumericIdentity(id),
			CIDRLabel: labels.NewLabelsFromModel(cidrLabels),
		})
	}
	return result, nil
}

func (f *MockIdentityAllocator) ReleaseCIDRIdentitiesByID(ctx context.Context, identities []identity.NumericIdentity) {
}

// GetIdentityCache returns the identity cache.
func (f *MockIdentityAllocator) GetIdentityCache() cache.IdentityCache {
	return f.IdentityCache
}
<|endoftext|>


<commit_before>package konnectors

import (
	"archive/tar"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/sirupsen/logrus"

	"github.com/spf13/afero"
)

func init() {
	jobs.AddWorker("konnector", &jobs.WorkerConfig{
		Concurrency:  runtime.NumCPU() * 2,
		MaxExecCount: 2,
		MaxExecTime:  200 * time.Second,
		Timeout:      200 * time.Second,
		WorkerFunc:   Worker,
		WorkerCommit: commit,
	})
}

// Options contains the options to execute a konnector.
type Options struct {
	Konnector    string `json:"konnector"`
	Account      string `json:"account"`
	FolderToSave string `json:"folder_to_save"`
}

// result stores the result of a konnector execution.
type result struct {
	DocID       string         `json:"_id,omitempty"`
	DocRev      string         `json:"_rev,omitempty"`
	CreatedAt   time.Time      `json:"last_execution"`
	LastSuccess time.Time      `json:"last_success"`
	Logs        []konnectorMsg `json:"logs"`
	Account     string         `json:"account"`
	State       string         `json:"state"`
	Error       string         `json:"error"`
}

func (r *result) ID() string         { return r.DocID }
func (r *result) Rev() string        { return r.DocRev }
func (r *result) DocType() string    { return consts.KonnectorResults }
func (r *result) Clone() couchdb.Doc { c := *r; return &c }
func (r *result) SetID(id string)    { r.DocID = id }
func (r *result) SetRev(rev string)  { r.DocRev = rev }

const konnectorMsgTypeError string = "error"

// const konnectorMsgTypeDebug string = "debug"
// const konnectorMsgTypeWarning string = "warning"
// const konnectorMsgTypeProgress string = "progress"

type konnectorMsg struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

type konnectorLogs struct {
	Slug     string         `json:"_id,omitempty"`
	DocRev   string         `json:"_rev,omitempty"`
	Messages []konnectorMsg `json:"logs"`
}

func (kl *konnectorLogs) ID() string         { return kl.Slug }
func (kl *konnectorLogs) Rev() string        { return kl.DocRev }
func (kl *konnectorLogs) DocType() string    { return consts.KonnectorLogs }
func (kl *konnectorLogs) Clone() couchdb.Doc { c := *kl; return &c }
func (kl *konnectorLogs) SetID(id string)    {}
func (kl *konnectorLogs) SetRev(rev string)  { kl.DocRev = rev }

// Worker is the worker that runs a konnector by executing an external process.
func Worker(ctx context.Context, m *jobs.Message) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	fields := struct {
		Account      string `json:"account"`
		FolderToSave string `json:"folder_to_save"`
	}{
		Account:      opts.Account,
		FolderToSave: opts.FolderToSave,
	}
	domain := ctx.Value(jobs.ContextDomainKey).(string)
	worker := ctx.Value(jobs.ContextWorkerKey).(string)
	jobID := fmt.Sprintf("%s/%s/%s", worker, slug, domain)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	man, err := apps.GetKonnectorBySlug(inst, slug)
	if err != nil {
		return err
	}
	if man.State() != apps.Ready {
		return errors.New("Konnector is not ready")
	}

	token := inst.BuildKonnectorToken(man)

	osFS := afero.NewOsFs()
	workDir, err := afero.TempDir(osFS, "", "konnector-"+slug)
	if err != nil {
		return err
	}
	defer osFS.RemoveAll(workDir)
	workFS := afero.NewBasePathFs(osFS, workDir)

	fileServer := inst.KonnectorsFileServer()
	tarFile, err := fileServer.Open(slug, man.Version(), apps.KonnectorArchiveName)
	if err != nil {
		return err
	}

	tr := tar.NewReader(tarFile)
	for {
		var hdr *tar.Header
		hdr, err = tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		dirname := path.Dir(hdr.Name)
		if dirname != "." {
			if err = workFS.MkdirAll(dirname, 0755); err != nil {
				return nil
			}
		}
		var f afero.File
		f, err = workFS.OpenFile(hdr.Name, os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			return err
		}
		_, err = io.Copy(f, tr)
		errc := f.Close()
		if err != nil {
			return err
		}
		if errc != nil {
			return errc
		}
	}

	fieldsJSON, err := json.Marshal(fields)
	if err != nil {
		return err
	}

	konnCmd := config.GetConfig().Konnectors.Cmd
	cmd := exec.CommandContext(ctx, konnCmd, workDir) // #nosec
	cmd.Env = []string{
		"COZY_URL=" + inst.PageURL("/", nil),
		"COZY_CREDENTIALS=" + token,
		"COZY_FIELDS=" + string(fieldsJSON),
		"COZY_TYPE=" + man.Type,
		"COZY_JOB_ID=" + jobID,
	}

	cmdErr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	scanErr := bufio.NewScanner(cmdErr)
	scanOut := bufio.NewScanner(cmdOut)
	scanOut.Buffer(nil, 256*1024)

	waitch := make(chan error)
	logsch := make(chan konnectorMsg, 10)

	var messages []konnectorMsg

	log := logger.WithDomain(domain)

	go doScanOut(jobID, scanOut, domain, logsch, log)
	go doScanErr(jobID, scanErr, log)
	if err = cmd.Start(); err != nil {
		return wrapErr(ctx, err)
	}

	go func() { waitch <- cmd.Wait() }()

	hub := realtime.GetHub()
	for msg := range logsch {
		// TODO: filter some of the messages
		messages = append(messages, msg)
		hub.Publish(&realtime.Event{
			Verb: realtime.EventCreate,
			Doc: couchdb.JSONDoc{Type: consts.JobEvents, M: map[string]interface{}{
				"type":    msg.Type,
				"message": msg.Message,
			}},
			Domain: domain,
		})
	}

	err = <-waitch
	if err != nil {
		err = wrapErr(ctx, err)
	}

	errLogs := couchdb.Upsert(inst, &konnectorLogs{
		Slug:     slug,
		Messages: messages,
	})
	if errLogs != nil {
		fmt.Println("Failed to save konnector logs", errLogs)
	}

	for _, msg := range messages {
		if msg.Type == konnectorMsgTypeError {
			// konnector err is more explicit
			return errors.New(msg.Message)
		}
	}

	return err
}

func doScanOut(jobID string,
	scanner *bufio.Scanner, domain string,
	logsch chan konnectorMsg, log *logrus.Entry) {

	defer close(logsch)

	for scanner.Scan() {
		linebb := scanner.Bytes()
		from := bytes.IndexByte(linebb, '{')
		to := bytes.LastIndexByte(linebb, '}')
		var msg konnectorMsg
		log.Infof("[konnector] %s: Stdout: %s", jobID, string(linebb))
		if from > -1 && from < to && to > -1 {
			err := json.Unmarshal(linebb[from:to+1], &msg)
			if err == nil {
				logsch <- msg
				continue
			}
		}
		log.Warnf("[konnector] %s: Could not parse as JSON", jobID)
		log.Debugf("[konnector] %s: %s", jobID, string(linebb))
	}

	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stdout: %s", jobID, err)
	}
}

func doScanErr(jobID string, scanner *bufio.Scanner, log *logrus.Entry) {
	for scanner.Scan() {
		log.Errorf("[konnector] %s: Stderr: %s", jobID, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stderr: %s", jobID, err)
	}
}

func commit(ctx context.Context, m *jobs.Message, errjob error) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	domain := ctx.Value(jobs.ContextDomainKey).(string)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	lastResult := &result{}
	err = couchdb.GetDoc(inst, consts.KonnectorResults, slug, lastResult)
	if err != nil {
		if !couchdb.IsNotFoundError(err) {
			return err
		}
		lastResult = nil
	}

	var state, errstr string
	var lastSuccess time.Time
	if errjob != nil {
		if lastResult != nil {
			lastSuccess = lastResult.LastSuccess
		}
		errstr = errjob.Error()
		state = jobs.Errored
	} else {
		lastSuccess = time.Now()
		state = jobs.Done
	}
	result := &result{
		DocID:       slug,
		Account:     opts.Account,
		CreatedAt:   time.Now(),
		LastSuccess: lastSuccess,
		State:       state,
		Error:       errstr,
	}
	if lastResult == nil {
		err = couchdb.CreateNamedDocWithDB(inst, result)
	} else {
		result.SetRev(lastResult.Rev())
		err = couchdb.UpdateDoc(inst, result)
	}
	// if err != nil {
	// 	return err
	// }

	return err
	// // if it is the first try we do not take into account an error, we bail.
	// if lastResult == nil {
	// 	return nil
	// }
	// // if the job has not errored, or the last one was already errored, we bail.
	// if state != jobs.Errored || lastResult.State == jobs.Errored {
	// 	return nil
	// }

	// konnectorURL := inst.SubDomain(consts.CollectSlug)
	// konnectorURL.Fragment = "/category/all/" + slug
	// mail := mails.Options{
	// 	Mode:         mails.ModeNoReply,
	// 	Subject:      inst.Translate("Error Konnector execution", domain),
	// 	TemplateName: "konnector_error_" + inst.Locale,
	// 	TemplateValues: map[string]string{
	// 		"KonnectorName": slug,
	// 		"KonnectorPage": konnectorURL.String(),
	// 	},
	// }
	// msg, err := jobs.NewMessage(jobs.JSONEncoding, &mail)
	// if err != nil {
	// 	return err
	// }
	// log := logger.WithDomain(domain)
	// log.Info("Konnector has failed definitively, should send mail.", mail)
	// _, err = stack.GetBroker().PushJob(&jobs.JobRequest{
	// 	Domain:     domain,
	// 	WorkerType: "sendmail",
	// 	Message:    msg,
	// })
	// return err
}

func wrapErr(ctx context.Context, err error) error {
	if ctx.Err() == context.DeadlineExceeded {
		return context.DeadlineExceeded
	}
	return err
}
<commit_msg>More idiomatic use of cmd.Wait()<commit_after>package konnectors

import (
	"archive/tar"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/sirupsen/logrus"

	"github.com/spf13/afero"
)

func init() {
	jobs.AddWorker("konnector", &jobs.WorkerConfig{
		Concurrency:  runtime.NumCPU() * 2,
		MaxExecCount: 2,
		MaxExecTime:  200 * time.Second,
		Timeout:      200 * time.Second,
		WorkerFunc:   Worker,
		WorkerCommit: commit,
	})
}

// Options contains the options to execute a konnector.
type Options struct {
	Konnector    string `json:"konnector"`
	Account      string `json:"account"`
	FolderToSave string `json:"folder_to_save"`
}

// result stores the result of a konnector execution.
type result struct {
	DocID       string         `json:"_id,omitempty"`
	DocRev      string         `json:"_rev,omitempty"`
	CreatedAt   time.Time      `json:"last_execution"`
	LastSuccess time.Time      `json:"last_success"`
	Logs        []konnectorMsg `json:"logs"`
	Account     string         `json:"account"`
	State       string         `json:"state"`
	Error       string         `json:"error"`
}

func (r *result) ID() string         { return r.DocID }
func (r *result) Rev() string        { return r.DocRev }
func (r *result) DocType() string    { return consts.KonnectorResults }
func (r *result) Clone() couchdb.Doc { c := *r; return &c }
func (r *result) SetID(id string)    { r.DocID = id }
func (r *result) SetRev(rev string)  { r.DocRev = rev }

const konnectorMsgTypeError string = "error"

// const konnectorMsgTypeDebug string = "debug"
// const konnectorMsgTypeWarning string = "warning"
// const konnectorMsgTypeProgress string = "progress"

type konnectorMsg struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

type konnectorLogs struct {
	Slug     string         `json:"_id,omitempty"`
	DocRev   string         `json:"_rev,omitempty"`
	Messages []konnectorMsg `json:"logs"`
}

func (kl *konnectorLogs) ID() string         { return kl.Slug }
func (kl *konnectorLogs) Rev() string        { return kl.DocRev }
func (kl *konnectorLogs) DocType() string    { return consts.KonnectorLogs }
func (kl *konnectorLogs) Clone() couchdb.Doc { c := *kl; return &c }
func (kl *konnectorLogs) SetID(id string)    {}
func (kl *konnectorLogs) SetRev(rev string)  { kl.DocRev = rev }

// Worker is the worker that runs a konnector by executing an external process.
func Worker(ctx context.Context, m *jobs.Message) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	fields := struct {
		Account      string `json:"account"`
		FolderToSave string `json:"folder_to_save"`
	}{
		Account:      opts.Account,
		FolderToSave: opts.FolderToSave,
	}
	domain := ctx.Value(jobs.ContextDomainKey).(string)
	worker := ctx.Value(jobs.ContextWorkerKey).(string)
	jobID := fmt.Sprintf("%s/%s/%s", worker, slug, domain)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	man, err := apps.GetKonnectorBySlug(inst, slug)
	if err != nil {
		return err
	}
	if man.State() != apps.Ready {
		return errors.New("Konnector is not ready")
	}

	token := inst.BuildKonnectorToken(man)

	osFS := afero.NewOsFs()
	workDir, err := afero.TempDir(osFS, "", "konnector-"+slug)
	if err != nil {
		return err
	}
	defer osFS.RemoveAll(workDir)
	workFS := afero.NewBasePathFs(osFS, workDir)

	fileServer := inst.KonnectorsFileServer()
	tarFile, err := fileServer.Open(slug, man.Version(), apps.KonnectorArchiveName)
	if err != nil {
		return err
	}

	tr := tar.NewReader(tarFile)
	for {
		var hdr *tar.Header
		hdr, err = tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		dirname := path.Dir(hdr.Name)
		if dirname != "." {
			if err = workFS.MkdirAll(dirname, 0755); err != nil {
				return nil
			}
		}
		var f afero.File
		f, err = workFS.OpenFile(hdr.Name, os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			return err
		}
		_, err = io.Copy(f, tr)
		errc := f.Close()
		if err != nil {
			return err
		}
		if errc != nil {
			return errc
		}
	}

	fieldsJSON, err := json.Marshal(fields)
	if err != nil {
		return err
	}

	konnCmd := config.GetConfig().Konnectors.Cmd
	cmd := exec.CommandContext(ctx, konnCmd, workDir) // #nosec
	cmd.Env = []string{
		"COZY_URL=" + inst.PageURL("/", nil),
		"COZY_CREDENTIALS=" + token,
		"COZY_FIELDS=" + string(fieldsJSON),
		"COZY_TYPE=" + man.Type,
		"COZY_JOB_ID=" + jobID,
	}

	cmdErr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	scanErr := bufio.NewScanner(cmdErr)
	scanOut := bufio.NewScanner(cmdOut)
	scanOut.Buffer(nil, 256*1024)

	var messages []konnectorMsg

	log := logger.WithDomain(domain)
	logsch := make(chan konnectorMsg, 10)

	if err = cmd.Start(); err != nil {
		return wrapErr(ctx, err)
	}

	go doScanOut(jobID, scanOut, domain, logsch, log)
	go doScanErr(jobID, scanErr, log)

	hub := realtime.GetHub()
	for msg := range logsch {
		// TODO: filter some of the messages
		messages = append(messages, msg)
		hub.Publish(&realtime.Event{
			Verb: realtime.EventCreate,
			Doc: couchdb.JSONDoc{Type: consts.JobEvents, M: map[string]interface{}{
				"type":    msg.Type,
				"message": msg.Message,
			}},
			Domain: domain,
		})
	}

	if err = cmd.Wait(); err != nil {
		err = wrapErr(ctx, err)
	}

	errLogs := couchdb.Upsert(inst, &konnectorLogs{
		Slug:     slug,
		Messages: messages,
	})
	if errLogs != nil {
		fmt.Println("Failed to save konnector logs", errLogs)
	}

	for _, msg := range messages {
		if msg.Type == konnectorMsgTypeError {
			// konnector err is more explicit
			return errors.New(msg.Message)
		}
	}

	return err
}

func doScanOut(jobID string,
	scanner *bufio.Scanner, domain string,
	logsch chan konnectorMsg, log *logrus.Entry) {

	defer close(logsch)

	for scanner.Scan() {
		linebb := scanner.Bytes()
		from := bytes.IndexByte(linebb, '{')
		to := bytes.LastIndexByte(linebb, '}')
		var msg konnectorMsg
		log.Infof("[konnector] %s: Stdout: %s", jobID, string(linebb))
		if from > -1 && from < to && to > -1 {
			err := json.Unmarshal(linebb[from:to+1], &msg)
			if err == nil {
				logsch <- msg
				continue
			}
		}
		log.Warnf("[konnector] %s: Could not parse as JSON", jobID)
		log.Debugf("[konnector] %s: %s", jobID, string(linebb))
	}

	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stdout: %s", jobID, err)
	}
}

func doScanErr(jobID string, scanner *bufio.Scanner, log *logrus.Entry) {
	for scanner.Scan() {
		log.Errorf("[konnector] %s: Stderr: %s", jobID, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stderr: %s", jobID, err)
	}
}

func commit(ctx context.Context, m *jobs.Message, errjob error) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	domain := ctx.Value(jobs.ContextDomainKey).(string)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	lastResult := &result{}
	err = couchdb.GetDoc(inst, consts.KonnectorResults, slug, lastResult)
	if err != nil {
		if !couchdb.IsNotFoundError(err) {
			return err
		}
		lastResult = nil
	}

	var state, errstr string
	var lastSuccess time.Time
	if errjob != nil {
		if lastResult != nil {
			lastSuccess = lastResult.LastSuccess
		}
		errstr = errjob.Error()
		state = jobs.Errored
	} else {
		lastSuccess = time.Now()
		state = jobs.Done
	}
	result := &result{
		DocID:       slug,
		Account:     opts.Account,
		CreatedAt:   time.Now(),
		LastSuccess: lastSuccess,
		State:       state,
		Error:       errstr,
	}
	if lastResult == nil {
		err = couchdb.CreateNamedDocWithDB(inst, result)
	} else {
		result.SetRev(lastResult.Rev())
		err = couchdb.UpdateDoc(inst, result)
	}
	// if err != nil {
	// 	return err
	// }

	return err
	// // if it is the first try we do not take into account an error, we bail.
	// if lastResult == nil {
	// 	return nil
	// }
	// // if the job has not errored, or the last one was already errored, we bail.
	// if state != jobs.Errored || lastResult.State == jobs.Errored {
	// 	return nil
	// }

	// konnectorURL := inst.SubDomain(consts.CollectSlug)
	// konnectorURL.Fragment = "/category/all/" + slug
	// mail := mails.Options{
	// 	Mode:         mails.ModeNoReply,
	// 	Subject:      inst.Translate("Error Konnector execution", domain),
	// 	TemplateName: "konnector_error_" + inst.Locale,
	// 	TemplateValues: map[string]string{
	// 		"KonnectorName": slug,
	// 		"KonnectorPage": konnectorURL.String(),
	// 	},
	// }
	// msg, err := jobs.NewMessage(jobs.JSONEncoding, &mail)
	// if err != nil {
	// 	return err
	// }
	// log := logger.WithDomain(domain)
	// log.Info("Konnector has failed definitively, should send mail.", mail)
	// _, err = stack.GetBroker().PushJob(&jobs.JobRequest{
	// 	Domain:     domain,
	// 	WorkerType: "sendmail",
	// 	Message:    msg,
	// })
	// return err
}

func wrapErr(ctx context.Context, err error) error {
	if ctx.Err() == context.DeadlineExceeded {
		return context.DeadlineExceeded
	}
	return err
}
<|endoftext|>


<commit_before>package itchio

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"net/http"
	"net/http/httptest"
	"net/url"
)

func testTools(code int, body string) (*httptest.Server, *Client) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(code)
		fmt.Fprintln(w, body)
	}))

	// Make a transport that reroutes all traffic to the example server
	transport := &http.Transport{
		Proxy: func(req *http.Request) (*url.URL, error) {
			return url.Parse(server.URL)
		},
	}

	// Make a http.Client with the transport
	httpClient := &http.Client{Transport: transport}

	client := &Client{
		Key:        "APIKEY",
		HTTPClient: httpClient,
		BaseURL:    server.URL,
	}

	return server, client
}

func Test_ListProfileGames(t *testing.T) {
	server, client := testTools(200, `{
		"games": [
		  {"url": "https://kenney.itch.io/barb", "id": 123, "min_price": 5000},
		  {"url": "https://leafo.itch.io/x-moon", "id": 456, "min_price": 12000}
		]
	}`)
	defer server.Close()

	games, err := client.ListProfileGames()
	assert.NoError(t, err)
	assert.EqualValues(t, len(games.Games), 2)
	assert.EqualValues(t, games.Games[0].ID, 123)
	assert.EqualValues(t, games.Games[0].URL, "https://kenney.itch.io/barb")
	assert.EqualValues(t, games.Games[0].MinPrice, 5000)
}

func Test_ListProfileGamesError(t *testing.T) {
	server, client := testTools(400, `{
		"errors": [
		  "invalid game"
		]
	}`)
	defer server.Close()

	_, err := client.ListProfileGames()
	assert.Error(t, err)
	assert.True(t, IsAPIError(err))
	assert.EqualValues(t, "itch.io API error (400): /profile/games: invalid game", err.Error())
}

func Test_ParseSpec(t *testing.T) {
	var spec *Spec
	var err error

	spec, err = ParseSpec("user/page:channel")
	assert.NoError(t, err)
	assert.Equal(t, spec.Target, "user/page")
	assert.Equal(t, spec.Channel, "channel")

	spec, err = ParseSpec("user/page")
	assert.NoError(t, err)
	assert.Equal(t, spec.Target, "user/page")
	assert.Equal(t, spec.Channel, "")

	err = spec.EnsureChannel()
	assert.Error(t, err)

	_, err = ParseSpec("a:b:c")
	assert.Error(t, err)
}
<commit_msg>use ClientWithKey to make a new client<commit_after>package itchio

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"net/http"
	"net/http/httptest"
	"net/url"
)

func testTools(code int, body string) (*httptest.Server, *Client) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(code)
		fmt.Fprintln(w, body)
	}))

	// Make a transport that reroutes all traffic to the example server
	transport := &http.Transport{
		Proxy: func(req *http.Request) (*url.URL, error) {
			return url.Parse(server.URL)
		},
	}

	// Make a http.Client with the transport
	httpClient := &http.Client{Transport: transport}

	client := ClientWithKey("APIKEY")
	client.HTTPClient = httpClient
	client.BaseURL = server.URL

	return server, client
}

func Test_ListProfileGames(t *testing.T) {
	server, client := testTools(200, `{
		"games": [
		  {"url": "https://kenney.itch.io/barb", "id": 123, "min_price": 5000},
		  {"url": "https://leafo.itch.io/x-moon", "id": 456, "min_price": 12000}
		]
	}`)
	defer server.Close()

	games, err := client.ListProfileGames()
	assert.NoError(t, err)
	assert.EqualValues(t, len(games.Games), 2)
	assert.EqualValues(t, games.Games[0].ID, 123)
	assert.EqualValues(t, games.Games[0].URL, "https://kenney.itch.io/barb")
	assert.EqualValues(t, games.Games[0].MinPrice, 5000)
}

func Test_ListProfileGamesError(t *testing.T) {
	server, client := testTools(400, `{
		"errors": [
		  "invalid game"
		]
	}`)
	defer server.Close()

	_, err := client.ListProfileGames()
	assert.Error(t, err)
	assert.True(t, IsAPIError(err))
	assert.EqualValues(t, "itch.io API error (400): /profile/games: invalid game", err.Error())
}

func Test_ParseSpec(t *testing.T) {
	var spec *Spec
	var err error

	spec, err = ParseSpec("user/page:channel")
	assert.NoError(t, err)
	assert.Equal(t, spec.Target, "user/page")
	assert.Equal(t, spec.Channel, "channel")

	spec, err = ParseSpec("user/page")
	assert.NoError(t, err)
	assert.Equal(t, spec.Target, "user/page")
	assert.Equal(t, spec.Channel, "")

	err = spec.EnsureChannel()
	assert.Error(t, err)

	_, err = ParseSpec("a:b:c")
	assert.Error(t, err)
}
<|endoftext|>