<commit_before>// Copyright 2016-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 loader

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/byteorder"
	"github.com/cilium/cilium/pkg/cgroups"
	"github.com/cilium/cilium/pkg/command/exec"
	"github.com/cilium/cilium/pkg/common"
	"github.com/cilium/cilium/pkg/datapath"
	"github.com/cilium/cilium/pkg/datapath/alignchecker"
	datapathOption "github.com/cilium/cilium/pkg/datapath/option"
	"github.com/cilium/cilium/pkg/datapath/prefilter"
	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/sysctl"
	"github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

const (
	initArgLib int = iota
	initArgRundir
	initArgIPv4NodeIP
	initArgIPv6NodeIP
	initArgMode
	initArgDevices
	initArgXDPDevice
	initArgXDPMode
	initArgMTU
	initArgIPSec
	initArgEncryptInterface
	initArgHostReachableServices
	initArgHostReachableServicesUDP
	initArgHostReachableServicesPeer
	initArgCgroupRoot
	initArgBpffsRoot
	initArgNodePort
	initArgNodePortBind
	initBPFCPU
	initArgNodePortIPv4Addrs
	initArgNodePortIPv6Addrs
	initArgNrCPUs
	initArgMax
)

// firstInitialization is true when Reinitialize() is called for the first
// time. It can only be accessed when GetCompilationLock() is being held.
var firstInitialization = true

func (l *Loader) writeNetdevHeader(dir string, o datapath.BaseProgramOwner) error {
	headerPath := filepath.Join(dir, common.NetdevHeaderFileName)
	log.WithField(logfields.Path, headerPath).Debug("writing configuration")

	f, err := os.Create(headerPath)
	if err != nil {
		return fmt.Errorf("failed to open file %s for writing: %s", headerPath, err)

	}
	defer f.Close()

	if err := l.templateCache.WriteNetdevConfig(f, o); err != nil {
		return err
	}
	return nil
}

// Must be called with option.Config.EnablePolicyMU locked.
func writePreFilterHeader(preFilter *prefilter.PreFilter, dir string) error {
	headerPath := filepath.Join(dir, common.PreFilterHeaderFileName)
	log.WithField(logfields.Path, headerPath).Debug("writing configuration")
	f, err := os.Create(headerPath)
	if err != nil {
		return fmt.Errorf("failed to open file %s for writing: %s", headerPath, err)

	}
	defer f.Close()
	fw := bufio.NewWriter(f)
	fmt.Fprint(fw, "/*\n")
	fmt.Fprintf(fw, " * XDP device: %s\n", option.Config.DevicePreFilter)
	fmt.Fprintf(fw, " * XDP mode: %s\n", option.Config.ModePreFilter)
	fmt.Fprint(fw, " */\n\n")
	preFilter.WriteConfig(fw)
	return fw.Flush()
}

// Reinitialize (re-)configures the base datapath configuration including global
// BPF programs, netfilter rule configuration and reserving routes in IPAM for
// locally detected prefixes. It may be run upon initial Cilium startup, after
// restore from a previous Cilium run, or during regular Cilium operation.
func (l *Loader) Reinitialize(ctx context.Context, o datapath.BaseProgramOwner, deviceMTU int, iptMgr datapath.IptablesManager, p datapath.Proxy, r datapath.RouteReserver) error {
	var (
		args []string
		ret  error
	)

	type setting struct {
		name      string
		val       string
		ignoreErr bool
	}

	args = make([]string, initArgMax)

	sysSettings := []setting{
		{"net.core.bpf_jit_enable", "1", true},
		{"net.ipv4.conf.all.rp_filter", "0", false},
		{"kernel.unprivileged_bpf_disabled", "1", true},
	}

	// Lock so that endpoints cannot be built while we are compile base programs.
	o.GetCompilationLock().Lock()
	defer o.GetCompilationLock().Unlock()
	defer func() { firstInitialization = false }()

	l.init(o.Datapath(), o.LocalConfig())

	if err := l.writeNetdevHeader("./", o); err != nil {
		log.WithError(err).Warn("Unable to write netdev header")
		return err
	}

	if option.Config.XDPDevice != "undefined" {
		args[initArgXDPDevice] = option.Config.XDPDevice
		args[initArgXDPMode] = option.Config.XDPMode
	} else {
		args[initArgXDPDevice] = "<nil>"
		args[initArgXDPMode] = "<nil>"
	}

	if option.Config.DevicePreFilter != "undefined" {
		scopedLog := log.WithField(logfields.XDPDevice, option.Config.XDPDevice)

		preFilter, err := prefilter.NewPreFilter()
		if err != nil {
			scopedLog.WithError(ret).Warn("Unable to init prefilter")
			return ret
		}

		if err := writePreFilterHeader(preFilter, "./"); err != nil {
			scopedLog.WithError(err).Warn("Unable to write prefilter header")
			return err
		}

		o.SetPrefilter(preFilter)
	}

	args[initArgLib] = option.Config.BpfDir
	args[initArgRundir] = option.Config.StateDir
	args[initArgCgroupRoot] = cgroups.GetCgroupRoot()
	args[initArgBpffsRoot] = bpf.GetMapRoot()

	if option.Config.EnableIPv4 {
		args[initArgIPv4NodeIP] = node.GetInternalIPv4().String()
	} else {
		args[initArgIPv4NodeIP] = "<nil>"
	}

	if option.Config.EnableIPv6 {
		args[initArgIPv6NodeIP] = node.GetIPv6().String()
		// Docker <17.05 has an issue which causes IPv6 to be disabled in the initns for all
		// interface (https://github.com/docker/libnetwork/issues/1720)
		// Enable IPv6 for now
		sysSettings = append(sysSettings,
			setting{"net.ipv6.conf.all.disable_ipv6", "0", false})
	} else {
		args[initArgIPv6NodeIP] = "<nil>"
	}

	args[initArgMTU] = fmt.Sprintf("%d", deviceMTU)

	if option.Config.EnableIPSec {
		args[initArgIPSec] = "true"
	} else {
		args[initArgIPSec] = "false"
	}

	if option.Config.EnableHostReachableServices {
		args[initArgHostReachableServices] = "true"
		if option.Config.EnableHostServicesUDP {
			args[initArgHostReachableServicesUDP] = "true"
		} else {
			args[initArgHostReachableServicesUDP] = "false"
		}
		if option.Config.EnableHostServicesPeer {
			args[initArgHostReachableServicesPeer] = "true"
		} else {
			args[initArgHostReachableServicesPeer] = "false"
		}
	} else {
		args[initArgHostReachableServices] = "false"
		args[initArgHostReachableServicesUDP] = "false"
		args[initArgHostReachableServicesPeer] = "false"
	}

	if option.Config.EncryptInterface != "" {
		args[initArgEncryptInterface] = option.Config.EncryptInterface
	} else {
		args[initArgEncryptInterface] = "<nil>"
	}

	if len(option.Config.Devices) != 0 {
		for _, device := range option.Config.Devices {
			_, err := netlink.LinkByName(device)
			if err != nil {
				log.WithError(err).WithField("device", device).Warn("Link does not exist")
				return err
			}
		}
		args[initArgDevices] = strings.Join(option.Config.Devices, ";")
	} else if option.Config.IsFlannelMasterDeviceSet() {
		args[initArgDevices] = option.Config.FlannelMasterDevice
	} else {
		args[initArgDevices] = "<nil>"
	}

	switch {
	case option.Config.EnableEndpointRoutes:
		args[initArgMode] = "routed"
	case option.Config.IsFlannelMasterDeviceSet():
		args[initArgMode] = "flannel"
	case option.Config.Tunnel != option.TunnelDisabled:
		args[initArgMode] = option.Config.Tunnel
	case option.Config.DatapathMode == datapathOption.DatapathModeIpvlan:
		args[initArgMode] = "ipvlan"
	default:
		args[initArgMode] = "direct"
	}

	if option.Config.EnableNodePort {
		args[initArgNodePort] = "true"
		if option.Config.EnableIPv4 {
			addrs := node.GetNodePortIPv4AddrsWithDevices()
			tmp := make([]string, 0, len(addrs))
			for iface, ipv4 := range addrs {
				tmp = append(tmp,
					fmt.Sprintf("%s=%#x", iface,
						byteorder.HostSliceToNetwork(ipv4, reflect.Uint32).(uint32)))
			}
			args[initArgNodePortIPv4Addrs] = strings.Join(tmp, ";")
		} else {
			args[initArgNodePortIPv4Addrs] = "<nil>"
		}
		if option.Config.EnableIPv6 {
			addrs := node.GetNodePortIPv6AddrsWithDevices()
			tmp := make([]string, 0, len(addrs))
			for iface, ipv6 := range addrs {
				tmp = append(tmp, fmt.Sprintf("%s=%s", iface, common.GoArray2CNoSpaces(ipv6)))
			}
			args[initArgNodePortIPv6Addrs] = strings.Join(tmp, ";")
		} else {
			args[initArgNodePortIPv6Addrs] = "<nil>"
		}
	} else {
		args[initArgNodePort] = "false"
		args[initArgNodePortIPv4Addrs] = "<nil>"
		args[initArgNodePortIPv6Addrs] = "<nil>"
	}

	if option.Config.NodePortBindProtection {
		args[initArgNodePortBind] = "true"
	} else {
		args[initArgNodePortBind] = "false"
	}

	args[initBPFCPU] = GetBPFCPU()
	args[initArgNrCPUs] = fmt.Sprintf("%d", common.GetNumPossibleCPUs(log))

	clockSource := []string{"ktime", "jiffies"}
	log.Infof("Setting up base BPF datapath (BPF %s instruction set, %s clock source)",
		args[initBPFCPU], clockSource[option.Config.ClockSource])

	for _, s := range sysSettings {
		log.Infof("Setting sysctl %s=%s", s.name, s.val)
		if err := sysctl.Write(s.name, s.val); err != nil {
			if !s.ignoreErr {
				return fmt.Errorf("Failed to sysctl -w %s=%s: %s", s.name, s.val, err)
			}
			log.WithError(err).WithFields(logrus.Fields{
				logfields.SysParamName:  s.name,
				logfields.SysParamValue: s.val,
			}).Warning("Failed to sysctl -w")
		}
	}

	for i, arg := range args {
		if arg == "" {
			log.Warningf("empty argument passed to bpf/init.sh at position %d", i)
		}
	}

	prog := filepath.Join(option.Config.BpfDir, "init.sh")
	ctx, cancel := context.WithTimeout(ctx, defaults.ExecTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, prog, args...)
	cmd.Env = bpf.Environment()
	if _, err := cmd.CombinedOutput(log, true); err != nil {
		return err
	}

	if l.canDisableDwarfRelocations {
		// Validate alignments of C and Go equivalent structs
		if err := alignchecker.CheckStructAlignments(defaults.AlignCheckerName); err != nil {
			log.WithError(err).Fatal("C and Go structs alignment check failed")
		}
	} else {
		log.Warning("Cannot check matching of C and Go common struct alignments due to old LLVM/clang version")
	}

	if !option.Config.IsFlannelMasterDeviceSet() {
		r.ReserveLocalRoutes()
	}

	if err := o.Datapath().Node().NodeConfigurationChanged(*o.LocalConfig()); err != nil {
		return err
	}

	if option.Config.InstallIptRules {
		if err := iptMgr.TransientRulesStart(option.Config.HostDevice); err != nil {
			log.WithError(err).Warning("failed to install transient iptables rules")
		}
	}
	// The iptables rules are only removed on the first initialization to
	// remove stale rules or when iptables is enabled. The first invocation
	// is silent as rules may not exist.
	if firstInitialization || option.Config.InstallIptRules {
		iptMgr.RemoveRules(firstInitialization)
	}
	if option.Config.InstallIptRules {
		err := iptMgr.InstallRules(option.Config.HostDevice)
		iptMgr.TransientRulesEnd(false)
		if err != nil {
			return err
		}
	}
	// Reinstall proxy rules for any running proxies
	if p != nil {
		p.ReinstallRules()
	}

	return nil
}
<commit_msg>cilium: disable timer migration<commit_after>// Copyright 2016-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 loader

import (
	"bufio"
	"context"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	"github.com/cilium/cilium/pkg/bpf"
	"github.com/cilium/cilium/pkg/byteorder"
	"github.com/cilium/cilium/pkg/cgroups"
	"github.com/cilium/cilium/pkg/command/exec"
	"github.com/cilium/cilium/pkg/common"
	"github.com/cilium/cilium/pkg/datapath"
	"github.com/cilium/cilium/pkg/datapath/alignchecker"
	datapathOption "github.com/cilium/cilium/pkg/datapath/option"
	"github.com/cilium/cilium/pkg/datapath/prefilter"
	"github.com/cilium/cilium/pkg/defaults"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/node"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/sysctl"
	"github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

const (
	initArgLib int = iota
	initArgRundir
	initArgIPv4NodeIP
	initArgIPv6NodeIP
	initArgMode
	initArgDevices
	initArgXDPDevice
	initArgXDPMode
	initArgMTU
	initArgIPSec
	initArgEncryptInterface
	initArgHostReachableServices
	initArgHostReachableServicesUDP
	initArgHostReachableServicesPeer
	initArgCgroupRoot
	initArgBpffsRoot
	initArgNodePort
	initArgNodePortBind
	initBPFCPU
	initArgNodePortIPv4Addrs
	initArgNodePortIPv6Addrs
	initArgNrCPUs
	initArgMax
)

// firstInitialization is true when Reinitialize() is called for the first
// time. It can only be accessed when GetCompilationLock() is being held.
var firstInitialization = true

func (l *Loader) writeNetdevHeader(dir string, o datapath.BaseProgramOwner) error {
	headerPath := filepath.Join(dir, common.NetdevHeaderFileName)
	log.WithField(logfields.Path, headerPath).Debug("writing configuration")

	f, err := os.Create(headerPath)
	if err != nil {
		return fmt.Errorf("failed to open file %s for writing: %s", headerPath, err)

	}
	defer f.Close()

	if err := l.templateCache.WriteNetdevConfig(f, o); err != nil {
		return err
	}
	return nil
}

// Must be called with option.Config.EnablePolicyMU locked.
func writePreFilterHeader(preFilter *prefilter.PreFilter, dir string) error {
	headerPath := filepath.Join(dir, common.PreFilterHeaderFileName)
	log.WithField(logfields.Path, headerPath).Debug("writing configuration")
	f, err := os.Create(headerPath)
	if err != nil {
		return fmt.Errorf("failed to open file %s for writing: %s", headerPath, err)

	}
	defer f.Close()
	fw := bufio.NewWriter(f)
	fmt.Fprint(fw, "/*\n")
	fmt.Fprintf(fw, " * XDP device: %s\n", option.Config.DevicePreFilter)
	fmt.Fprintf(fw, " * XDP mode: %s\n", option.Config.ModePreFilter)
	fmt.Fprint(fw, " */\n\n")
	preFilter.WriteConfig(fw)
	return fw.Flush()
}

// Reinitialize (re-)configures the base datapath configuration including global
// BPF programs, netfilter rule configuration and reserving routes in IPAM for
// locally detected prefixes. It may be run upon initial Cilium startup, after
// restore from a previous Cilium run, or during regular Cilium operation.
func (l *Loader) Reinitialize(ctx context.Context, o datapath.BaseProgramOwner, deviceMTU int, iptMgr datapath.IptablesManager, p datapath.Proxy, r datapath.RouteReserver) error {
	var (
		args []string
		ret  error
	)

	type setting struct {
		name      string
		val       string
		ignoreErr bool
	}

	args = make([]string, initArgMax)

	sysSettings := []setting{
		{"net.core.bpf_jit_enable", "1", true},
		{"net.ipv4.conf.all.rp_filter", "0", false},
		{"kernel.unprivileged_bpf_disabled", "1", true},
		{"kernel.timer_migration", "0", true},
	}

	// Lock so that endpoints cannot be built while we are compile base programs.
	o.GetCompilationLock().Lock()
	defer o.GetCompilationLock().Unlock()
	defer func() { firstInitialization = false }()

	l.init(o.Datapath(), o.LocalConfig())

	if err := l.writeNetdevHeader("./", o); err != nil {
		log.WithError(err).Warn("Unable to write netdev header")
		return err
	}

	if option.Config.XDPDevice != "undefined" {
		args[initArgXDPDevice] = option.Config.XDPDevice
		args[initArgXDPMode] = option.Config.XDPMode
	} else {
		args[initArgXDPDevice] = "<nil>"
		args[initArgXDPMode] = "<nil>"
	}

	if option.Config.DevicePreFilter != "undefined" {
		scopedLog := log.WithField(logfields.XDPDevice, option.Config.XDPDevice)

		preFilter, err := prefilter.NewPreFilter()
		if err != nil {
			scopedLog.WithError(ret).Warn("Unable to init prefilter")
			return ret
		}

		if err := writePreFilterHeader(preFilter, "./"); err != nil {
			scopedLog.WithError(err).Warn("Unable to write prefilter header")
			return err
		}

		o.SetPrefilter(preFilter)
	}

	args[initArgLib] = option.Config.BpfDir
	args[initArgRundir] = option.Config.StateDir
	args[initArgCgroupRoot] = cgroups.GetCgroupRoot()
	args[initArgBpffsRoot] = bpf.GetMapRoot()

	if option.Config.EnableIPv4 {
		args[initArgIPv4NodeIP] = node.GetInternalIPv4().String()
	} else {
		args[initArgIPv4NodeIP] = "<nil>"
	}

	if option.Config.EnableIPv6 {
		args[initArgIPv6NodeIP] = node.GetIPv6().String()
		// Docker <17.05 has an issue which causes IPv6 to be disabled in the initns for all
		// interface (https://github.com/docker/libnetwork/issues/1720)
		// Enable IPv6 for now
		sysSettings = append(sysSettings,
			setting{"net.ipv6.conf.all.disable_ipv6", "0", false})
	} else {
		args[initArgIPv6NodeIP] = "<nil>"
	}

	args[initArgMTU] = fmt.Sprintf("%d", deviceMTU)

	if option.Config.EnableIPSec {
		args[initArgIPSec] = "true"
	} else {
		args[initArgIPSec] = "false"
	}

	if option.Config.EnableHostReachableServices {
		args[initArgHostReachableServices] = "true"
		if option.Config.EnableHostServicesUDP {
			args[initArgHostReachableServicesUDP] = "true"
		} else {
			args[initArgHostReachableServicesUDP] = "false"
		}
		if option.Config.EnableHostServicesPeer {
			args[initArgHostReachableServicesPeer] = "true"
		} else {
			args[initArgHostReachableServicesPeer] = "false"
		}
	} else {
		args[initArgHostReachableServices] = "false"
		args[initArgHostReachableServicesUDP] = "false"
		args[initArgHostReachableServicesPeer] = "false"
	}

	if option.Config.EncryptInterface != "" {
		args[initArgEncryptInterface] = option.Config.EncryptInterface
	} else {
		args[initArgEncryptInterface] = "<nil>"
	}

	if len(option.Config.Devices) != 0 {
		for _, device := range option.Config.Devices {
			_, err := netlink.LinkByName(device)
			if err != nil {
				log.WithError(err).WithField("device", device).Warn("Link does not exist")
				return err
			}
		}
		args[initArgDevices] = strings.Join(option.Config.Devices, ";")
	} else if option.Config.IsFlannelMasterDeviceSet() {
		args[initArgDevices] = option.Config.FlannelMasterDevice
	} else {
		args[initArgDevices] = "<nil>"
	}

	switch {
	case option.Config.EnableEndpointRoutes:
		args[initArgMode] = "routed"
	case option.Config.IsFlannelMasterDeviceSet():
		args[initArgMode] = "flannel"
	case option.Config.Tunnel != option.TunnelDisabled:
		args[initArgMode] = option.Config.Tunnel
	case option.Config.DatapathMode == datapathOption.DatapathModeIpvlan:
		args[initArgMode] = "ipvlan"
	default:
		args[initArgMode] = "direct"
	}

	if option.Config.EnableNodePort {
		args[initArgNodePort] = "true"
		if option.Config.EnableIPv4 {
			addrs := node.GetNodePortIPv4AddrsWithDevices()
			tmp := make([]string, 0, len(addrs))
			for iface, ipv4 := range addrs {
				tmp = append(tmp,
					fmt.Sprintf("%s=%#x", iface,
						byteorder.HostSliceToNetwork(ipv4, reflect.Uint32).(uint32)))
			}
			args[initArgNodePortIPv4Addrs] = strings.Join(tmp, ";")
		} else {
			args[initArgNodePortIPv4Addrs] = "<nil>"
		}
		if option.Config.EnableIPv6 {
			addrs := node.GetNodePortIPv6AddrsWithDevices()
			tmp := make([]string, 0, len(addrs))
			for iface, ipv6 := range addrs {
				tmp = append(tmp, fmt.Sprintf("%s=%s", iface, common.GoArray2CNoSpaces(ipv6)))
			}
			args[initArgNodePortIPv6Addrs] = strings.Join(tmp, ";")
		} else {
			args[initArgNodePortIPv6Addrs] = "<nil>"
		}
	} else {
		args[initArgNodePort] = "false"
		args[initArgNodePortIPv4Addrs] = "<nil>"
		args[initArgNodePortIPv6Addrs] = "<nil>"
	}

	if option.Config.NodePortBindProtection {
		args[initArgNodePortBind] = "true"
	} else {
		args[initArgNodePortBind] = "false"
	}

	args[initBPFCPU] = GetBPFCPU()
	args[initArgNrCPUs] = fmt.Sprintf("%d", common.GetNumPossibleCPUs(log))

	clockSource := []string{"ktime", "jiffies"}
	log.Infof("Setting up base BPF datapath (BPF %s instruction set, %s clock source)",
		args[initBPFCPU], clockSource[option.Config.ClockSource])

	for _, s := range sysSettings {
		log.Infof("Setting sysctl %s=%s", s.name, s.val)
		if err := sysctl.Write(s.name, s.val); err != nil {
			if !s.ignoreErr {
				return fmt.Errorf("Failed to sysctl -w %s=%s: %s", s.name, s.val, err)
			}
			log.WithError(err).WithFields(logrus.Fields{
				logfields.SysParamName:  s.name,
				logfields.SysParamValue: s.val,
			}).Warning("Failed to sysctl -w")
		}
	}

	for i, arg := range args {
		if arg == "" {
			log.Warningf("empty argument passed to bpf/init.sh at position %d", i)
		}
	}

	prog := filepath.Join(option.Config.BpfDir, "init.sh")
	ctx, cancel := context.WithTimeout(ctx, defaults.ExecTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, prog, args...)
	cmd.Env = bpf.Environment()
	if _, err := cmd.CombinedOutput(log, true); err != nil {
		return err
	}

	if l.canDisableDwarfRelocations {
		// Validate alignments of C and Go equivalent structs
		if err := alignchecker.CheckStructAlignments(defaults.AlignCheckerName); err != nil {
			log.WithError(err).Fatal("C and Go structs alignment check failed")
		}
	} else {
		log.Warning("Cannot check matching of C and Go common struct alignments due to old LLVM/clang version")
	}

	if !option.Config.IsFlannelMasterDeviceSet() {
		r.ReserveLocalRoutes()
	}

	if err := o.Datapath().Node().NodeConfigurationChanged(*o.LocalConfig()); err != nil {
		return err
	}

	if option.Config.InstallIptRules {
		if err := iptMgr.TransientRulesStart(option.Config.HostDevice); err != nil {
			log.WithError(err).Warning("failed to install transient iptables rules")
		}
	}
	// The iptables rules are only removed on the first initialization to
	// remove stale rules or when iptables is enabled. The first invocation
	// is silent as rules may not exist.
	if firstInitialization || option.Config.InstallIptRules {
		iptMgr.RemoveRules(firstInitialization)
	}
	if option.Config.InstallIptRules {
		err := iptMgr.InstallRules(option.Config.HostDevice)
		iptMgr.TransientRulesEnd(false)
		if err != nil {
			return err
		}
	}
	// Reinstall proxy rules for any running proxies
	if p != nil {
		p.ReinstallRules()
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package endpoints

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"regexp"
	"sync"
	"testing"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	auditinternal "k8s.io/apiserver/pkg/apis/audit"
	genericapitesting "k8s.io/apiserver/pkg/endpoints/testing"
	"k8s.io/apiserver/pkg/registry/rest"
)

type fakeAuditSink struct {
	lock   sync.Mutex
	events []*auditinternal.Event
}

func (s *fakeAuditSink) ProcessEvents(evs ...*auditinternal.Event) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.events = append(s.events, evs...)
}

func (s *fakeAuditSink) Events() []*auditinternal.Event {
	s.lock.Lock()
	defer s.lock.Unlock()
	return append([]*auditinternal.Event{}, s.events...)
}

func TestAudit(t *testing.T) {
	type eventCheck func(events []*auditinternal.Event) error

	// fixtures
	simpleFoo := &genericapitesting.Simple{Other: "foo"}
	simpleFooJSON, _ := runtime.Encode(testCodec, simpleFoo)

	simpleCPrime := &genericapitesting.Simple{
		ObjectMeta: metav1.ObjectMeta{Name: "c", Namespace: "other"},
		Other:      "bla",
	}
	simpleCPrimeJSON, _ := runtime.Encode(testCodec, simpleCPrime)

	// event checks
	noRequestBody := func(i int) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].RequestObject == nil {
				return nil
			}
			return fmt.Errorf("expected RequestBody to be nil, got non-nil '%s'", events[i].RequestObject.Raw)
		}
	}
	requestBodyIs := func(i int, text string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].RequestObject == nil {
				if text != "" {
					return fmt.Errorf("expected RequestBody %q, got <nil>", text)
				}
				return nil
			}
			if string(events[i].RequestObject.Raw) != text {
				return fmt.Errorf("expected RequestBody %q, got %q", text, string(events[i].RequestObject.Raw))
			}
			return nil
		}
	}
	requestBodyMatches := func(i int, pattern string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if matched, _ := regexp.Match(pattern, events[i].RequestObject.Raw); !matched {
				return fmt.Errorf("expected RequestBody to match %q, but didn't: %q", pattern, string(events[i].RequestObject.Raw))
			}
			return nil
		}
	}
	noResponseBody := func(i int) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].ResponseObject == nil {
				return nil
			}
			return fmt.Errorf("expected ResponseBody to be nil, got non-nil '%s'", events[i].ResponseObject.Raw)
		}
	}
	responseBodyMatches := func(i int, pattern string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if matched, _ := regexp.Match(pattern, events[i].ResponseObject.Raw); !matched {
				return fmt.Errorf("expected ResponseBody to match %q, but didn't: %q", pattern, string(events[i].ResponseObject.Raw))
			}
			return nil
		}
	}

	for _, test := range []struct {
		desc   string
		req    func(server string) (*http.Request, error)
		linker runtime.SelfLinker
		code   int
		events int
		checks []eventCheck
	}{
		{
			"get",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"name":"c".*}`),
			},
		},
		{
			"list",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple?labelSelector=a%3Dfoobar", nil)
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple",
				namespace:   "other",
			},
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"name":"a".*"name":"b".*}`),
			},
		},
		{
			"create",
			func(server string) (*http.Request, error) {
				return http.NewRequest("POST", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			201,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleFooJSON)),
				responseBodyMatches(0, `{.*"foo".*}`),
			},
		},
		{
			"not-allowed-named-create",
			func(server string) (*http.Request, error) {
				return http.NewRequest("POST", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/named", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			405,
			2,
			[]eventCheck{
				noRequestBody(0),  // the 405 is thrown long before the create handler would be executed
				noResponseBody(0), // the 405 is thrown long before the create handler would be executed
			},
		},
		{
			"delete",
			func(server string) (*http.Request, error) {
				return http.NewRequest("DELETE", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/a", nil)
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"kind":"Status".*"status":"Success".*}`),
			},
		},
		{
			"delete-with-options-in-body",
			func(server string) (*http.Request, error) {
				return http.NewRequest("DELETE", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/a", bytes.NewBuffer([]byte(`{"kind":"DeleteOptions"}`)))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				requestBodyMatches(0, "DeleteOptions"),
				responseBodyMatches(0, `{.*"kind":"Status".*"status":"Success".*}`),
			},
		},
		{
			"update",
			func(server string) (*http.Request, error) {
				return http.NewRequest("PUT", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewBuffer(simpleCPrimeJSON))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleCPrimeJSON)),
				responseBodyMatches(0, `{.*"bla".*}`),
			},
		},
		{
			"update-wrong-namespace",
			func(server string) (*http.Request, error) {
				return http.NewRequest("PUT", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/c", bytes.NewBuffer(simpleCPrimeJSON))
			},
			selfLinker,
			400,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleCPrimeJSON)),
				responseBodyMatches(0, `"Status".*"status":"Failure".*"code":400}`),
			},
		},
		{
			"patch",
			func(server string) (*http.Request, error) {
				req, _ := http.NewRequest("PATCH", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewReader([]byte(`{"labels":{"foo":"bar"}}`)))
				req.Header.Set("Content-Type", "application/merge-patch+json; charset=UTF-8")
				return req, nil
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple/c",
				name:        "c",
				namespace:   "other",
			},
			200,
			2,
			[]eventCheck{
				requestBodyIs(0, `{"labels":{"foo":"bar"}}`),
				responseBodyMatches(0, `"name":"c".*"labels":{"foo":"bar"}`),
			},
		},
		{
			"watch",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple?watch=true", nil)
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple",
				namespace:   "other",
			},
			200,
			3,
			[]eventCheck{
				noRequestBody(0),
				noResponseBody(0),
			},
		},
	} {
		sink := &fakeAuditSink{}
		handler := handleInternal(map[string]rest.Storage{
			"simple": &SimpleRESTStorage{
				list: []genericapitesting.Simple{
					{
						ObjectMeta: metav1.ObjectMeta{Name: "a", Namespace: "other"},
						Other:      "foo",
					},
					{
						ObjectMeta: metav1.ObjectMeta{Name: "b", Namespace: "other"},
						Other:      "foo",
					},
				},
				item: genericapitesting.Simple{
					ObjectMeta: metav1.ObjectMeta{Name: "c", Namespace: "other", UID: "uid"},
					Other:      "foo",
				},
			},
		}, admissionControl, selfLinker, sink)

		server := httptest.NewServer(handler)
		defer server.Close()
		client := http.Client{Timeout: 2 * time.Second}

		req, err := test.req(server.URL)
		if err != nil {
			t.Errorf("[%s] error creating the request: %v", test.desc, err)
		}

		response, err := client.Do(req)
		if err != nil {
			t.Errorf("[%s] error: %v", test.desc, err)
		}

		if response.StatusCode != test.code {
			t.Errorf("[%s] expected http code %d, got %#v", test.desc, test.code, response)
		}

		// close body because the handler might block in Flush, unable to send the remaining event.
		response.Body.Close()

		// wait for events to arrive, at least the given number in the test
		events := []*auditinternal.Event{}
		err = wait.Poll(50*time.Millisecond, wait.ForeverTestTimeout, wait.ConditionFunc(func() (done bool, err error) {
			events = sink.Events()
			return len(events) >= test.events, nil
		}))
		if err != nil {
			t.Errorf("[%s] timeout waiting for events", test.desc)
		}

		if got := len(events); got != test.events {
			t.Errorf("[%s] expected %d audit events, got %d", test.desc, test.events, got)
		} else {
			for i, check := range test.checks {
				err := check(events)
				if err != nil {
					t.Errorf("[%s,%d] %v", test.desc, i, err)
				}
			}
		}

		if len(events) > 0 {
			status := events[len(events)-1].ResponseStatus
			if status == nil {
				t.Errorf("[%s] expected non-nil ResponseStatus in last event", test.desc)
			} else if int(status.Code) != test.code {
				t.Errorf("[%s] expected ResponseStatus.Code=%d, got %d", test.desc, test.code, status.Code)
			}
		}
	}
}
<commit_msg>add ut for audit useragent<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package endpoints

import (
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"regexp"
	"sync"
	"testing"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	auditinternal "k8s.io/apiserver/pkg/apis/audit"
	genericapitesting "k8s.io/apiserver/pkg/endpoints/testing"
	"k8s.io/apiserver/pkg/registry/rest"
)

type fakeAuditSink struct {
	lock   sync.Mutex
	events []*auditinternal.Event
}

func (s *fakeAuditSink) ProcessEvents(evs ...*auditinternal.Event) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.events = append(s.events, evs...)
}

func (s *fakeAuditSink) Events() []*auditinternal.Event {
	s.lock.Lock()
	defer s.lock.Unlock()
	return append([]*auditinternal.Event{}, s.events...)
}

func TestAudit(t *testing.T) {
	type eventCheck func(events []*auditinternal.Event) error

	// fixtures
	simpleFoo := &genericapitesting.Simple{Other: "foo"}
	simpleFooJSON, _ := runtime.Encode(testCodec, simpleFoo)

	simpleCPrime := &genericapitesting.Simple{
		ObjectMeta: metav1.ObjectMeta{Name: "c", Namespace: "other"},
		Other:      "bla",
	}
	simpleCPrimeJSON, _ := runtime.Encode(testCodec, simpleCPrime)
	userAgent := "audit-test"

	// event checks
	noRequestBody := func(i int) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].RequestObject == nil {
				return nil
			}
			return fmt.Errorf("expected RequestBody to be nil, got non-nil '%s'", events[i].RequestObject.Raw)
		}
	}
	requestBodyIs := func(i int, text string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].RequestObject == nil {
				if text != "" {
					return fmt.Errorf("expected RequestBody %q, got <nil>", text)
				}
				return nil
			}
			if string(events[i].RequestObject.Raw) != text {
				return fmt.Errorf("expected RequestBody %q, got %q", text, string(events[i].RequestObject.Raw))
			}
			return nil
		}
	}
	requestBodyMatches := func(i int, pattern string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if matched, _ := regexp.Match(pattern, events[i].RequestObject.Raw); !matched {
				return fmt.Errorf("expected RequestBody to match %q, but didn't: %q", pattern, string(events[i].RequestObject.Raw))
			}
			return nil
		}
	}
	noResponseBody := func(i int) eventCheck {
		return func(events []*auditinternal.Event) error {
			if events[i].ResponseObject == nil {
				return nil
			}
			return fmt.Errorf("expected ResponseBody to be nil, got non-nil '%s'", events[i].ResponseObject.Raw)
		}
	}
	responseBodyMatches := func(i int, pattern string) eventCheck {
		return func(events []*auditinternal.Event) error {
			if matched, _ := regexp.Match(pattern, events[i].ResponseObject.Raw); !matched {
				return fmt.Errorf("expected ResponseBody to match %q, but didn't: %q", pattern, string(events[i].ResponseObject.Raw))
			}
			return nil
		}
	}
	requestUserAgentMatches := func(userAgent string) eventCheck {
		return func(events []*auditinternal.Event) error {
			for i := range events {
				if events[i].UserAgent != userAgent {
					return fmt.Errorf("expected request user agent to match %q, but got: %q", userAgent, events[i].UserAgent)
				}
			}
			return nil
		}
	}

	for _, test := range []struct {
		desc   string
		req    func(server string) (*http.Request, error)
		linker runtime.SelfLinker
		code   int
		events int
		checks []eventCheck
	}{
		{
			"get",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"name":"c".*}`),
			},
		},
		{
			"list",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple?labelSelector=a%3Dfoobar", nil)
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple",
				namespace:   "other",
			},
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"name":"a".*"name":"b".*}`),
			},
		},
		{
			"create",
			func(server string) (*http.Request, error) {
				return http.NewRequest("POST", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			201,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleFooJSON)),
				responseBodyMatches(0, `{.*"foo".*}`),
			},
		},
		{
			"not-allowed-named-create",
			func(server string) (*http.Request, error) {
				return http.NewRequest("POST", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/named", bytes.NewBuffer(simpleFooJSON))
			},
			selfLinker,
			405,
			2,
			[]eventCheck{
				noRequestBody(0),  // the 405 is thrown long before the create handler would be executed
				noResponseBody(0), // the 405 is thrown long before the create handler would be executed
			},
		},
		{
			"delete",
			func(server string) (*http.Request, error) {
				return http.NewRequest("DELETE", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/a", nil)
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				noRequestBody(0),
				responseBodyMatches(0, `{.*"kind":"Status".*"status":"Success".*}`),
			},
		},
		{
			"delete-with-options-in-body",
			func(server string) (*http.Request, error) {
				return http.NewRequest("DELETE", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/a", bytes.NewBuffer([]byte(`{"kind":"DeleteOptions"}`)))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				requestBodyMatches(0, "DeleteOptions"),
				responseBodyMatches(0, `{.*"kind":"Status".*"status":"Success".*}`),
			},
		},
		{
			"update",
			func(server string) (*http.Request, error) {
				return http.NewRequest("PUT", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewBuffer(simpleCPrimeJSON))
			},
			selfLinker,
			200,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleCPrimeJSON)),
				responseBodyMatches(0, `{.*"bla".*}`),
			},
		},
		{
			"update-wrong-namespace",
			func(server string) (*http.Request, error) {
				return http.NewRequest("PUT", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/default/simple/c", bytes.NewBuffer(simpleCPrimeJSON))
			},
			selfLinker,
			400,
			2,
			[]eventCheck{
				requestBodyIs(0, string(simpleCPrimeJSON)),
				responseBodyMatches(0, `"Status".*"status":"Failure".*"code":400}`),
			},
		},
		{
			"patch",
			func(server string) (*http.Request, error) {
				req, _ := http.NewRequest("PATCH", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple/c", bytes.NewReader([]byte(`{"labels":{"foo":"bar"}}`)))
				req.Header.Set("Content-Type", "application/merge-patch+json; charset=UTF-8")
				return req, nil
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple/c",
				name:        "c",
				namespace:   "other",
			},
			200,
			2,
			[]eventCheck{
				requestBodyIs(0, `{"labels":{"foo":"bar"}}`),
				responseBodyMatches(0, `"name":"c".*"labels":{"foo":"bar"}`),
			},
		},
		{
			"watch",
			func(server string) (*http.Request, error) {
				return http.NewRequest("GET", server+"/"+prefix+"/"+testGroupVersion.Group+"/"+testGroupVersion.Version+"/namespaces/other/simple?watch=true", nil)
			},
			&setTestSelfLinker{
				t:           t,
				expectedSet: "/" + prefix + "/" + testGroupVersion.Group + "/" + testGroupVersion.Version + "/namespaces/other/simple",
				namespace:   "other",
			},
			200,
			3,
			[]eventCheck{
				noRequestBody(0),
				noResponseBody(0),
			},
		},
	} {
		sink := &fakeAuditSink{}
		handler := handleInternal(map[string]rest.Storage{
			"simple": &SimpleRESTStorage{
				list: []genericapitesting.Simple{
					{
						ObjectMeta: metav1.ObjectMeta{Name: "a", Namespace: "other"},
						Other:      "foo",
					},
					{
						ObjectMeta: metav1.ObjectMeta{Name: "b", Namespace: "other"},
						Other:      "foo",
					},
				},
				item: genericapitesting.Simple{
					ObjectMeta: metav1.ObjectMeta{Name: "c", Namespace: "other", UID: "uid"},
					Other:      "foo",
				},
			},
		}, admissionControl, selfLinker, sink)

		server := httptest.NewServer(handler)
		defer server.Close()
		client := http.Client{Timeout: 2 * time.Second}

		req, err := test.req(server.URL)
		if err != nil {
			t.Errorf("[%s] error creating the request: %v", test.desc, err)
		}

		req.Header.Set("User-Agent", userAgent)

		response, err := client.Do(req)
		if err != nil {
			t.Errorf("[%s] error: %v", test.desc, err)
		}

		if response.StatusCode != test.code {
			t.Errorf("[%s] expected http code %d, got %#v", test.desc, test.code, response)
		}

		// close body because the handler might block in Flush, unable to send the remaining event.
		response.Body.Close()

		// wait for events to arrive, at least the given number in the test
		events := []*auditinternal.Event{}
		err = wait.Poll(50*time.Millisecond, wait.ForeverTestTimeout, wait.ConditionFunc(func() (done bool, err error) {
			events = sink.Events()
			return len(events) >= test.events, nil
		}))
		if err != nil {
			t.Errorf("[%s] timeout waiting for events", test.desc)
		}

		if got := len(events); got != test.events {
			t.Errorf("[%s] expected %d audit events, got %d", test.desc, test.events, got)
		} else {
			for i, check := range test.checks {
				err := check(events)
				if err != nil {
					t.Errorf("[%s,%d] %v", test.desc, i, err)
				}
			}

			if err := requestUserAgentMatches(userAgent)(events); err != nil {
				t.Errorf("[%s] %v", test.desc, err)
			}
		}

		if len(events) > 0 {
			status := events[len(events)-1].ResponseStatus
			if status == nil {
				t.Errorf("[%s] expected non-nil ResponseStatus in last event", test.desc)
			} else if int(status.Code) != test.code {
				t.Errorf("[%s] expected ResponseStatus.Code=%d, got %d", test.desc, test.code, status.Code)
			}
		}
	}
}
<|endoftext|>


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

package fsm

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

var (
	templateReplaceCaptureRE = regexp.MustCompile(`\$\{?([a-zA-Z0-9_\$]+)\}?`)
)

type TemplateFormatter struct {
	captureIndexes []int
	captureCount   int
	fmtString      string
}

// NewTemplateFormatter instantialize a TemplateFormatter
// from given template string and the maxium amount of captures.
func NewTemplateFormatter(template string, captureCount int) *TemplateFormatter {
	matches := templateReplaceCaptureRE.FindAllStringSubmatch(template, -1)
	if len(matches) == 0 {
		// if no regex reference found, keep it as it is
		return &TemplateFormatter{captureCount: 0, fmtString: template}
	}

	var indexes []int
	valueFormatter := template
	for _, match := range matches {
		idx, err := strconv.Atoi(match[len(match)-1])
		if err != nil || idx > captureCount || idx < 1 {
			// if index larger than captured count or using unsupported named capture group,
			// replace with empty string
			valueFormatter = strings.Replace(valueFormatter, match[0], "", -1)
		} else {
			valueFormatter = strings.Replace(valueFormatter, match[0], "%s", -1)
			// note: the regex reference variable $? starts from 1
			indexes = append(indexes, idx-1)
		}
	}
	return &TemplateFormatter{
		captureIndexes: indexes,
		captureCount:   len(indexes),
		fmtString:      valueFormatter,
	}
}

// Format accepts a list containing captured strings and return the formatted string
// using the template stored in current TemplateFormatter.
func (formatter *TemplateFormatter) Format(captures []string) string {
	if formatter.captureCount == 0 {
		// no label substitution, keep as it is
		return formatter.fmtString
	}
	indexes := formatter.captureIndexes
	vargs := make([]interface{}, formatter.captureCount)
	for i, idx := range indexes {
		vargs[i] = captures[idx]
	}
	return fmt.Sprintf(formatter.fmtString, vargs...)
}
<commit_msg>Copy edit comments in fsm/formatter<commit_after>// Copyright 2018 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fsm

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

var (
	templateReplaceCaptureRE = regexp.MustCompile(`\$\{?([a-zA-Z0-9_\$]+)\}?`)
)

type TemplateFormatter struct {
	captureIndexes []int
	captureCount   int
	fmtString      string
}

// NewTemplateFormatter instantiates a TemplateFormatter
// from given template string and the maximum amount of captures.
func NewTemplateFormatter(template string, captureCount int) *TemplateFormatter {
	matches := templateReplaceCaptureRE.FindAllStringSubmatch(template, -1)
	if len(matches) == 0 {
		// if no regex reference found, keep it as it is
		return &TemplateFormatter{captureCount: 0, fmtString: template}
	}

	var indexes []int
	valueFormatter := template
	for _, match := range matches {
		idx, err := strconv.Atoi(match[len(match)-1])
		if err != nil || idx > captureCount || idx < 1 {
			// if index larger than captured count or using unsupported named capture group,
			// replace with empty string
			valueFormatter = strings.Replace(valueFormatter, match[0], "", -1)
		} else {
			valueFormatter = strings.Replace(valueFormatter, match[0], "%s", -1)
			// note: the regex reference variable $? starts from 1
			indexes = append(indexes, idx-1)
		}
	}
	return &TemplateFormatter{
		captureIndexes: indexes,
		captureCount:   len(indexes),
		fmtString:      valueFormatter,
	}
}

// Format accepts a list containing captured strings and returns the formatted
// string using the template stored in current TemplateFormatter.
func (formatter *TemplateFormatter) Format(captures []string) string {
	if formatter.captureCount == 0 {
		// no label substitution, keep as it is
		return formatter.fmtString
	}
	indexes := formatter.captureIndexes
	vargs := make([]interface{}, formatter.captureCount)
	for i, idx := range indexes {
		vargs[i] = captures[idx]
	}
	return fmt.Sprintf(formatter.fmtString, vargs...)
}
<|endoftext|>


<commit_before>package game

import (
	"log"
	"math"

	"buildblast/physics"
	"buildblast/coords"
)

type ControlState struct {
	Forward         bool
	Left            bool
	Right           bool
	Back            bool
	Jump            bool
	ActivateLeft    bool
	ActivateRight   bool
	Lat             float64
	Lon             float64

	Timestamp       float64 // In ms
}

var PLAYER_HEIGHT = 1.75;
var PLAYER_EYE_HEIGHT = 1.6;
var PLAYER_BODY_HEIGHT = 1.3;
var PLAYER_HALF_EXTENTS = coords.Vec3{
	0.2,
	PLAYER_HEIGHT / 2,
	0.2,
};
var PLAYER_CENTER_OFFSET = coords.Vec3{
	0,
	PLAYER_BODY_HEIGHT/2 - PLAYER_EYE_HEIGHT,
	0,
};

// Gameplay state defaults
var PLAYER_MAX_HP = 100;

type Player struct {
	pos       coords.World
	look      coords.Vec3
	vy        float64
	box       physics.Box
	controls  ControlState
	history   *PlayerHistory
	world     *World
	name      string

	// Gameplay state
	hp        int

	// Inventory
	inventory []Item
	itemLeft  int
	itemRight int
}

func NewPlayer(world *World, name string) *Player {
	// 0 -> (w * h) - 1 are the bag slots, 2 more for left equip, and 2 more again for right.
	// [0, w*h - 1] = bag slots
	// [w*h]        = left equip
	// [w*h + 1]    = left reserve
	// [w*h + 2]    = right equip
	// [w*h + 3]    = right reserve
	inv := make([]Item, INV_WIDTH * INV_HEIGHT + 4)
	inv[0] = NewItem(ITEM_GUN)
	inv[1] = NewItem(ITEM_SHOVEL)
	inv[2] = NewItem(ITEM_DIRT)
	inv[3] = NewItem(ITEM_STONE)
	inv[INV_WIDTH * INV_HEIGHT] = NewItem(ITEM_GUN)
	inv[INV_WIDTH * INV_HEIGHT + 2] = NewItem(ITEM_SHOVEL)
	inv[INV_WIDTH * INV_HEIGHT + 3] = NewItem(ITEM_DIRT)

	return &Player{
		pos: world.generator.Spawn(),
		history: NewPlayerHistory(),
		hp: PLAYER_MAX_HP,
		inventory: inv,
		itemLeft: INV_WIDTH * INV_HEIGHT,
		itemRight: INV_WIDTH * INV_HEIGHT + 2,
		world: world,
		name: name,
	}
}

func (p *Player) Pos() coords.World {
	return p.pos
}

func (p *Player) ID() string {
	return "player-" + p.name
}

func (p *Player) Inventory() []Item {
	return p.inventory
}

func (p *Player) Tick(w *World) {}

func (p *Player) SetActiveItems(left, right int) {
	p.itemLeft = left;
	p.itemRight = right;
}

func (p *Player) MoveItems(from, to int) []Item {
	temp := p.inventory[from]
	p.inventory[from] = p.inventory[to]
	p.inventory[to] = temp
	return p.inventory
}

func (p *Player) AddItem(kind byte) {
	firstOpenSpace := -1;
	// Find the item
	for i, item := range p.inventory {
		if firstOpenSpace < 0 && item.kind == ITEM_NIL {
			firstOpenSpace = i
		}
		if item.kind == kind && item.num < MAX_STACK {
			p.inventory[i].num++
			return
		}
	}
	// TODO: Handle no space left
	if firstOpenSpace >= 0 {
		p.inventory[firstOpenSpace] = NewItem(kind)
	}
}

func (p *Player) RemoveItem(kind byte) {
	// Find the item
	for i, item := range p.inventory {
		if item.kind == kind {
			if item.num == 1 {
				p.inventory[i] = NewItem(ITEM_NIL)
				return
			}
			p.inventory[i].num--
			return
		}
	}
}

func (p *Player) ClientTick(controls ControlState) (coords.World, float64, int, *coords.World) {
	dt := (controls.Timestamp - p.controls.Timestamp) / 1000

	if dt > 1.0 {
		log.Println("WARN: Attempt to simulate step with dt of ", dt, " which is too large. Clipping to 1.0s")
		dt = 1.0
	}
	if dt < 0.0 {
		log.Println("WARN: Attempting to simulate step with negative dt of ", dt, " this is probably wrong.")
	}

	p.updateLook(controls)

	hitPos := p.simulateBlaster(dt, controls)
	p.simulateMovement(dt, controls)

	p.controls = controls
	p.history.Add(controls.Timestamp, p.pos)

	return p.pos, p.vy, p.hp, hitPos
}

func (p *Player) simulateMovement(dt float64, controls ControlState) {
	p.vy += dt * -9.81

	fw := 0.0
	if controls.Forward {
		fw = 1 * dt * 10
	} else if controls.Back {
		fw = -1 * dt * 10
	}

	rt := 0.0
	if controls.Right {
		rt = 1 * dt * 10
	} else if controls.Left {
		rt = -1 * dt * 10
	}

	cos := math.Cos
	sin := math.Sin

	move := coords.Vec3{
		X: -cos(controls.Lon) * fw + sin(controls.Lon) * rt,
		Y: p.vy * dt,
		Z: -sin(controls.Lon) * fw - cos(controls.Lon) * rt,
	}

	box := p.Box()

	move = box.AttemptMove(p.world, move)

	if (move.Y == 0) {
		if (controls.Jump) {
			p.vy = 6
		} else {
			p.vy = 0
		}
	}

	p.pos.X += move.X
	p.pos.Y += move.Y
	p.pos.Z += move.Z
}

func (p *Player) updateLook(controls ControlState) {
	cos := math.Cos
	sin := math.Sin

	lat := controls.Lat
	lon := controls.Lon

	p.look.X = sin(lat) * cos(lon)
	p.look.Y = cos(lat)
	p.look.Z = sin(lat) * sin(lon)
}

func (p *Player) simulateBlaster(dt float64, controls ControlState) *coords.World {
	shootingLeft := controls.ActivateLeft && p.inventory[p.itemLeft].Shootable()
	shootingRight := controls.ActivateRight && p.inventory[p.itemRight].Shootable()
	if !shootingLeft && !shootingRight {
		return nil
	}

	// They were holding it down last frame
	shootingLeftLast := p.controls.ActivateLeft && p.inventory[p.itemLeft].Shootable()
	shootingRightLast := p.controls.ActivateRight && p.inventory[p.itemRight].Shootable()
	if (shootingLeft && shootingLeftLast) || (shootingRight && shootingRightLast) {
		return nil
	}

	ray := physics.NewRay(p.pos, p.look)
	hitPos, hitEntity := p.world.FindFirstIntersect(p, controls.Timestamp, ray)
	if hitEntity != nil {
		p.world.DamageEntity(p.name, 10, hitEntity)
	}
	return hitPos
}

func (p *Player) Box() *physics.Box {
	return physics.NewBoxOffset(
		p.pos,
		PLAYER_HALF_EXTENTS,
		PLAYER_CENTER_OFFSET)
}

func (p *Player) BoxAt(t float64) *physics.Box {
	return physics.NewBoxOffset(
		p.history.PositionAt(t),
		PLAYER_HALF_EXTENTS,
		PLAYER_CENTER_OFFSET)
}

func (p *Player) Damage(amount int) {
	p.hp -= amount
}

func (p *Player) Dead() bool {
	return p.hp <= 0
}

func (p *Player) Respawn() {
	p.pos = p.world.generator.Spawn()
	p.hp = PLAYER_MAX_HP
	p.history.Clear()
}
<commit_msg>Take the item from your hand<commit_after>package game

import (
	"log"
	"math"

	"buildblast/physics"
	"buildblast/coords"
)

type ControlState struct {
	Forward         bool
	Left            bool
	Right           bool
	Back            bool
	Jump            bool
	ActivateLeft    bool
	ActivateRight   bool
	Lat             float64
	Lon             float64

	Timestamp       float64 // In ms
}

var PLAYER_HEIGHT = 1.75;
var PLAYER_EYE_HEIGHT = 1.6;
var PLAYER_BODY_HEIGHT = 1.3;
var PLAYER_HALF_EXTENTS = coords.Vec3{
	0.2,
	PLAYER_HEIGHT / 2,
	0.2,
};
var PLAYER_CENTER_OFFSET = coords.Vec3{
	0,
	PLAYER_BODY_HEIGHT/2 - PLAYER_EYE_HEIGHT,
	0,
};

// Gameplay state defaults
var PLAYER_MAX_HP = 100;

type Player struct {
	pos       coords.World
	look      coords.Vec3
	vy        float64
	box       physics.Box
	controls  ControlState
	history   *PlayerHistory
	world     *World
	name      string

	// Gameplay state
	hp        int

	// Inventory
	inventory []Item
	itemLeft  int
	itemRight int
}

func NewPlayer(world *World, name string) *Player {
	// 0 -> (w * h) - 1 are the bag slots, 2 more for left equip, and 2 more again for right.
	// [0, w*h - 1] = bag slots
	// [w*h]        = left equip
	// [w*h + 1]    = left reserve
	// [w*h + 2]    = right equip
	// [w*h + 3]    = right reserve
	inv := make([]Item, INV_WIDTH * INV_HEIGHT + 4)
	inv[0] = NewItem(ITEM_GUN)
	inv[1] = NewItem(ITEM_SHOVEL)
	inv[2] = NewItem(ITEM_DIRT)
	inv[3] = NewItem(ITEM_STONE)
	inv[INV_WIDTH * INV_HEIGHT] = NewItem(ITEM_GUN)
	inv[INV_WIDTH * INV_HEIGHT + 2] = NewItem(ITEM_SHOVEL)
	inv[INV_WIDTH * INV_HEIGHT + 3] = NewItem(ITEM_DIRT)

	return &Player{
		pos: world.generator.Spawn(),
		history: NewPlayerHistory(),
		hp: PLAYER_MAX_HP,
		inventory: inv,
		itemLeft: INV_WIDTH * INV_HEIGHT,
		itemRight: INV_WIDTH * INV_HEIGHT + 2,
		world: world,
		name: name,
	}
}

func (p *Player) Pos() coords.World {
	return p.pos
}

func (p *Player) ID() string {
	return "player-" + p.name
}

func (p *Player) Inventory() []Item {
	return p.inventory
}

func (p *Player) Tick(w *World) {}

func (p *Player) SetActiveItems(left, right int) {
	p.itemLeft = left;
	p.itemRight = right;
}

func (p *Player) MoveItems(from, to int) []Item {
	temp := p.inventory[from]
	p.inventory[from] = p.inventory[to]
	p.inventory[to] = temp
	return p.inventory
}

func (p *Player) AddItem(kind byte) {
	firstOpenSpace := -1;
	// Find the item
	for i, item := range p.inventory {
		if firstOpenSpace < 0 && item.kind == ITEM_NIL {
			firstOpenSpace = i
		}
		if item.kind == kind && item.num < MAX_STACK {
			p.inventory[i].num++
			return
		}
	}
	// TODO: Handle no space left
	if firstOpenSpace >= 0 {
		p.inventory[firstOpenSpace] = NewItem(kind)
	}
}

func (p *Player) RemoveItem(kind byte) {
	// Find the item

	// Check your hands first
	if p.inventory[p.itemLeft].kind == kind {
		p.lowerStack(p.itemLeft)
		return
	}
	if p.inventory[p.itemRight].kind == kind {
		p.lowerStack(p.itemRight)
		return
	}

	for i, item := range p.inventory {
		if item.kind == kind {
			p.lowerStack(i)
			return
		}
	}
}

func (p *Player) lowerStack(i int) {
	if p.inventory[i].num == 1 {
		p.inventory[i] = NewItem(ITEM_NIL)
		return
	}
	p.inventory[i].num--
}

func (p *Player) ClientTick(controls ControlState) (coords.World, float64, int, *coords.World) {
	dt := (controls.Timestamp - p.controls.Timestamp) / 1000

	if dt > 1.0 {
		log.Println("WARN: Attempt to simulate step with dt of ", dt, " which is too large. Clipping to 1.0s")
		dt = 1.0
	}
	if dt < 0.0 {
		log.Println("WARN: Attempting to simulate step with negative dt of ", dt, " this is probably wrong.")
	}

	p.updateLook(controls)

	hitPos := p.simulateBlaster(dt, controls)
	p.simulateMovement(dt, controls)

	p.controls = controls
	p.history.Add(controls.Timestamp, p.pos)

	return p.pos, p.vy, p.hp, hitPos
}

func (p *Player) simulateMovement(dt float64, controls ControlState) {
	p.vy += dt * -9.81

	fw := 0.0
	if controls.Forward {
		fw = 1 * dt * 10
	} else if controls.Back {
		fw = -1 * dt * 10
	}

	rt := 0.0
	if controls.Right {
		rt = 1 * dt * 10
	} else if controls.Left {
		rt = -1 * dt * 10
	}

	cos := math.Cos
	sin := math.Sin

	move := coords.Vec3{
		X: -cos(controls.Lon) * fw + sin(controls.Lon) * rt,
		Y: p.vy * dt,
		Z: -sin(controls.Lon) * fw - cos(controls.Lon) * rt,
	}

	box := p.Box()

	move = box.AttemptMove(p.world, move)

	if (move.Y == 0) {
		if (controls.Jump) {
			p.vy = 6
		} else {
			p.vy = 0
		}
	}

	p.pos.X += move.X
	p.pos.Y += move.Y
	p.pos.Z += move.Z
}

func (p *Player) updateLook(controls ControlState) {
	cos := math.Cos
	sin := math.Sin

	lat := controls.Lat
	lon := controls.Lon

	p.look.X = sin(lat) * cos(lon)
	p.look.Y = cos(lat)
	p.look.Z = sin(lat) * sin(lon)
}

func (p *Player) simulateBlaster(dt float64, controls ControlState) *coords.World {
	shootingLeft := controls.ActivateLeft && p.inventory[p.itemLeft].Shootable()
	shootingRight := controls.ActivateRight && p.inventory[p.itemRight].Shootable()
	if !shootingLeft && !shootingRight {
		return nil
	}

	// They were holding it down last frame
	shootingLeftLast := p.controls.ActivateLeft && p.inventory[p.itemLeft].Shootable()
	shootingRightLast := p.controls.ActivateRight && p.inventory[p.itemRight].Shootable()
	if (shootingLeft && shootingLeftLast) || (shootingRight && shootingRightLast) {
		return nil
	}

	ray := physics.NewRay(p.pos, p.look)
	hitPos, hitEntity := p.world.FindFirstIntersect(p, controls.Timestamp, ray)
	if hitEntity != nil {
		p.world.DamageEntity(p.name, 10, hitEntity)
	}
	return hitPos
}

func (p *Player) Box() *physics.Box {
	return physics.NewBoxOffset(
		p.pos,
		PLAYER_HALF_EXTENTS,
		PLAYER_CENTER_OFFSET)
}

func (p *Player) BoxAt(t float64) *physics.Box {
	return physics.NewBoxOffset(
		p.history.PositionAt(t),
		PLAYER_HALF_EXTENTS,
		PLAYER_CENTER_OFFSET)
}

func (p *Player) Damage(amount int) {
	p.hp -= amount
}

func (p *Player) Dead() bool {
	return p.hp <= 0
}

func (p *Player) Respawn() {
	p.pos = p.world.generator.Spawn()
	p.hp = PLAYER_MAX_HP
	p.history.Clear()
}
<|endoftext|>


<commit_before>package proxy

import (
	"github.com/fagongzi/gateway/pkg/filter"
)

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailers",
	"Transfer-Encoding",
}

// HeadersFilter HeadersFilter
type HeadersFilter struct {
	filter.BaseFilter
}

func newHeadersFilter() filter.Filter {
	return &HeadersFilter{}
}

// Init init filter
func (f *HeadersFilter) Init(cfg string) error {
	return nil
}

// Name return name of this filter
func (f *HeadersFilter) Name() string {
	return FilterHeader
}

// Pre execute before proxy
func (f *HeadersFilter) Pre(c filter.Context) (statusCode int, err error) {
	for _, h := range hopHeaders {
		c.ForwardRequest().Header.Del(h)
	}

	c.ForwardRequest().Header.SetHost(c.Server().Addr)
	return f.BaseFilter.Pre(c)
}

// Post execute after proxy
func (f *HeadersFilter) Post(c filter.Context) (statusCode int, err error) {
	for _, h := range hopHeaders {
		c.Response().Header.Del(h)
	}

	// 需要合并处理的，不做header的复制，由proxy做合并
	if len(c.API().Nodes) == 1 {
		c.OriginRequest().Response.Header.Reset()
		c.Response().Header.CopyTo(&c.OriginRequest().Response.Header)
	}

	return f.BaseFilter.Post(c)
}
<commit_msg>fix: keepalive<commit_after>package proxy

import (
	"github.com/fagongzi/gateway/pkg/filter"
)

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te",
	"Trailers",
	"Transfer-Encoding",
}

// HeadersFilter HeadersFilter
type HeadersFilter struct {
	filter.BaseFilter
}

func newHeadersFilter() filter.Filter {
	return &HeadersFilter{}
}

// Init init filter
func (f *HeadersFilter) Init(cfg string) error {
	return nil
}

// Name return name of this filter
func (f *HeadersFilter) Name() string {
	return FilterHeader
}

// Pre execute before proxy
func (f *HeadersFilter) Pre(c filter.Context) (statusCode int, err error) {
	for _, h := range hopHeaders {
		c.ForwardRequest().Header.Del(h)
	}

	c.ForwardRequest().Header.SetHost(c.Server().Addr)
	return f.BaseFilter.Pre(c)
}

// Post execute after proxy
func (f *HeadersFilter) Post(c filter.Context) (statusCode int, err error) {
	for _, h := range hopHeaders {
		c.Response().Header.Del(h)
	}

	// 需要合并处理的，不做header的复制，由proxy做合并
	if len(c.API().Nodes) == 1 {
		c.OriginRequest().Response.Header.Reset()
		c.Response().Header.CopyTo(&c.OriginRequest().Response.Header)
	}

	return f.BaseFilter.Post(c)
}
<|endoftext|>


<commit_before>package system

import (
	"os"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/mem"
	"github.com/shirou/gopsutil/net"
)

type PS interface {
	CPUTimes(perCPU, totalCPU bool) ([]cpu.TimesStat, error)
	DiskUsage(mountPointFilter []string, fstypeExclude []string) ([]*disk.UsageStat, []*disk.PartitionStat, error)
	NetIO() ([]net.IOCountersStat, error)
	NetProto() ([]net.ProtoCountersStat, error)
	DiskIO(names []string) (map[string]disk.IOCountersStat, error)
	VMStat() (*mem.VirtualMemoryStat, error)
	SwapStat() (*mem.SwapMemoryStat, error)
	NetConnections() ([]net.ConnectionStat, error)
}

type PSDiskDeps interface {
	Partitions(all bool) ([]disk.PartitionStat, error)
	OSGetenv(key string) string
	OSStat(name string) (os.FileInfo, error)
	PSDiskUsage(path string) (*disk.UsageStat, error)
}

func add(acc telegraf.Accumulator,
	name string, val float64, tags map[string]string) {
	if val >= 0 {
		acc.AddFields(name, map[string]interface{}{"value": val}, tags)
	}
}

func newSystemPS() *systemPS {
	return &systemPS{&systemPSDisk{}}
}

type systemPS struct {
	PSDiskDeps
}

type systemPSDisk struct{}

func (s *systemPS) CPUTimes(perCPU, totalCPU bool) ([]cpu.TimesStat, error) {
	var cpuTimes []cpu.TimesStat
	if perCPU {
		if perCPUTimes, err := cpu.Times(true); err == nil {
			cpuTimes = append(cpuTimes, perCPUTimes...)
		} else {
			return nil, err
		}
	}
	if totalCPU {
		if totalCPUTimes, err := cpu.Times(false); err == nil {
			cpuTimes = append(cpuTimes, totalCPUTimes...)
		} else {
			return nil, err
		}
	}
	return cpuTimes, nil
}

func (s *systemPS) DiskUsage(
	mountPointFilter []string,
	fstypeExclude []string,
) ([]*disk.UsageStat, []*disk.PartitionStat, error) {
	parts, err := s.Partitions(true)
	if err != nil {
		return nil, nil, err
	}

	// Make a "set" out of the filter slice
	mountPointFilterSet := make(map[string]bool)
	for _, filter := range mountPointFilter {
		mountPointFilterSet[filter] = true
	}
	fstypeExcludeSet := make(map[string]bool)
	for _, filter := range fstypeExclude {
		fstypeExcludeSet[filter] = true
	}

	var usage []*disk.UsageStat
	var partitions []*disk.PartitionStat

	for i := range parts {
		p := parts[i]

		if len(mountPointFilter) > 0 {
			// If the mount point is not a member of the filter set,
			// don't gather info on it.
			if _, ok := mountPointFilterSet[p.Mountpoint]; !ok {
				continue
			}
		}

		// If the mount point is a member of the exclude set,
		// don't gather info on it.
		if _, ok := fstypeExcludeSet[p.Fstype]; ok {
			continue
		}

		mountpoint := s.OSGetenv("HOST_MOUNT_PREFIX") + p.Mountpoint
		if _, err := s.OSStat(mountpoint); err != nil {
			continue
		}
		du, err := s.PSDiskUsage(mountpoint)
		if err != nil {
			continue
		}
		du.Path = p.Mountpoint
		du.Fstype = p.Fstype
		usage = append(usage, du)
		partitions = append(partitions, &p)
	}

	return usage, partitions, nil
}

func (s *systemPS) NetProto() ([]net.ProtoCountersStat, error) {
	return net.ProtoCounters(nil)
}

func (s *systemPS) NetIO() ([]net.IOCountersStat, error) {
	return net.IOCounters(true)
}

func (s *systemPS) NetConnections() ([]net.ConnectionStat, error) {
	return net.Connections("all")
}

func (s *systemPS) DiskIO(names []string) (map[string]disk.IOCountersStat, error) {
	m, err := disk.IOCounters(names...)
	if err == internal.NotImplementedError {
		return nil, nil
	}

	return m, err
}

func (s *systemPS) VMStat() (*mem.VirtualMemoryStat, error) {
	return mem.VirtualMemory()
}

func (s *systemPS) SwapStat() (*mem.SwapMemoryStat, error) {
	return mem.SwapMemory()
}

func (s *systemPSDisk) Partitions(all bool) ([]disk.PartitionStat, error) {
	return disk.Partitions(all)
}

func (s *systemPSDisk) OSGetenv(key string) string {
	return os.Getenv(key)
}

func (s *systemPSDisk) OSStat(name string) (os.FileInfo, error) {
	return os.Stat(name)
}

func (s *systemPSDisk) PSDiskUsage(path string) (*disk.UsageStat, error) {
	return disk.Usage(path)
}
<commit_msg>Fix a bug getting disk usage inside docker<commit_after>package system

import (
	"os"
	"strings"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/mem"
	"github.com/shirou/gopsutil/net"
)

type PS interface {
	CPUTimes(perCPU, totalCPU bool) ([]cpu.TimesStat, error)
	DiskUsage(mountPointFilter []string, fstypeExclude []string) ([]*disk.UsageStat, []*disk.PartitionStat, error)
	NetIO() ([]net.IOCountersStat, error)
	NetProto() ([]net.ProtoCountersStat, error)
	DiskIO(names []string) (map[string]disk.IOCountersStat, error)
	VMStat() (*mem.VirtualMemoryStat, error)
	SwapStat() (*mem.SwapMemoryStat, error)
	NetConnections() ([]net.ConnectionStat, error)
}

type PSDiskDeps interface {
	Partitions(all bool) ([]disk.PartitionStat, error)
	OSGetenv(key string) string
	OSStat(name string) (os.FileInfo, error)
	PSDiskUsage(path string) (*disk.UsageStat, error)
}

func add(acc telegraf.Accumulator,
	name string, val float64, tags map[string]string) {
	if val >= 0 {
		acc.AddFields(name, map[string]interface{}{"value": val}, tags)
	}
}

func newSystemPS() *systemPS {
	return &systemPS{&systemPSDisk{}}
}

type systemPS struct {
	PSDiskDeps
}

type systemPSDisk struct{}

func (s *systemPS) CPUTimes(perCPU, totalCPU bool) ([]cpu.TimesStat, error) {
	var cpuTimes []cpu.TimesStat
	if perCPU {
		if perCPUTimes, err := cpu.Times(true); err == nil {
			cpuTimes = append(cpuTimes, perCPUTimes...)
		} else {
			return nil, err
		}
	}
	if totalCPU {
		if totalCPUTimes, err := cpu.Times(false); err == nil {
			cpuTimes = append(cpuTimes, totalCPUTimes...)
		} else {
			return nil, err
		}
	}
	return cpuTimes, nil
}

func (s *systemPS) DiskUsage(
	mountPointFilter []string,
	fstypeExclude []string,
) ([]*disk.UsageStat, []*disk.PartitionStat, error) {
	parts, err := s.Partitions(true)
	if err != nil {
		return nil, nil, err
	}

	// Make a "set" out of the filter slice
	mountPointFilterSet := make(map[string]bool)
	for _, filter := range mountPointFilter {
		mountPointFilterSet[filter] = true
	}
	fstypeExcludeSet := make(map[string]bool)
	for _, filter := range fstypeExclude {
		fstypeExcludeSet[filter] = true
	}

	var usage []*disk.UsageStat
	var partitions []*disk.PartitionStat
	hostMountPrefix := s.OSGetenv("HOST_MOUNT_PREFIX")

	for i := range parts {
		p := parts[i]

		if len(mountPointFilter) > 0 {
			// If the mount point is not a member of the filter set,
			// don't gather info on it.
			if _, ok := mountPointFilterSet[p.Mountpoint]; !ok {
				continue
			}
		}

		// If there's a host mount prefix, exclude other paths
		if len(hostMountPrefix) > 0 && !strings.HasPrefix(p.Mountpoint, hostMountPrefix) {
			continue
		}

		// If the mount point is a member of the exclude set,
		// don't gather info on it.
		if _, ok := fstypeExcludeSet[p.Fstype]; ok {
			continue
		}

		if _, err := s.OSStat(p.Mountpoint); err != nil {
			continue
		}
		du, err := s.PSDiskUsage(p.Mountpoint)
		if err != nil {
			continue
		}
		du.Path = p.Mountpoint
		du.Fstype = p.Fstype
		usage = append(usage, du)
		partitions = append(partitions, &p)
	}

	return usage, partitions, nil
}

func (s *systemPS) NetProto() ([]net.ProtoCountersStat, error) {
	return net.ProtoCounters(nil)
}

func (s *systemPS) NetIO() ([]net.IOCountersStat, error) {
	return net.IOCounters(true)
}

func (s *systemPS) NetConnections() ([]net.ConnectionStat, error) {
	return net.Connections("all")
}

func (s *systemPS) DiskIO(names []string) (map[string]disk.IOCountersStat, error) {
	m, err := disk.IOCounters(names...)
	if err == internal.NotImplementedError {
		return nil, nil
	}

	return m, err
}

func (s *systemPS) VMStat() (*mem.VirtualMemoryStat, error) {
	return mem.VirtualMemory()
}

func (s *systemPS) SwapStat() (*mem.SwapMemoryStat, error) {
	return mem.SwapMemory()
}

func (s *systemPSDisk) Partitions(all bool) ([]disk.PartitionStat, error) {
	return disk.Partitions(all)
}

func (s *systemPSDisk) OSGetenv(key string) string {
	return os.Getenv(key)
}

func (s *systemPSDisk) OSStat(name string) (os.FileInfo, error) {
	return os.Stat(name)
}

func (s *systemPSDisk) PSDiskUsage(path string) (*disk.UsageStat, error) {
	return disk.Usage(path)
}
<|endoftext|>


<commit_before>package duktape

/*
#cgo linux LDFLAGS: -lm

# include "duktape.h"
extern duk_ret_t goFinalize(duk_context *ctx);
extern duk_ret_t goCall(duk_context *ctx);
*/
import "C"
import "sync"
import "errors"
import "unsafe"

const goFuncProp = "goFuncData"
const goObjProp = "goObjData"
const (
	DUK_ENUM_INCLUDE_NONENUMERABLE = C.DUK_ENUM_INCLUDE_NONENUMERABLE
	DUK_ENUM_INCLUDE_INTERNAL      = C.DUK_ENUM_INCLUDE_INTERNAL
	DUK_ENUM_OWN_PROPERTIES_ONLY   = C.DUK_ENUM_OWN_PROPERTIES_ONLY
	DUK_ENUM_ARRAY_INDICES_ONLY    = C.DUK_ENUM_ARRAY_INDICES_ONLY
	DUK_ENUM_SORT_ARRAY_INDICES    = C.DUK_ENUM_SORT_ARRAY_INDICES
	DUK_ENUM_NO_PROXY_BEHAVIOR     = C.DUK_ENUM_NO_PROXY_BEHAVIOR
)

const (
	DUK_RET_UNIMPLEMENTED_ERROR = C.DUK_RET_UNIMPLEMENTED_ERROR
	DUK_RET_UNSUPPORTED_ERROR   = C.DUK_RET_UNSUPPORTED_ERROR
	DUK_RET_INTERNAL_ERROR      = C.DUK_RET_INTERNAL_ERROR
	DUK_RET_ALLOC_ERROR         = C.DUK_RET_ALLOC_ERROR
	DUK_RET_ASSERTION_ERROR     = C.DUK_RET_ASSERTION_ERROR
	DUK_RET_API_ERROR           = C.DUK_RET_API_ERROR
	DUK_RET_UNCAUGHT_ERROR      = C.DUK_RET_UNCAUGHT_ERROR
	DUK_RET_ERROR               = C.DUK_RET_ERROR
	DUK_RET_EVAL_ERROR          = C.DUK_RET_EVAL_ERROR
	DUK_RET_RANGE_ERROR         = C.DUK_RET_RANGE_ERROR
	DUK_RET_REFERENCE_ERROR     = C.DUK_RET_REFERENCE_ERROR
	DUK_RET_SYNTAX_ERROR        = C.DUK_RET_SYNTAX_ERROR
	DUK_RET_TYPE_ERROR          = C.DUK_RET_TYPE_ERROR
	DUK_RET_URI_ERROR           = C.DUK_RET_URI_ERROR
)

const (
	DUK_TYPE_NONE Type = iota
	DUK_TYPE_UNDEFINED
	DUK_TYPE_NULL
	DUK_TYPE_BOOLEAN
	DUK_TYPE_NUMBER
	DUK_TYPE_STRING
	DUK_TYPE_OBJECT
	DUK_TYPE_BUFFER
	DUK_TYPE_POINTER
)

const (
	DUK_ERR_UNIMPLEMENTED_ERROR = C.DUK_ERR_UNIMPLEMENTED_ERROR
	DUK_ERR_UNSUPPORTED_ERROR   = C.DUK_ERR_UNSUPPORTED_ERROR
	DUK_ERR_INTERNAL_ERROR      = C.DUK_ERR_INTERNAL_ERROR
	DUK_ERR_ALLOC_ERROR         = C.DUK_ERR_ALLOC_ERROR
	DUK_ERR_ASSERTION_ERROR     = C.DUK_ERR_ASSERTION_ERROR
	DUK_ERR_API_ERROR           = C.DUK_ERR_API_ERROR
	DUK_ERR_UNCAUGHT_ERROR      = C.DUK_ERR_UNCAUGHT_ERROR
	DUK_ERR_ERROR               = C.DUK_ERR_ERROR
	DUK_ERR_EVAL_ERROR          = C.DUK_ERR_EVAL_ERROR
	DUK_ERR_RANGE_ERROR         = C.DUK_ERR_RANGE_ERROR
	DUK_ERR_REFERENCE_ERROR     = C.DUK_ERR_REFERENCE_ERROR
	DUK_ERR_SYNTAX_ERROR        = C.DUK_ERR_SYNTAX_ERROR
	DUK_ERR_TYPE_ERROR          = C.DUK_ERR_TYPE_ERROR
	DUK_ERR_URI_ERROR           = C.DUK_ERR_URI_ERROR
)

type Type int

func (t Type) IsNone() bool      { return t == DUK_TYPE_NONE }
func (t Type) IsUndefined() bool { return t == DUK_TYPE_UNDEFINED }
func (t Type) IsNull() bool      { return t == DUK_TYPE_NULL }
func (t Type) IsBool() bool      { return t == DUK_TYPE_BOOLEAN }
func (t Type) IsNumber() bool    { return t == DUK_TYPE_NUMBER }
func (t Type) IsString() bool    { return t == DUK_TYPE_STRING }
func (t Type) IsObject() bool    { return t == DUK_TYPE_OBJECT }
func (t Type) IsBuffer() bool    { return t == DUK_TYPE_BUFFER }
func (t Type) IsPointer() bool   { return t == DUK_TYPE_POINTER }

var objectMutex sync.Mutex
var objectMap map[unsafe.Pointer]interface{} = make(map[unsafe.Pointer]interface{})

type Context struct {
	duk_context unsafe.Pointer
}

// Returns initialized duktape context object
func NewContext() *Context {
	ctx := &Context{
		// TODO: "A caller SHOULD implement a fatal error handler in most applications."
		duk_context: C.duk_create_heap(nil, nil, nil, nil, nil),
	}
	return ctx
}

func (d *Context) PutInternalPropString(objIndex int, key string) bool {
	cKey := C.CString("\xff" + key) // \xff as the first char designates an internal property
	defer C.free(unsafe.Pointer(cKey))
	return int(C.duk_put_prop_string(d.duk_context, C.duk_idx_t(objIndex), cKey)) == 1
}

func (d *Context) GetInternalPropString(objIndex int, key string) bool {
	cKey := C.CString("\xff" + key) // \xff as the first char designates an internal property
	defer C.free(unsafe.Pointer(cKey))
	return int(C.duk_get_prop_string(d.duk_context, C.duk_idx_t(objIndex), cKey)) == 1
}

//export goFinalize
func goFinalize(ctx unsafe.Pointer) C.duk_ret_t {
	d := &Context{ctx}
	d.PushCurrentFunction()
	d.GetInternalPropString(-1, goFuncProp)
	if !Type(d.GetType(-1)).IsPointer() {
		d.Pop2()
		return C.duk_ret_t(C.DUK_RET_TYPE_ERROR)
	}
	key := d.GetPointer(-1)
	d.Pop2()
	objectMutex.Lock()
	delete(objectMap, key)
	objectMutex.Unlock()
	C.free(key)
	return C.duk_ret_t(0)
}

func (d *Context) putGoObjectRef(prop string, o interface{}) {
	key := C.malloc(1) // guaranteed to be unique until freed

	objectMutex.Lock()
	objectMap[key] = o
	objectMutex.Unlock()

	d.PushCFunction((*[0]byte)(C.goFinalize), 1)
	d.PushPointer(key)
	d.PutInternalPropString(-2, goFuncProp)

	d.SetFinalizer(-2)

	d.PushPointer(key)
	d.PutInternalPropString(-2, prop)
}

func (d *Context) PushGoObject(o interface{}) {
	d.PushObject()
	d.putGoObjectRef(goObjProp, o)
}

func (d *Context) getGoObjectRef(objIndex int, prop string) interface{} {
	d.GetInternalPropString(objIndex, prop)
	if !Type(d.GetType(-1)).IsPointer() {
		d.Pop()
		return nil
	}
	key := d.GetPointer(-1)
	d.Pop()
	objectMutex.Lock()
	defer objectMutex.Unlock()
	return objectMap[key]
}

func (d *Context) GetGoObject(objIndex int) interface{} {
	return d.getGoObjectRef(objIndex, goObjProp)
}

//export goCall
func goCall(ctx unsafe.Pointer) C.duk_ret_t {
	d := &Context{ctx}

	/*
		d.PushContextDump()
		log.Printf("goCall context: %s", d.GetString(-1))
		d.Pop()
	*/

	d.PushCurrentFunction()
	if fd, _ := d.getGoObjectRef(-1, goFuncProp).(*GoFuncData); fd == nil {
		d.Pop()
		return C.duk_ret_t(C.DUK_RET_TYPE_ERROR)
	} else {
		d.Pop()
		return C.duk_ret_t(fd.f(d))
	}
}

type GoFunc func(d *Context) int
type GoFuncData struct {
	f GoFunc
}

// Push goCall with its "goFuncData" property set to fd
func (d *Context) PushGoFunc(f GoFunc) {
	fd := &GoFuncData{f}
	d.PushCFunction((*[0]byte)(C.goCall), C.DUK_VARARGS)
	d.putGoObjectRef(goFuncProp, fd)
}

type MethodSuite map[string]GoFunc

func (d *Context) EvalWith(source string, suite MethodSuite) error {
	if err := d.PevalString(source); err != 0 {
		return errors.New(d.SafeToString(-1))
	}

	d.PushObject()

	for prop, f := range suite {
		d.PushGoFunc(f)
		d.PutPropString(-2, prop)
	}

	if err := d.Pcall(1); err != 0 {
		return errors.New(d.SafeToString(-1))
	}

	return nil
}

// TBD: panic handling.
// When a goroutine panics, mark the context as panicking and throw a special JS
// exception, saving the value passed to panic() to the context.
// When the exception is caught, retrieve the value that was passed to panic(),
// and call panic() again.
<commit_msg>add DUK_COMPILE_* flags<commit_after>package duktape

/*
#cgo linux LDFLAGS: -lm

# include "duktape.h"
extern duk_ret_t goFinalize(duk_context *ctx);
extern duk_ret_t goCall(duk_context *ctx);
*/
import "C"
import "sync"
import "errors"
import "unsafe"

const goFuncProp = "goFuncData"
const goObjProp = "goObjData"
const (
	DUK_ENUM_INCLUDE_NONENUMERABLE = C.DUK_ENUM_INCLUDE_NONENUMERABLE
	DUK_ENUM_INCLUDE_INTERNAL      = C.DUK_ENUM_INCLUDE_INTERNAL
	DUK_ENUM_OWN_PROPERTIES_ONLY   = C.DUK_ENUM_OWN_PROPERTIES_ONLY
	DUK_ENUM_ARRAY_INDICES_ONLY    = C.DUK_ENUM_ARRAY_INDICES_ONLY
	DUK_ENUM_SORT_ARRAY_INDICES    = C.DUK_ENUM_SORT_ARRAY_INDICES
	DUK_ENUM_NO_PROXY_BEHAVIOR     = C.DUK_ENUM_NO_PROXY_BEHAVIOR
)

const (
	DUK_RET_UNIMPLEMENTED_ERROR = C.DUK_RET_UNIMPLEMENTED_ERROR
	DUK_RET_UNSUPPORTED_ERROR   = C.DUK_RET_UNSUPPORTED_ERROR
	DUK_RET_INTERNAL_ERROR      = C.DUK_RET_INTERNAL_ERROR
	DUK_RET_ALLOC_ERROR         = C.DUK_RET_ALLOC_ERROR
	DUK_RET_ASSERTION_ERROR     = C.DUK_RET_ASSERTION_ERROR
	DUK_RET_API_ERROR           = C.DUK_RET_API_ERROR
	DUK_RET_UNCAUGHT_ERROR      = C.DUK_RET_UNCAUGHT_ERROR
	DUK_RET_ERROR               = C.DUK_RET_ERROR
	DUK_RET_EVAL_ERROR          = C.DUK_RET_EVAL_ERROR
	DUK_RET_RANGE_ERROR         = C.DUK_RET_RANGE_ERROR
	DUK_RET_REFERENCE_ERROR     = C.DUK_RET_REFERENCE_ERROR
	DUK_RET_SYNTAX_ERROR        = C.DUK_RET_SYNTAX_ERROR
	DUK_RET_TYPE_ERROR          = C.DUK_RET_TYPE_ERROR
	DUK_RET_URI_ERROR           = C.DUK_RET_URI_ERROR
)

const (
	DUK_TYPE_NONE Type = iota
	DUK_TYPE_UNDEFINED
	DUK_TYPE_NULL
	DUK_TYPE_BOOLEAN
	DUK_TYPE_NUMBER
	DUK_TYPE_STRING
	DUK_TYPE_OBJECT
	DUK_TYPE_BUFFER
	DUK_TYPE_POINTER
)

const (
	DUK_COMPILE_EVAL     = C.DUK_COMPILE_EVAL
	DUK_COMPILE_FUNCTION = C.DUK_COMPILE_FUNCTION
	DUK_COMPILE_STRICT   = C.DUK_COMPILE_STRICT
)

const (
	DUK_ERR_UNIMPLEMENTED_ERROR = C.DUK_ERR_UNIMPLEMENTED_ERROR
	DUK_ERR_UNSUPPORTED_ERROR   = C.DUK_ERR_UNSUPPORTED_ERROR
	DUK_ERR_INTERNAL_ERROR      = C.DUK_ERR_INTERNAL_ERROR
	DUK_ERR_ALLOC_ERROR         = C.DUK_ERR_ALLOC_ERROR
	DUK_ERR_ASSERTION_ERROR     = C.DUK_ERR_ASSERTION_ERROR
	DUK_ERR_API_ERROR           = C.DUK_ERR_API_ERROR
	DUK_ERR_UNCAUGHT_ERROR      = C.DUK_ERR_UNCAUGHT_ERROR
	DUK_ERR_ERROR               = C.DUK_ERR_ERROR
	DUK_ERR_EVAL_ERROR          = C.DUK_ERR_EVAL_ERROR
	DUK_ERR_RANGE_ERROR         = C.DUK_ERR_RANGE_ERROR
	DUK_ERR_REFERENCE_ERROR     = C.DUK_ERR_REFERENCE_ERROR
	DUK_ERR_SYNTAX_ERROR        = C.DUK_ERR_SYNTAX_ERROR
	DUK_ERR_TYPE_ERROR          = C.DUK_ERR_TYPE_ERROR
	DUK_ERR_URI_ERROR           = C.DUK_ERR_URI_ERROR
)

type Type int

func (t Type) IsNone() bool      { return t == DUK_TYPE_NONE }
func (t Type) IsUndefined() bool { return t == DUK_TYPE_UNDEFINED }
func (t Type) IsNull() bool      { return t == DUK_TYPE_NULL }
func (t Type) IsBool() bool      { return t == DUK_TYPE_BOOLEAN }
func (t Type) IsNumber() bool    { return t == DUK_TYPE_NUMBER }
func (t Type) IsString() bool    { return t == DUK_TYPE_STRING }
func (t Type) IsObject() bool    { return t == DUK_TYPE_OBJECT }
func (t Type) IsBuffer() bool    { return t == DUK_TYPE_BUFFER }
func (t Type) IsPointer() bool   { return t == DUK_TYPE_POINTER }

var objectMutex sync.Mutex
var objectMap map[unsafe.Pointer]interface{} = make(map[unsafe.Pointer]interface{})

type Context struct {
	duk_context unsafe.Pointer
}

// Returns initialized duktape context object
func NewContext() *Context {
	ctx := &Context{
		// TODO: "A caller SHOULD implement a fatal error handler in most applications."
		duk_context: C.duk_create_heap(nil, nil, nil, nil, nil),
	}
	return ctx
}

func (d *Context) PutInternalPropString(objIndex int, key string) bool {
	cKey := C.CString("\xff" + key) // \xff as the first char designates an internal property
	defer C.free(unsafe.Pointer(cKey))
	return int(C.duk_put_prop_string(d.duk_context, C.duk_idx_t(objIndex), cKey)) == 1
}

func (d *Context) GetInternalPropString(objIndex int, key string) bool {
	cKey := C.CString("\xff" + key) // \xff as the first char designates an internal property
	defer C.free(unsafe.Pointer(cKey))
	return int(C.duk_get_prop_string(d.duk_context, C.duk_idx_t(objIndex), cKey)) == 1
}

//export goFinalize
func goFinalize(ctx unsafe.Pointer) C.duk_ret_t {
	d := &Context{ctx}
	d.PushCurrentFunction()
	d.GetInternalPropString(-1, goFuncProp)
	if !Type(d.GetType(-1)).IsPointer() {
		d.Pop2()
		return C.duk_ret_t(C.DUK_RET_TYPE_ERROR)
	}
	key := d.GetPointer(-1)
	d.Pop2()
	objectMutex.Lock()
	delete(objectMap, key)
	objectMutex.Unlock()
	C.free(key)
	return C.duk_ret_t(0)
}

func (d *Context) putGoObjectRef(prop string, o interface{}) {
	key := C.malloc(1) // guaranteed to be unique until freed

	objectMutex.Lock()
	objectMap[key] = o
	objectMutex.Unlock()

	d.PushCFunction((*[0]byte)(C.goFinalize), 1)
	d.PushPointer(key)
	d.PutInternalPropString(-2, goFuncProp)

	d.SetFinalizer(-2)

	d.PushPointer(key)
	d.PutInternalPropString(-2, prop)
}

func (d *Context) PushGoObject(o interface{}) {
	d.PushObject()
	d.putGoObjectRef(goObjProp, o)
}

func (d *Context) getGoObjectRef(objIndex int, prop string) interface{} {
	d.GetInternalPropString(objIndex, prop)
	if !Type(d.GetType(-1)).IsPointer() {
		d.Pop()
		return nil
	}
	key := d.GetPointer(-1)
	d.Pop()
	objectMutex.Lock()
	defer objectMutex.Unlock()
	return objectMap[key]
}

func (d *Context) GetGoObject(objIndex int) interface{} {
	return d.getGoObjectRef(objIndex, goObjProp)
}

//export goCall
func goCall(ctx unsafe.Pointer) C.duk_ret_t {
	d := &Context{ctx}

	/*
		d.PushContextDump()
		log.Printf("goCall context: %s", d.GetString(-1))
		d.Pop()
	*/

	d.PushCurrentFunction()
	if fd, _ := d.getGoObjectRef(-1, goFuncProp).(*GoFuncData); fd == nil {
		d.Pop()
		return C.duk_ret_t(C.DUK_RET_TYPE_ERROR)
	} else {
		d.Pop()
		return C.duk_ret_t(fd.f(d))
	}
}

type GoFunc func(d *Context) int
type GoFuncData struct {
	f GoFunc
}

// Push goCall with its "goFuncData" property set to fd
func (d *Context) PushGoFunc(f GoFunc) {
	fd := &GoFuncData{f}
	d.PushCFunction((*[0]byte)(C.goCall), C.DUK_VARARGS)
	d.putGoObjectRef(goFuncProp, fd)
}

type MethodSuite map[string]GoFunc

func (d *Context) EvalWith(source string, suite MethodSuite) error {
	if err := d.PevalString(source); err != 0 {
		return errors.New(d.SafeToString(-1))
	}

	d.PushObject()

	for prop, f := range suite {
		d.PushGoFunc(f)
		d.PutPropString(-2, prop)
	}

	if err := d.Pcall(1); err != 0 {
		return errors.New(d.SafeToString(-1))
	}

	return nil
}

// TBD: panic handling.
// When a goroutine panics, mark the context as panicking and throw a special JS
// exception, saving the value passed to panic() to the context.
// When the exception is caught, retrieve the value that was passed to panic(),
// and call panic() again.
<|endoftext|>


<commit_before><commit_msg>Add unit tests<commit_after><|endoftext|>


<commit_before>package monexec

import (
	"time"
	"os/exec"
	"syscall"
	"os"
	"io"
	"log"
	"bufio"
	"sync"
	"context"
	"errors"
)

const RestartInfinity = -1
const DefaultRestartTimeout = 5 * time.Second
const DefaultStopTimeout = 5 * time.Second

type BaseInfo struct {
	ID          string
	Command     string
	Args        []string
	Environment map[string]string
	Workdir     string
	StopTimeout time.Duration
}

func (b *BaseInfo) GetID() string {
	id := b.ID
	if id == "" {
		id = b.Command
	}
	return id
}

func (b *BaseInfo) WithID(id string) *BaseInfo {
	b.ID = id
	return b
}

func (b *BaseInfo) Arg(arg string) *BaseInfo {
	b.Args = append(b.Args, arg)
	return b
}

func (b *BaseInfo) Env(arg, value string) *BaseInfo {
	if b.Environment == nil {
		b.Environment = make(map[string]string)
	}
	b.Environment[arg] = value
	return b
}

type Delayed struct {
	BaseInfo
	RestartDelay time.Duration
}

type Executable struct {
	Delayed
	RestartCount int
	Critical     bool
}

func (exe *Executable) stopOrKill(logger *log.Logger, cmd *exec.Cmd) {
	ch := make(chan struct{}, 1)
	logger.Println("Sending SIGKILL")
	cmd.Process.Kill()
	go func() {
		cmd.Wait()
		ch <- struct{}{}
	}()

	select {
	case <-ch:
		logger.Println("Process gracefull stopped")
	case <-time.After(exe.StopTimeout):
		logger.Println("Process gracefull shutdown waiting timeout")
		cmd.Process.Signal(syscall.SIGTERM)
	}
}

func (exe *Executable) Run(logger *log.Logger, stop <-chan struct{}) error {
	cmd := exec.Command(exe.Command, exe.Args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGTERM,
	}
	for _, param := range os.Environ() {
		cmd.Env = append(cmd.Env, param)
	}
	if exe.Environment != nil {
		for k, v := range exe.Environment {
			cmd.Env = append(cmd.Env, k+"="+v)
		}
	}
	if exe.Workdir != "" {
		cmd.Dir = exe.Workdir
	}

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

	err = cmd.Start()
	if err != nil {
		return err
	}
	defer cmd.Process.Kill()
	logger.Println("Started with PID", cmd.Process.Pid)
	res := make(chan error, 1)
	go dumpToLogger(io.MultiReader(stdout, stderr), logger)
	go func() { res <- cmd.Wait() }()
	select {
	case <-stop:
		exe.stopOrKill(logger, cmd)
		return nil
	case err := <-res:
		return err
	}
}

func (exe *Executable) Daemonize(stop <-chan struct{}) error {
	var err error
	logger := log.New(os.Stdout, "["+exe.GetID()+"] ", log.LstdFlags)
LOOP:
	for {
		err = exe.Run(logger, stop)
		if err == nil {
			logger.Println("Process finished successfully")
		} else {
			logger.Println("Process finished unsuccessfully:", err)
		}
		if exe.Critical && err != nil {
			logger.Println("Process is critical - finishing work on error immediatly")
			break LOOP
		}

		if exe.RestartCount > 0 {
			logger.Println("Restarts left:", exe.RestartCount)
			exe.RestartCount--
		} else if exe.RestartCount == 0 {
			logger.Println("Process restart limit exceeded")
			break
		}

		logger.Println("Waiting", exe.RestartDelay, "before restart")
		select {
		case <-stop:
			logger.Println("Gracefull shutdown restart loop")
			break LOOP
		case <-time.After(exe.RestartDelay):
		}
	}
	if err != nil {
		logger.Println("Restart loop finished with error:", err)
	} else {
		logger.Println("Restart loop finished without error")
	}
	return err
}

func dumpToLogger(reader io.Reader, logger *log.Logger) {
	scanner := bufio.NewReader(reader)
	for {
		line, _, err := scanner.ReadLine()
		if err != nil {
			break
		}
		logger.Println("out:", string(line))
	}
}

type Monitor struct {
	Executables []*Executable
}

func (m *Monitor) Add(ex *Executable) *Executable {
	m.Executables = append(m.Executables, ex)
	return ex
}

func (m *Monitor) Oneshot(command string, args ...string) *Executable {
	exe := &Executable{RestartCount: 0, Critical: false}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	return m.Add(exe)
}

func (m *Monitor) Critical(command string, args ...string) *Executable {
	exe := &Executable{RestartCount: 0, Critical: true}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	return m.Add(exe)
}

func (m *Monitor) Restart(maxRetries int, command string, args ...string) *Executable {
	exe := &Executable{RestartCount: maxRetries, Critical: false}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	exe.RestartDelay = DefaultRestartTimeout
	return m.Add(exe)
}

func (m *Monitor) Forever(command string, args ...string) *Executable {
	return m.Restart(RestartInfinity, command, args...)
}

func (m *Monitor) Run(ctx context.Context) error {
	done := make(chan struct{})

	wg := sync.WaitGroup{}
	errs := make([]error, len(m.Executables))
	for i, exe := range m.Executables {
		wg.Add(1)
		go func(i int, exe *Executable) {
			defer wg.Done()
			errs[i] = exe.Daemonize(firstDone(ctx.Done(), done))
			if errs[i] != nil && exe.Critical {
				close(done)
			}
		}(i, exe)
	}
	wg.Wait()
	return m.joinErrors(errs)
}

func (m *Monitor) joinErrors(errs []error) error {
	errT := ""
	for i, err := range errs {
		if errT != "" {
			errT += "\n"
		}
		if err != nil {
			errT += m.Executables[i].GetID() + ": " + err.Error()
		}
	}
	if errT != "" {
		return errors.New(errT)
	}
	return nil
}

func firstDone(chanA, chanB <-chan struct{}) <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		select {
		case <-chanA:
			close(ch)
		case <-chanB:
			close(ch)
		}
	}()
	return ch
}
<commit_msg>prepare for yaml support<commit_after>package monexec

import (
	"time"
	"os/exec"
	"syscall"
	"os"
	"io"
	"log"
	"bufio"
	"sync"
	"context"
	"errors"
)

const RestartInfinity = -1
const DefaultRestartTimeout = 5 * time.Second
const DefaultStopTimeout = 5 * time.Second

type BaseInfo struct {
	ID          string            `yaml:"id"`
	Command     string            `yaml:"command"`
	Args        []string          `yaml:"args"`
	Environment map[string]string `yaml:"environment"`
	Workdir     string            `yaml:"workdir"`
	StopTimeout time.Duration     `yaml:"stop_timeout"`
}

func (b *BaseInfo) GetID() string {
	id := b.ID
	if id == "" {
		id = b.Command
	}
	return id
}

func (b *BaseInfo) WithID(id string) *BaseInfo {
	b.ID = id
	return b
}

func (b *BaseInfo) Arg(arg string) *BaseInfo {
	b.Args = append(b.Args, arg)
	return b
}

func (b *BaseInfo) Env(arg, value string) *BaseInfo {
	if b.Environment == nil {
		b.Environment = make(map[string]string)
	}
	b.Environment[arg] = value
	return b
}

type Delayed struct {
	BaseInfo                   `yaml:",squash"`
	RestartDelay time.Duration `yaml:"restart_delay"`
}

type Executable struct {
	Delayed       `yaml:",squash"`
	Retries  int  `yaml:"retries"`
	Critical bool `yaml:"critical"`
}

func (exe *Executable) stopOrKill(logger *log.Logger, cmd *exec.Cmd) {
	ch := make(chan struct{}, 1)
	logger.Println("Sending SIGKILL")
	cmd.Process.Kill()
	go func() {
		cmd.Wait()
		ch <- struct{}{}
	}()

	select {
	case <-ch:
		logger.Println("Process gracefull stopped")
	case <-time.After(exe.StopTimeout):
		logger.Println("Process gracefull shutdown waiting timeout")
		cmd.Process.Signal(syscall.SIGTERM)
	}
}

func (exe *Executable) Run(logger *log.Logger, stop <-chan struct{}) error {
	cmd := exec.Command(exe.Command, exe.Args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGTERM,
	}
	for _, param := range os.Environ() {
		cmd.Env = append(cmd.Env, param)
	}
	if exe.Environment != nil {
		for k, v := range exe.Environment {
			cmd.Env = append(cmd.Env, k+"="+v)
		}
	}
	if exe.Workdir != "" {
		cmd.Dir = exe.Workdir
	}

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

	err = cmd.Start()
	if err != nil {
		return err
	}
	defer cmd.Process.Kill()
	logger.Println("Started with PID", cmd.Process.Pid)
	res := make(chan error, 1)
	go dumpToLogger(io.MultiReader(stdout, stderr), logger)
	go func() { res <- cmd.Wait() }()
	select {
	case <-stop:
		exe.stopOrKill(logger, cmd)
		return nil
	case err := <-res:
		return err
	}
}

func (exe *Executable) Daemonize(stop <-chan struct{}) error {
	var err error
	logger := log.New(os.Stdout, "["+exe.GetID()+"] ", log.LstdFlags)
LOOP:
	for {
		err = exe.Run(logger, stop)
		if err == nil {
			logger.Println("Process finished successfully")
		} else {
			logger.Println("Process finished unsuccessfully:", err)
		}
		if exe.Critical && err != nil {
			logger.Println("Process is critical - finishing work on error immediatly")
			break LOOP
		}

		if exe.Retries > 0 {
			logger.Println("Restarts left:", exe.Retries)
			exe.Retries--
		} else if exe.Retries == 0 {
			logger.Println("Process restart limit exceeded")
			break
		}

		logger.Println("Waiting", exe.RestartDelay, "before restart")
		select {
		case <-stop:
			logger.Println("Gracefull shutdown restart loop")
			break LOOP
		case <-time.After(exe.RestartDelay):
		}
	}
	if err != nil {
		logger.Println("Restart loop finished with error:", err)
	} else {
		logger.Println("Restart loop finished without error")
	}
	return err
}

func dumpToLogger(reader io.Reader, logger *log.Logger) {
	scanner := bufio.NewReader(reader)
	for {
		line, _, err := scanner.ReadLine()
		if err != nil {
			break
		}
		logger.Println("out:", string(line))
	}
}

type Monitor struct {
	Executables []*Executable
}

func (m *Monitor) Add(ex *Executable) *Executable {
	m.Executables = append(m.Executables, ex)
	return ex
}

func (m *Monitor) Oneshot(command string, args ...string) *Executable {
	exe := &Executable{Retries: 0, Critical: false}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	return m.Add(exe)
}

func (m *Monitor) Critical(command string, args ...string) *Executable {
	exe := &Executable{Retries: 0, Critical: true}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	return m.Add(exe)
}

func (m *Monitor) Restart(maxRetries int, command string, args ...string) *Executable {
	exe := &Executable{Retries: maxRetries, Critical: false}
	exe.Command = command
	exe.Args = append(exe.Args, args...)
	exe.StopTimeout = DefaultStopTimeout
	exe.RestartDelay = DefaultRestartTimeout
	return m.Add(exe)
}

func (m *Monitor) Forever(command string, args ...string) *Executable {
	return m.Restart(RestartInfinity, command, args...)
}

func (m *Monitor) Run(ctx context.Context) error {
	done := make(chan struct{})

	wg := sync.WaitGroup{}
	errs := make([]error, len(m.Executables))
	for i, exe := range m.Executables {
		wg.Add(1)
		go func(i int, exe *Executable) {
			defer wg.Done()
			errs[i] = exe.Daemonize(firstDone(ctx.Done(), done))
			if errs[i] != nil && exe.Critical {
				close(done)
			}
		}(i, exe)
	}
	wg.Wait()
	return m.joinErrors(errs)
}

func (m *Monitor) joinErrors(errs []error) error {
	errT := ""
	for i, err := range errs {
		if errT != "" {
			errT += "\n"
		}
		if err != nil {
			errT += m.Executables[i].GetID() + ": " + err.Error()
		}
	}
	if errT != "" {
		return errors.New(errT)
	}
	return nil
}

func firstDone(chanA, chanB <-chan struct{}) <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		select {
		case <-chanA:
			close(ch)
		case <-chanB:
			close(ch)
		}
	}()
	return ch
}
<|endoftext|>