<commit_before>package plugin

import (
	"fmt"
	"net"
	osexec "os/exec"
	"strings"
	"time"

	log "github.com/golang/glog"

	"github.com/openshift/origin/pkg/sdn/plugin/cniserver"

	osclient "github.com/openshift/origin/pkg/client"
	osapi "github.com/openshift/origin/pkg/sdn/api"
	"github.com/openshift/origin/pkg/util/ipcmd"
	"github.com/openshift/origin/pkg/util/netutils"
	"github.com/openshift/origin/pkg/util/ovs"

	docker "github.com/fsouza/go-dockerclient"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	knetwork "k8s.io/kubernetes/pkg/kubelet/network"
	"k8s.io/kubernetes/pkg/labels"
	kexec "k8s.io/kubernetes/pkg/util/exec"
	kubeutilnet "k8s.io/kubernetes/pkg/util/net"
	kwait "k8s.io/kubernetes/pkg/util/wait"
)

type OsdnNode struct {
	multitenant        bool
	kClient            *kclient.Client
	osClient           *osclient.Client
	ovs                *ovs.Interface
	networkInfo        *NetworkInfo
	podManager         *podManager
	localSubnetCIDR    string
	localIP            string
	hostName           string
	podNetworkReady    chan struct{}
	kubeletInitReady   chan struct{}
	vnids              *nodeVNIDMap
	iptablesSyncPeriod time.Duration
	mtu                uint32
	egressPolicies     map[uint32][]*osapi.EgressNetworkPolicy

	host             knetwork.Host
	kubeletCniPlugin knetwork.NetworkPlugin

	clearLbr0IptablesRule bool
}

// Called by higher layers to create the plugin SDN node instance
func NewNodePlugin(pluginName string, osClient *osclient.Client, kClient *kclient.Client, hostname string, selfIP string, iptablesSyncPeriod time.Duration, mtu uint32) (*OsdnNode, error) {
	if !osapi.IsOpenShiftNetworkPlugin(pluginName) {
		return nil, nil
	}

	log.Infof("Initializing SDN node of type %q with configured hostname %q (IP %q), iptables sync period %q", pluginName, hostname, selfIP, iptablesSyncPeriod.String())
	if hostname == "" {
		output, err := kexec.New().Command("uname", "-n").CombinedOutput()
		if err != nil {
			return nil, err
		}
		hostname = strings.TrimSpace(string(output))
		log.Infof("Resolved hostname to %q", hostname)
	}
	if selfIP == "" {
		var err error
		selfIP, err = netutils.GetNodeIP(hostname)
		if err != nil {
			log.V(5).Infof("Failed to determine node address from hostname %s; using default interface (%v)", hostname, err)
			var defaultIP net.IP
			defaultIP, err = kubeutilnet.ChooseHostInterface()
			if err != nil {
				return nil, err
			}
			selfIP = defaultIP.String()
			log.Infof("Resolved IP address to %q", selfIP)
		}
	}

	ovsif, err := ovs.New(kexec.New(), BR)
	if err != nil {
		return nil, err
	}

	plugin := &OsdnNode{
		multitenant:        osapi.IsOpenShiftMultitenantNetworkPlugin(pluginName),
		kClient:            kClient,
		osClient:           osClient,
		ovs:                ovsif,
		localIP:            selfIP,
		hostName:           hostname,
		vnids:              newNodeVNIDMap(),
		podNetworkReady:    make(chan struct{}),
		kubeletInitReady:   make(chan struct{}),
		iptablesSyncPeriod: iptablesSyncPeriod,
		mtu:                mtu,
		egressPolicies:     make(map[uint32][]*osapi.EgressNetworkPolicy),
	}

	if err := plugin.dockerPreCNICleanup(); err != nil {
		return nil, err
	}

	return plugin, nil
}

// Detect whether we are upgrading from a pre-CNI openshift and clean up
// interfaces and iptables rules that are no longer required
func (node *OsdnNode) dockerPreCNICleanup() error {
	exec := kexec.New()
	itx := ipcmd.NewTransaction(exec, "lbr0")
	itx.SetLink("down")
	if err := itx.EndTransaction(); err != nil {
		// no cleanup required
		return nil
	}

	node.clearLbr0IptablesRule = true

	// Restart docker to kill old pods and make it use docker0 again.
	// "systemctl restart" will bail out (unnecessarily) in the
	// OpenShift-in-a-container case, so we work around that by sending
	// the messages by hand.
	if err := osexec.Command("dbus-send", "--system", "--print-reply", "--reply-timeout=2000", "--type=method_call", "--dest=org.freedesktop.systemd1", "/org/freedesktop/systemd1", "org.freedesktop.systemd1.Manager.Reload"); err != nil {
		log.Error(err)
	}
	if err := osexec.Command("dbus-send", "--system", "--print-reply", "--reply-timeout=2000", "--type=method_call", "--dest=org.freedesktop.systemd1", "/org/freedesktop/systemd1", "org.freedesktop.systemd1.Manager.RestartUnit", "string:'docker.service' string:'replace'"); err != nil {
		log.Error(err)
	}

	// Delete pre-CNI interfaces
	for _, intf := range []string{"lbr0", "vovsbr", "vlinuxbr"} {
		itx := ipcmd.NewTransaction(exec, intf)
		itx.DeleteLink()
		itx.IgnoreError()
		itx.EndTransaction()
	}

	// Wait until docker has restarted since kubelet will exit it docker isn't running
	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		return fmt.Errorf("failed to get docker client: %v", err)
	}
	err = kwait.ExponentialBackoff(
		kwait.Backoff{
			Duration: 100 * time.Millisecond,
			Factor:   1.2,
			Steps:    6,
		},
		func() (bool, error) {
			if err := dockerClient.Ping(); err != nil {
				// wait longer
				return false, nil
			}
			return true, nil
		})
	if err != nil {
		return fmt.Errorf("failed to connect to docker after SDN cleanup restart: %v", err)
	}

	log.Infof("Cleaned up left-over openshift-sdn docker bridge and interfaces")

	return nil
}

func (node *OsdnNode) Start() error {
	var err error
	node.networkInfo, err = getNetworkInfo(node.osClient)
	if err != nil {
		return fmt.Errorf("Failed to get network information: %v", err)
	}

	nodeIPTables := newNodeIPTables(node.networkInfo.ClusterNetwork.String(), node.iptablesSyncPeriod)
	if err = nodeIPTables.Setup(); err != nil {
		return fmt.Errorf("Failed to set up iptables: %v", err)
	}

	node.localSubnetCIDR, err = node.getLocalSubnet()
	if err != nil {
		return err
	}

	networkChanged, err := node.SetupSDN()
	if err != nil {
		return err
	}

	err = node.SubnetStartNode()
	if err != nil {
		return err
	}

	if node.multitenant {
		if err = node.VnidStartNode(); err != nil {
			return err
		}
		if err = node.SetupEgressNetworkPolicy(); err != nil {
			return err
		}
	}

	// Wait for kubelet to init the plugin so we get a knetwork.Host
	log.V(5).Infof("Waiting for kubelet network plugin initialization")
	<-node.kubeletInitReady

	log.V(5).Infof("Creating and initializing openshift-sdn pod manager")
	node.podManager, err = newPodManager(node.host, node.multitenant, node.localSubnetCIDR, node.networkInfo, node.kClient, node.vnids, node.mtu)
	if err != nil {
		return err
	}
	if err := node.podManager.Start(cniserver.CNIServerSocketPath); err != nil {
		return err
	}

	if networkChanged {
		var pods []kapi.Pod
		pods, err = node.GetLocalPods(kapi.NamespaceAll)
		if err != nil {
			return err
		}
		for _, p := range pods {
			err = node.UpdatePod(p)
			if err != nil {
				log.Warningf("Could not update pod %q: %s", p.Name, err)
			}
		}
	}

	log.V(5).Infof("openshift-sdn network plugin ready")
	node.markPodNetworkReady()

	return nil
}

// FIXME: this should eventually go into kubelet via a CNI UPDATE/CHANGE action
// See https://github.com/containernetworking/cni/issues/89
func (node *OsdnNode) UpdatePod(pod kapi.Pod) error {
	req := &cniserver.PodRequest{
		Command:      cniserver.CNI_UPDATE,
		PodNamespace: pod.Namespace,
		PodName:      pod.Name,
		ContainerId:  getPodContainerID(&pod),
		// netns is read from docker if needed, since we don't get it from kubelet
		Result: make(chan *cniserver.PodResult),
	}

	// Send request and wait for the result
	_, err := node.podManager.handleCNIRequest(req)
	return err
}

func (node *OsdnNode) GetLocalPods(namespace string) ([]kapi.Pod, error) {
	fieldSelector := fields.Set{"spec.nodeName": node.hostName}.AsSelector()
	opts := kapi.ListOptions{
		LabelSelector: labels.Everything(),
		FieldSelector: fieldSelector,
	}
	podList, err := node.kClient.Pods(namespace).List(opts)
	if err != nil {
		return nil, err
	}

	// Filter running pods
	pods := make([]kapi.Pod, 0, len(podList.Items))
	for _, pod := range podList.Items {
		if pod.Status.Phase == kapi.PodRunning {
			pods = append(pods, pod)
		}
	}
	return pods, nil
}

func (node *OsdnNode) markPodNetworkReady() {
	close(node.podNetworkReady)
}

func (node *OsdnNode) IsPodNetworkReady() error {
	select {
	case <-node.podNetworkReady:
		return nil
	default:
		return fmt.Errorf("SDN pod network is not ready")
	}
}
<commit_msg>Fix the command to restart docker on 1.3->1.4 upgrade<commit_after>package plugin

import (
	"fmt"
	"net"
	osexec "os/exec"
	"strings"
	"time"

	log "github.com/golang/glog"

	"github.com/openshift/origin/pkg/sdn/plugin/cniserver"

	osclient "github.com/openshift/origin/pkg/client"
	osapi "github.com/openshift/origin/pkg/sdn/api"
	"github.com/openshift/origin/pkg/util/ipcmd"
	"github.com/openshift/origin/pkg/util/netutils"
	"github.com/openshift/origin/pkg/util/ovs"

	docker "github.com/fsouza/go-dockerclient"

	kapi "k8s.io/kubernetes/pkg/api"
	kclient "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	knetwork "k8s.io/kubernetes/pkg/kubelet/network"
	"k8s.io/kubernetes/pkg/labels"
	kexec "k8s.io/kubernetes/pkg/util/exec"
	kubeutilnet "k8s.io/kubernetes/pkg/util/net"
	kwait "k8s.io/kubernetes/pkg/util/wait"
)

type OsdnNode struct {
	multitenant        bool
	kClient            *kclient.Client
	osClient           *osclient.Client
	ovs                *ovs.Interface
	networkInfo        *NetworkInfo
	podManager         *podManager
	localSubnetCIDR    string
	localIP            string
	hostName           string
	podNetworkReady    chan struct{}
	kubeletInitReady   chan struct{}
	vnids              *nodeVNIDMap
	iptablesSyncPeriod time.Duration
	mtu                uint32
	egressPolicies     map[uint32][]*osapi.EgressNetworkPolicy

	host             knetwork.Host
	kubeletCniPlugin knetwork.NetworkPlugin

	clearLbr0IptablesRule bool
}

// Called by higher layers to create the plugin SDN node instance
func NewNodePlugin(pluginName string, osClient *osclient.Client, kClient *kclient.Client, hostname string, selfIP string, iptablesSyncPeriod time.Duration, mtu uint32) (*OsdnNode, error) {
	if !osapi.IsOpenShiftNetworkPlugin(pluginName) {
		return nil, nil
	}

	log.Infof("Initializing SDN node of type %q with configured hostname %q (IP %q), iptables sync period %q", pluginName, hostname, selfIP, iptablesSyncPeriod.String())
	if hostname == "" {
		output, err := kexec.New().Command("uname", "-n").CombinedOutput()
		if err != nil {
			return nil, err
		}
		hostname = strings.TrimSpace(string(output))
		log.Infof("Resolved hostname to %q", hostname)
	}
	if selfIP == "" {
		var err error
		selfIP, err = netutils.GetNodeIP(hostname)
		if err != nil {
			log.V(5).Infof("Failed to determine node address from hostname %s; using default interface (%v)", hostname, err)
			var defaultIP net.IP
			defaultIP, err = kubeutilnet.ChooseHostInterface()
			if err != nil {
				return nil, err
			}
			selfIP = defaultIP.String()
			log.Infof("Resolved IP address to %q", selfIP)
		}
	}

	ovsif, err := ovs.New(kexec.New(), BR)
	if err != nil {
		return nil, err
	}

	plugin := &OsdnNode{
		multitenant:        osapi.IsOpenShiftMultitenantNetworkPlugin(pluginName),
		kClient:            kClient,
		osClient:           osClient,
		ovs:                ovsif,
		localIP:            selfIP,
		hostName:           hostname,
		vnids:              newNodeVNIDMap(),
		podNetworkReady:    make(chan struct{}),
		kubeletInitReady:   make(chan struct{}),
		iptablesSyncPeriod: iptablesSyncPeriod,
		mtu:                mtu,
		egressPolicies:     make(map[uint32][]*osapi.EgressNetworkPolicy),
	}

	if err := plugin.dockerPreCNICleanup(); err != nil {
		return nil, err
	}

	return plugin, nil
}

// Detect whether we are upgrading from a pre-CNI openshift and clean up
// interfaces and iptables rules that are no longer required
func (node *OsdnNode) dockerPreCNICleanup() error {
	exec := kexec.New()
	itx := ipcmd.NewTransaction(exec, "lbr0")
	itx.SetLink("down")
	if err := itx.EndTransaction(); err != nil {
		// no cleanup required
		return nil
	}

	node.clearLbr0IptablesRule = true

	// Restart docker to kill old pods and make it use docker0 again.
	// "systemctl restart" will bail out (unnecessarily) in the
	// OpenShift-in-a-container case, so we work around that by sending
	// the messages by hand.
	if _, err := osexec.Command("dbus-send", "--system", "--print-reply", "--reply-timeout=2000", "--type=method_call", "--dest=org.freedesktop.systemd1", "/org/freedesktop/systemd1", "org.freedesktop.systemd1.Manager.Reload").CombinedOutput(); err != nil {
		log.Error(err)
	}
	if _, err := osexec.Command("dbus-send", "--system", "--print-reply", "--reply-timeout=2000", "--type=method_call", "--dest=org.freedesktop.systemd1", "/org/freedesktop/systemd1", "org.freedesktop.systemd1.Manager.RestartUnit", "string:'docker.service' string:'replace'").CombinedOutput(); err != nil {
		log.Error(err)
	}

	// Delete pre-CNI interfaces
	for _, intf := range []string{"lbr0", "vovsbr", "vlinuxbr"} {
		itx := ipcmd.NewTransaction(exec, intf)
		itx.DeleteLink()
		itx.IgnoreError()
		itx.EndTransaction()
	}

	// Wait until docker has restarted since kubelet will exit it docker isn't running
	dockerClient, err := docker.NewClientFromEnv()
	if err != nil {
		return fmt.Errorf("failed to get docker client: %v", err)
	}
	err = kwait.ExponentialBackoff(
		kwait.Backoff{
			Duration: 100 * time.Millisecond,
			Factor:   1.2,
			Steps:    6,
		},
		func() (bool, error) {
			if err := dockerClient.Ping(); err != nil {
				// wait longer
				return false, nil
			}
			return true, nil
		})
	if err != nil {
		return fmt.Errorf("failed to connect to docker after SDN cleanup restart: %v", err)
	}

	log.Infof("Cleaned up left-over openshift-sdn docker bridge and interfaces")

	return nil
}

func (node *OsdnNode) Start() error {
	var err error
	node.networkInfo, err = getNetworkInfo(node.osClient)
	if err != nil {
		return fmt.Errorf("Failed to get network information: %v", err)
	}

	nodeIPTables := newNodeIPTables(node.networkInfo.ClusterNetwork.String(), node.iptablesSyncPeriod)
	if err = nodeIPTables.Setup(); err != nil {
		return fmt.Errorf("Failed to set up iptables: %v", err)
	}

	node.localSubnetCIDR, err = node.getLocalSubnet()
	if err != nil {
		return err
	}

	networkChanged, err := node.SetupSDN()
	if err != nil {
		return err
	}

	err = node.SubnetStartNode()
	if err != nil {
		return err
	}

	if node.multitenant {
		if err = node.VnidStartNode(); err != nil {
			return err
		}
		if err = node.SetupEgressNetworkPolicy(); err != nil {
			return err
		}
	}

	// Wait for kubelet to init the plugin so we get a knetwork.Host
	log.V(5).Infof("Waiting for kubelet network plugin initialization")
	<-node.kubeletInitReady

	log.V(5).Infof("Creating and initializing openshift-sdn pod manager")
	node.podManager, err = newPodManager(node.host, node.multitenant, node.localSubnetCIDR, node.networkInfo, node.kClient, node.vnids, node.mtu)
	if err != nil {
		return err
	}
	if err := node.podManager.Start(cniserver.CNIServerSocketPath); err != nil {
		return err
	}

	if networkChanged {
		var pods []kapi.Pod
		pods, err = node.GetLocalPods(kapi.NamespaceAll)
		if err != nil {
			return err
		}
		for _, p := range pods {
			err = node.UpdatePod(p)
			if err != nil {
				log.Warningf("Could not update pod %q: %s", p.Name, err)
			}
		}
	}

	log.V(5).Infof("openshift-sdn network plugin ready")
	node.markPodNetworkReady()

	return nil
}

// FIXME: this should eventually go into kubelet via a CNI UPDATE/CHANGE action
// See https://github.com/containernetworking/cni/issues/89
func (node *OsdnNode) UpdatePod(pod kapi.Pod) error {
	req := &cniserver.PodRequest{
		Command:      cniserver.CNI_UPDATE,
		PodNamespace: pod.Namespace,
		PodName:      pod.Name,
		ContainerId:  getPodContainerID(&pod),
		// netns is read from docker if needed, since we don't get it from kubelet
		Result: make(chan *cniserver.PodResult),
	}

	// Send request and wait for the result
	_, err := node.podManager.handleCNIRequest(req)
	return err
}

func (node *OsdnNode) GetLocalPods(namespace string) ([]kapi.Pod, error) {
	fieldSelector := fields.Set{"spec.nodeName": node.hostName}.AsSelector()
	opts := kapi.ListOptions{
		LabelSelector: labels.Everything(),
		FieldSelector: fieldSelector,
	}
	podList, err := node.kClient.Pods(namespace).List(opts)
	if err != nil {
		return nil, err
	}

	// Filter running pods
	pods := make([]kapi.Pod, 0, len(podList.Items))
	for _, pod := range podList.Items {
		if pod.Status.Phase == kapi.PodRunning {
			pods = append(pods, pod)
		}
	}
	return pods, nil
}

func (node *OsdnNode) markPodNetworkReady() {
	close(node.podNetworkReady)
}

func (node *OsdnNode) IsPodNetworkReady() error {
	select {
	case <-node.podNetworkReady:
		return nil
	default:
		return fmt.Errorf("SDN pod network is not ready")
	}
}
<|endoftext|>


<commit_before>package secrets

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/containers/storage/pkg/idtools"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/opencontainers/selinux/go-selinux/label"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var (
	// DefaultMountsFile holds the default mount paths in the form
	// "host_path:container_path"
	DefaultMountsFile = "/usr/share/containers/mounts.conf"
	// OverrideMountsFile holds the default mount paths in the form
	// "host_path:container_path" overridden by the user
	OverrideMountsFile = "/etc/containers/mounts.conf"
	// UserOverrideMountsFile holds the default mount paths in the form
	// "host_path:container_path" overridden by the rootless user
	UserOverrideMountsFile = filepath.Join(os.Getenv("HOME"), ".config/containers/mounts.conf")
)

// secretData stores the name of the file and the content read from it
type secretData struct {
	name string
	data []byte
}

// saveTo saves secret data to given directory
func (s secretData) saveTo(dir string) error {
	path := filepath.Join(dir, s.name)
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil && !os.IsExist(err) {
		return err
	}
	return ioutil.WriteFile(path, s.data, 0700)
}

func readAll(root, prefix string) ([]secretData, error) {
	path := filepath.Join(root, prefix)

	data := []secretData{}

	files, err := ioutil.ReadDir(path)
	if err != nil {
		if os.IsNotExist(err) {
			return data, nil
		}

		return nil, err
	}

	for _, f := range files {
		fileData, err := readFile(root, filepath.Join(prefix, f.Name()))
		if err != nil {
			// If the file did not exist, might be a dangling symlink
			// Ignore the error
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		data = append(data, fileData...)
	}

	return data, nil
}

func readFile(root, name string) ([]secretData, error) {
	path := filepath.Join(root, name)

	s, err := os.Stat(path)
	if err != nil {
		return nil, err
	}

	if s.IsDir() {
		dirData, err := readAll(root, name)
		if err != nil {
			return nil, err
		}
		return dirData, nil
	}
	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return []secretData{{name: name, data: bytes}}, nil
}

func getHostSecretData(hostDir string) ([]secretData, error) {
	var allSecrets []secretData
	hostSecrets, err := readAll(hostDir, "")
	if err != nil {
		return nil, errors.Wrapf(err, "failed to read secrets from %q", hostDir)
	}
	return append(allSecrets, hostSecrets...), nil
}

func getMounts(filePath string) []string {
	file, err := os.Open(filePath)
	if err != nil {
		// This is expected on most systems
		logrus.Debugf("file %q not found, skipping...", filePath)
		return nil
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		logrus.Errorf("error reading file %q, %v skipping...", filePath, err)
		return nil
	}
	var mounts []string
	for scanner.Scan() {
		if strings.HasPrefix(strings.TrimSpace(scanner.Text()), "/") {
			mounts = append(mounts, scanner.Text())
		} else {
			logrus.Debugf("skipping unrecognized mount in %v: %q",
				filePath, scanner.Text())
		}
	}
	return mounts
}

// getHostAndCtrDir separates the host:container paths
func getMountsMap(path string) (string, string, error) {
	arr := strings.SplitN(path, ":", 2)
	if len(arr) == 2 {
		return arr[0], arr[1], nil
	}
	return "", "", errors.Errorf("unable to get host and container dir")
}

// SecretMounts copies, adds, and mounts the secrets to the container root filesystem
func SecretMounts(mountLabel, containerWorkingDir, mountFile string, rootless bool) []rspec.Mount {
	return SecretMountsWithUIDGID(mountLabel, containerWorkingDir, mountFile, containerWorkingDir, 0, 0, rootless)
}

// SecretMountsWithUIDGID specifies the uid/gid of the owner
func SecretMountsWithUIDGID(mountLabel, containerWorkingDir, mountFile, mountPrefix string, uid, gid int, rootless bool) []rspec.Mount {
	var (
		secretMounts []rspec.Mount
		mountFiles   []string
	)
	// Add secrets from paths given in the mounts.conf files
	// mountFile will have a value if the hidden --default-mounts-file flag is set
	// Note for testing purposes only
	if mountFile == "" {
		mountFiles = append(mountFiles, []string{OverrideMountsFile, DefaultMountsFile}...)
		if rootless {
			mountFiles = append([]string{UserOverrideMountsFile}, mountFiles...)
		}
	} else {
		mountFiles = append(mountFiles, mountFile)
	}
	for _, file := range mountFiles {
		if _, err := os.Stat(file); err == nil {
			mounts, err := addSecretsFromMountsFile(file, mountLabel, containerWorkingDir, mountPrefix, uid, gid)
			if err != nil {
				logrus.Warnf("error mounting secrets, skipping: %v", err)
			}
			secretMounts = mounts
			break
		}
	}

	// Add FIPS mode secret if /etc/system-fips exists on the host
	_, err := os.Stat("/etc/system-fips")
	if err == nil {
		if err := addFIPSModeSecret(&secretMounts, containerWorkingDir, mountPrefix, mountLabel, uid, gid); err != nil {
			logrus.Errorf("error adding FIPS mode secret to container: %v", err)
		}
	} else if os.IsNotExist(err) {
		logrus.Debug("/etc/system-fips does not exist on host, not mounting FIPS mode secret")
	} else {
		logrus.Errorf("stat /etc/system-fips failed for FIPS mode secret: %v", err)
	}
	return secretMounts
}

func rchown(chowndir string, uid, gid int) error {
	return filepath.Walk(chowndir, func(filePath string, f os.FileInfo, err error) error {
		return os.Lchown(filePath, uid, gid)
	})
}

// addSecretsFromMountsFile copies the contents of host directory to container directory
// and returns a list of mounts
func addSecretsFromMountsFile(filePath, mountLabel, containerWorkingDir, mountPrefix string, uid, gid int) ([]rspec.Mount, error) {
	var mounts []rspec.Mount
	defaultMountsPaths := getMounts(filePath)
	for _, path := range defaultMountsPaths {
		hostDir, ctrDir, err := getMountsMap(path)
		if err != nil {
			return nil, err
		}
		// skip if the hostDir path doesn't exist
		if _, err = os.Stat(hostDir); err != nil {
			if os.IsNotExist(err) {
				logrus.Warnf("Path %q from %q doesn't exist, skipping", hostDir, filePath)
				continue
			}
			return nil, errors.Wrapf(err, "failed to stat %q", hostDir)
		}

		ctrDirOnHost := filepath.Join(containerWorkingDir, ctrDir)

		// In the event of a restart, don't want to copy secrets over again as they already would exist in ctrDirOnHost
		_, err = os.Stat(ctrDirOnHost)
		if os.IsNotExist(err) {
			if err = os.MkdirAll(ctrDirOnHost, 0755); err != nil {
				return nil, errors.Wrapf(err, "making container directory %q failed", ctrDirOnHost)
			}
			hostDir, err = resolveSymbolicLink(hostDir)
			if err != nil {
				return nil, err
			}

			data, err := getHostSecretData(hostDir)
			if err != nil {
				return nil, errors.Wrapf(err, "getting host secret data failed")
			}
			for _, s := range data {
				if err := s.saveTo(ctrDirOnHost); err != nil {
					return nil, errors.Wrapf(err, "error saving data to container filesystem on host %q", ctrDirOnHost)
				}
			}

			err = label.Relabel(ctrDirOnHost, mountLabel, false)
			if err != nil {
				return nil, errors.Wrap(err, "error applying correct labels")
			}
			if uid != 0 || gid != 0 {
				if err := rchown(ctrDirOnHost, uid, gid); err != nil {
					return nil, err
				}
			}
		} else if err != nil {
			return nil, errors.Wrapf(err, "error getting status of %q", ctrDirOnHost)
		}

		m := rspec.Mount{
			Source:      filepath.Join(mountPrefix, ctrDir),
			Destination: ctrDir,
			Type:        "bind",
			Options:     []string{"bind", "rprivate"},
		}

		mounts = append(mounts, m)
	}
	return mounts, nil
}

// addFIPSModeSecret creates /run/secrets/system-fips in the container
// root filesystem if /etc/system-fips exists on hosts.
// This enables the container to be FIPS compliant and run openssl in
// FIPS mode as the host is also in FIPS mode.
func addFIPSModeSecret(mounts *[]rspec.Mount, containerWorkingDir, mountPrefix, mountLabel string, uid, gid int) error {
	secretsDir := "/run/secrets"
	ctrDirOnHost := filepath.Join(containerWorkingDir, secretsDir)
	if _, err := os.Stat(ctrDirOnHost); os.IsNotExist(err) {
		if err = idtools.MkdirAllAs(ctrDirOnHost, 0755, uid, gid); err != nil {
			return errors.Wrapf(err, "making container directory on host failed")
		}
		if err = label.Relabel(ctrDirOnHost, mountLabel, false); err != nil {
			return errors.Wrap(err, "error applying correct labels")
		}
	}
	fipsFile := filepath.Join(ctrDirOnHost, "system-fips")
	// In the event of restart, it is possible for the FIPS mode file to already exist
	if _, err := os.Stat(fipsFile); os.IsNotExist(err) {
		file, err := os.Create(fipsFile)
		if err != nil {
			return errors.Wrapf(err, "error creating system-fips file in container for FIPS mode")
		}
		defer file.Close()
	}

	if !mountExists(*mounts, secretsDir) {
		m := rspec.Mount{
			Source:      filepath.Join(mountPrefix, secretsDir),
			Destination: secretsDir,
			Type:        "bind",
			Options:     []string{"bind", "rprivate"},
		}
		*mounts = append(*mounts, m)
	}

	return nil
}

// mountExists checks if a mount already exists in the spec
func mountExists(mounts []rspec.Mount, dest string) bool {
	for _, mount := range mounts {
		if mount.Destination == dest {
			return true
		}
	}
	return false
}

// resolveSymbolicLink resolves a possbile symlink path. If the path is a symlink, returns resolved
// path; if not, returns the original path.
func resolveSymbolicLink(path string) (string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return "", err
	}
	if info.Mode()&os.ModeSymlink != os.ModeSymlink {
		return path, nil
	}
	return filepath.EvalSymlinks(path)
}
<commit_msg>Add support for file secret mounts<commit_after>package secrets

import (
	"bufio"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/containers/storage/pkg/idtools"
	rspec "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/opencontainers/selinux/go-selinux/label"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

var (
	// DefaultMountsFile holds the default mount paths in the form
	// "host_path:container_path"
	DefaultMountsFile = "/usr/share/containers/mounts.conf"
	// OverrideMountsFile holds the default mount paths in the form
	// "host_path:container_path" overridden by the user
	OverrideMountsFile = "/etc/containers/mounts.conf"
	// UserOverrideMountsFile holds the default mount paths in the form
	// "host_path:container_path" overridden by the rootless user
	UserOverrideMountsFile = filepath.Join(os.Getenv("HOME"), ".config/containers/mounts.conf")
)

// secretData stores the name of the file and the content read from it
type secretData struct {
	name string
	data []byte
}

// saveTo saves secret data to given directory
func (s secretData) saveTo(dir string) error {
	path := filepath.Join(dir, s.name)
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil && !os.IsExist(err) {
		return err
	}
	return ioutil.WriteFile(path, s.data, 0700)
}

func readAll(root, prefix string) ([]secretData, error) {
	path := filepath.Join(root, prefix)

	data := []secretData{}

	files, err := ioutil.ReadDir(path)
	if err != nil {
		if os.IsNotExist(err) {
			return data, nil
		}

		return nil, err
	}

	for _, f := range files {
		fileData, err := readFile(root, filepath.Join(prefix, f.Name()))
		if err != nil {
			// If the file did not exist, might be a dangling symlink
			// Ignore the error
			if os.IsNotExist(err) {
				continue
			}
			return nil, err
		}
		data = append(data, fileData...)
	}

	return data, nil
}

func readFile(root, name string) ([]secretData, error) {
	path := filepath.Join(root, name)

	s, err := os.Stat(path)
	if err != nil {
		return nil, err
	}

	if s.IsDir() {
		dirData, err := readAll(root, name)
		if err != nil {
			return nil, err
		}
		return dirData, nil
	}
	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return []secretData{{name: name, data: bytes}}, nil
}

func getHostSecretData(hostDir string) ([]secretData, error) {
	var allSecrets []secretData
	hostSecrets, err := readAll(hostDir, "")
	if err != nil {
		return nil, errors.Wrapf(err, "failed to read secrets from %q", hostDir)
	}
	return append(allSecrets, hostSecrets...), nil
}

func getMounts(filePath string) []string {
	file, err := os.Open(filePath)
	if err != nil {
		// This is expected on most systems
		logrus.Debugf("file %q not found, skipping...", filePath)
		return nil
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		logrus.Errorf("error reading file %q, %v skipping...", filePath, err)
		return nil
	}
	var mounts []string
	for scanner.Scan() {
		if strings.HasPrefix(strings.TrimSpace(scanner.Text()), "/") {
			mounts = append(mounts, scanner.Text())
		} else {
			logrus.Debugf("skipping unrecognized mount in %v: %q",
				filePath, scanner.Text())
		}
	}
	return mounts
}

// getHostAndCtrDir separates the host:container paths
func getMountsMap(path string) (string, string, error) {
	arr := strings.SplitN(path, ":", 2)
	if len(arr) == 2 {
		return arr[0], arr[1], nil
	}
	return "", "", errors.Errorf("unable to get host and container dir")
}

// SecretMounts copies, adds, and mounts the secrets to the container root filesystem
func SecretMounts(mountLabel, containerWorkingDir, mountFile string, rootless bool) []rspec.Mount {
	return SecretMountsWithUIDGID(mountLabel, containerWorkingDir, mountFile, containerWorkingDir, 0, 0, rootless)
}

// SecretMountsWithUIDGID specifies the uid/gid of the owner
func SecretMountsWithUIDGID(mountLabel, containerWorkingDir, mountFile, mountPrefix string, uid, gid int, rootless bool) []rspec.Mount {
	var (
		secretMounts []rspec.Mount
		mountFiles   []string
	)
	// Add secrets from paths given in the mounts.conf files
	// mountFile will have a value if the hidden --default-mounts-file flag is set
	// Note for testing purposes only
	if mountFile == "" {
		mountFiles = append(mountFiles, []string{OverrideMountsFile, DefaultMountsFile}...)
		if rootless {
			mountFiles = append([]string{UserOverrideMountsFile}, mountFiles...)
		}
	} else {
		mountFiles = append(mountFiles, mountFile)
	}
	for _, file := range mountFiles {
		if _, err := os.Stat(file); err == nil {
			mounts, err := addSecretsFromMountsFile(file, mountLabel, containerWorkingDir, mountPrefix, uid, gid)
			if err != nil {
				logrus.Warnf("error mounting secrets, skipping: %v", err)
			}
			secretMounts = mounts
			break
		}
	}

	// Add FIPS mode secret if /etc/system-fips exists on the host
	_, err := os.Stat("/etc/system-fips")
	if err == nil {
		if err := addFIPSModeSecret(&secretMounts, containerWorkingDir, mountPrefix, mountLabel, uid, gid); err != nil {
			logrus.Errorf("error adding FIPS mode secret to container: %v", err)
		}
	} else if os.IsNotExist(err) {
		logrus.Debug("/etc/system-fips does not exist on host, not mounting FIPS mode secret")
	} else {
		logrus.Errorf("stat /etc/system-fips failed for FIPS mode secret: %v", err)
	}
	return secretMounts
}

func rchown(chowndir string, uid, gid int) error {
	return filepath.Walk(chowndir, func(filePath string, f os.FileInfo, err error) error {
		return os.Lchown(filePath, uid, gid)
	})
}

// addSecretsFromMountsFile copies the contents of host directory to container directory
// and returns a list of mounts
func addSecretsFromMountsFile(filePath, mountLabel, containerWorkingDir, mountPrefix string, uid, gid int) ([]rspec.Mount, error) {
	var mounts []rspec.Mount
	defaultMountsPaths := getMounts(filePath)
	for _, path := range defaultMountsPaths {
		hostDirOrFile, ctrDirOrFile, err := getMountsMap(path)
		if err != nil {
			return nil, err
		}
		// skip if the hostDirOrFile path doesn't exist
		fileInfo, err := os.Stat(hostDirOrFile)
		if err != nil {
			if os.IsNotExist(err) {
				logrus.Warnf("Path %q from %q doesn't exist, skipping", hostDirOrFile, filePath)
				continue
			}
			return nil, errors.Wrapf(err, "failed to stat %q", hostDirOrFile)
		}

		ctrDirOrFileOnHost := filepath.Join(containerWorkingDir, ctrDirOrFile)

		// In the event of a restart, don't want to copy secrets over again as they already would exist in ctrDirOrFileOnHost
		_, err = os.Stat(ctrDirOrFileOnHost)
		if os.IsNotExist(err) {

			hostDirOrFile, err = resolveSymbolicLink(hostDirOrFile)
			if err != nil {
				return nil, err
			}

			switch mode := fileInfo.Mode(); {
			case mode.IsDir():
				if err = os.MkdirAll(ctrDirOrFileOnHost, 0755); err != nil {
					return nil, errors.Wrapf(err, "making container directory %q failed", ctrDirOrFileOnHost)
				}
				data, err := getHostSecretData(hostDirOrFile)
				if err != nil {
					return nil, errors.Wrapf(err, "getting host secret data failed")
				}
				for _, s := range data {
					if err := s.saveTo(ctrDirOrFileOnHost); err != nil {
						return nil, errors.Wrapf(err, "error saving data to container filesystem on host %q", ctrDirOrFileOnHost)
					}
				}
			case mode.IsRegular():
				data, err := readFile("", hostDirOrFile)
				if err != nil {
					return nil, errors.Wrapf(err, "error reading file %q", hostDirOrFile)

				}
				for _, s := range data {
					if err := os.MkdirAll(filepath.Dir(ctrDirOrFileOnHost), 0700); err != nil {
						return nil, err
					}
					if err := ioutil.WriteFile(ctrDirOrFileOnHost, s.data, 0700); err != nil {
						return nil, errors.Wrapf(err, "error saving data to container filesystem on host %q", ctrDirOrFileOnHost)
					}
				}
			default:
				return nil, errors.Errorf("unsupported file type for: %q", hostDirOrFile)
			}

			err = label.Relabel(ctrDirOrFileOnHost, mountLabel, false)
			if err != nil {
				return nil, errors.Wrap(err, "error applying correct labels")
			}
			if uid != 0 || gid != 0 {
				if err := rchown(ctrDirOrFileOnHost, uid, gid); err != nil {
					return nil, err
				}
			}
		} else if err != nil {
			return nil, errors.Wrapf(err, "error getting status of %q", ctrDirOrFileOnHost)
		}

		m := rspec.Mount{
			Source:      filepath.Join(mountPrefix, ctrDirOrFile),
			Destination: ctrDirOrFile,
			Type:        "bind",
			Options:     []string{"bind", "rprivate"},
		}

		mounts = append(mounts, m)
	}
	return mounts, nil
}

// addFIPSModeSecret creates /run/secrets/system-fips in the container
// root filesystem if /etc/system-fips exists on hosts.
// This enables the container to be FIPS compliant and run openssl in
// FIPS mode as the host is also in FIPS mode.
func addFIPSModeSecret(mounts *[]rspec.Mount, containerWorkingDir, mountPrefix, mountLabel string, uid, gid int) error {
	secretsDir := "/run/secrets"
	ctrDirOnHost := filepath.Join(containerWorkingDir, secretsDir)
	if _, err := os.Stat(ctrDirOnHost); os.IsNotExist(err) {
		if err = idtools.MkdirAllAs(ctrDirOnHost, 0755, uid, gid); err != nil {
			return errors.Wrapf(err, "making container directory on host failed")
		}
		if err = label.Relabel(ctrDirOnHost, mountLabel, false); err != nil {
			return errors.Wrap(err, "error applying correct labels")
		}
	}
	fipsFile := filepath.Join(ctrDirOnHost, "system-fips")
	// In the event of restart, it is possible for the FIPS mode file to already exist
	if _, err := os.Stat(fipsFile); os.IsNotExist(err) {
		file, err := os.Create(fipsFile)
		if err != nil {
			return errors.Wrapf(err, "error creating system-fips file in container for FIPS mode")
		}
		defer file.Close()
	}

	if !mountExists(*mounts, secretsDir) {
		m := rspec.Mount{
			Source:      filepath.Join(mountPrefix, secretsDir),
			Destination: secretsDir,
			Type:        "bind",
			Options:     []string{"bind", "rprivate"},
		}
		*mounts = append(*mounts, m)
	}

	return nil
}

// mountExists checks if a mount already exists in the spec
func mountExists(mounts []rspec.Mount, dest string) bool {
	for _, mount := range mounts {
		if mount.Destination == dest {
			return true
		}
	}
	return false
}

// resolveSymbolicLink resolves a possbile symlink path. If the path is a symlink, returns resolved
// path; if not, returns the original path.
func resolveSymbolicLink(path string) (string, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return "", err
	}
	if info.Mode()&os.ModeSymlink != os.ModeSymlink {
		return path, nil
	}
	return filepath.EvalSymlinks(path)
}
<|endoftext|>


<commit_before>package util

import (
	"testing"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/hltcoe/goncrete"
	"github.com/maxthomas/hardhat/pkg/sample"
	"github.com/stretchr/testify/assert"
)

var (
	fact = thrift.NewTCompactProtocolFactory()
)

func TestSerPool(t *testing.T) {
	pool := SerializerPool(fact)

	nItems := 100
	comms := make(chan *goncrete.Communication, nItems)
	for i := 0; i < nItems; i++ {
		comms <- sample.Communication()
	}
	close(comms)
	for i := 0; i < 3; i++ {
		go func() {
			for item := range comms {
				ser := pool.Get().(*thrift.TSerializer)

				bytez, err := ser.Write(item)
				if assert.NoError(t, err) {
					assert.FailNow(t, "failed parallel serialization")
				}

				assert.NotNil(t, bytez)
				ser.Transport.Reset()
				pool.Put(ser)
			}
		}()
	}
}
<commit_msg>fix up a test<commit_after>package util

import (
	"sync"
	"testing"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/hltcoe/goncrete"
	"github.com/maxthomas/hardhat/pkg/sample"
	"github.com/stretchr/testify/assert"
)

var (
	fact = thrift.NewTCompactProtocolFactory()
)

func TestSerPool(t *testing.T) {
	pool := SerializerPool(fact)

	nItems := 100
	comms := make(chan *goncrete.Communication, nItems)
	for i := 0; i < nItems; i++ {
		comms <- sample.Communication()
	}
	close(comms)

	errs := make(chan error, 4)
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ser := pool.Get().(*thrift.TSerializer)
			defer pool.Put(ser)
			ser.Transport.Reset()

			for item := range comms {
				bytez, err := ser.Write(item)
				if !assert.NoError(t, err) {
					errs <- err
					return
				}

				if !assert.NotNil(t, bytez) {
					errs <- err
					return
				}
			}
		}()
	}
	wg.Wait()
	if len(errs) > 0 {
		firstErr := <-errs
		assert.FailNow(t, "at least one error", firstErr.Error())
	}
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 version // import "k8s.io/helm/pkg/version"

import "k8s.io/helm/pkg/proto/hapi/version"

var (
	// Version is the current version of the Helm.
	// Update this whenever making a new release.
	// The version is of the format Major.Minor.Patch[-Prerelease][+BuildMetadata]
	//
	// Increment major number for new feature additions and behavioral changes.
	// Increment minor number for bug fixes and performance enhancements.
	// Increment patch number for critical fixes to existing releases.
	Version = "v2.10"

	// BuildMetadata is extra build time data
	BuildMetadata = "unreleased"
	// GitCommit is the git sha1
	GitCommit = ""
	// GitTreeState is the state of the git tree
	GitTreeState = ""
)

// GetVersion returns the semver string of the version
func GetVersion() string {
	if BuildMetadata == "" {
		return Version
	}
	return Version + "+" + BuildMetadata
}

// GetVersionProto returns protobuf representing the version
func GetVersionProto() *version.Version {
	return &version.Version{
		SemVer:       GetVersion(),
		GitCommit:    GitCommit,
		GitTreeState: GitTreeState,
	}
}
<commit_msg>bump version to v2.11 (#4700)<commit_after>/*
Copyright The Helm 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 version // import "k8s.io/helm/pkg/version"

import "k8s.io/helm/pkg/proto/hapi/version"

var (
	// Version is the current version of the Helm.
	// Update this whenever making a new release.
	// The version is of the format Major.Minor.Patch[-Prerelease][+BuildMetadata]
	//
	// Increment major number for new feature additions and behavioral changes.
	// Increment minor number for bug fixes and performance enhancements.
	// Increment patch number for critical fixes to existing releases.
	Version = "v2.11"

	// BuildMetadata is extra build time data
	BuildMetadata = "unreleased"
	// GitCommit is the git sha1
	GitCommit = ""
	// GitTreeState is the state of the git tree
	GitTreeState = ""
)

// GetVersion returns the semver string of the version
func GetVersion() string {
	if BuildMetadata == "" {
		return Version
	}
	return Version + "+" + BuildMetadata
}

// GetVersionProto returns protobuf representing the version
func GetVersionProto() *version.Version {
	return &version.Version{
		SemVer:       GetVersion(),
		GitCommit:    GitCommit,
		GitTreeState: GitTreeState,
	}
}
<|endoftext|>


<commit_before>package contracts

type Loggable interface {
	Hello(msg string)

	//
	//	Info(msg string, args ...interface{})
	//	//
	//	Debug(msg string, args ...interface{})
	//
	//	//
	//	Warn(msg string, args ...interface{})
	//
	//	//
	//	Error(msg string, args ...interface{})
	//
	//	//
	//	Trace(msg string, args ...interface{})
	//
	//	//
	//	Fatal(msg string, args ...interface{})
}
<commit_msg>Added logger constructor<commit_after>package contracts

type Loggable interface {
	NewLogger() Loggable

	Info(msg string, args ...interface{})

	Debug(msg string, args ...interface{})

	Warn(msg string, args ...interface{})

	Error(msg string, args ...interface{})

	Fatal(msg string, args ...interface{})

	Panic(msg string, args ...interface{})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"

	"github.com/citadel/citadel"
	"github.com/citadel/citadel/cluster"
	"github.com/citadel/citadel/scheduler"
	"github.com/ehazlett/interlock"
)

const (
	haproxyTmpl = `# managed by interlock
global
    {{ if .Config.SyslogAddr }}log {{ .Config.SyslogAddr }} local0
    log-send-hostname{{ end }}
    maxconn {{ .Config.MaxConn }}
    pidfile {{ .Config.PidPath }}

defaults
    mode http
    retries 3
    option redispatch
    option httplog
    option dontlognull
    timeout connect {{ .Config.ConnectTimeout }}
    timeout client {{ .Config.ClientTimeout }}
    timeout server {{ .Config.ServerTimeout }}

frontend http-default
    bind *:{{ .Config.Port }}
    monitor-uri /haproxy?monitor
    {{ if .Config.StatsUser }}stats realm Stats
    stats auth {{ .Config.StatsUser }}:{{ .Config.StatsPassword }}{{ end }}
    stats enable
    stats uri /haproxy?stats
    stats refresh 5s
    {{ range $host := .Hosts }}acl is_{{ $host.Name }} hdr_beg(host) {{ $host.Domain }}
    use_backend {{ $host.Name }} if is_{{ $host.Name }}
    {{ end }}
{{ range $host := .Hosts }}backend {{ $host.Name }}
    http-response add-header X-Request-Start t=%[date]
    balance roundrobin
    option forwardfor
    {{ range $option := $host.BackendOptions }}option {{ $option }}
    {{ end }}
    {{ if $host.Check }}option {{ $host.Check }}{{ end }}
    {{ range $i,$up := $host.Upstreams }}server {{ $host.Name }}_{{ $i }} {{ $up.Addr }} check inter {{ $up.CheckInterval }}
    {{ end }}
{{ end }}`
)

type (
	Manager struct {
		mux      sync.Mutex
		config   *interlock.Config
		engines  []*citadel.Engine
		cluster  *cluster.Cluster
		proxyCmd *exec.Cmd
	}
)

func NewManager(cfg *interlock.Config) (*Manager, error) {
	engines := []*citadel.Engine{}
	for _, e := range cfg.InterlockEngines {
		if err := e.Engine.Connect(nil); err != nil {
			return nil, err
		}
		engines = append(engines, e.Engine)
	}
	m := &Manager{
		config:  cfg,
		engines: engines,
	}
	if err := m.init(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Manager) init() error {
	c, err := cluster.New(scheduler.NewResourceManager(), m.engines...)
	if err != nil {
		return err
	}
	m.cluster = c
	return nil
}

func (m *Manager) writeConfig(config *interlock.ProxyConfig) error {
	m.mux.Lock()
	defer m.mux.Unlock()
	f, err := os.OpenFile(m.config.ProxyConfigPath, os.O_WRONLY|os.O_TRUNC, 0664)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		ff, fErr := os.Create(m.config.ProxyConfigPath)
		defer ff.Close()
		if fErr != nil {
			return fErr
		}
		f = ff
	}
	defer f.Close()
	t := template.New("haproxy")
	tmpl, err := t.Parse(haproxyTmpl)
	if err != nil {
		return err
	}
	var c bytes.Buffer
	if err := tmpl.Execute(&c, config); err != nil {
		return err
	}
	_, fErr := f.Write(c.Bytes())
	if fErr != nil {
		return fErr
	}
	f.Sync()
	return nil
}

func (m *Manager) GenerateProxyConfig(isKillEvent bool) (*interlock.ProxyConfig, error) {
	containers, err := m.cluster.ListContainers()
	if err != nil {
		return nil, err
	}
	var hosts []*interlock.Host
	proxyUpstreams := map[string][]*interlock.Upstream{}
	hostChecks := map[string]string{}
	hostBackendOptions := map[string][]string{}
	for _, cnt := range containers {
		if cnt.Image.Domainname == "" {
			continue
		}
		cntId := cnt.ID[:12]
		// load interlock data
		env := cnt.Image.Environment
		interlockData := &interlock.InterlockData{}
		if key, ok := env["INTERLOCK_DATA"]; ok {
			b := bytes.NewBufferString(key)
			if err := json.NewDecoder(b).Decode(&interlockData); err != nil {
				logger.Warnf("%s: unable to parse interlock data: %s", cntId, err)
			}
		}
		hostname := cnt.Image.Hostname
		domain := cnt.Image.Domainname
		if hostname != domain && hostname != "" {
			domain = fmt.Sprintf("%s.%s", hostname, domain)
		}
		if interlockData.Check != "" {
			if val, ok := hostChecks[domain]; ok {
				// check existing host check for different values
				if val != interlockData.Check {
					logger.Warnf("conflicting check specified for %s", domain)
				}
			} else {
				hostChecks[domain] = interlockData.Check
				logger.Infof("using custom check for %s: %s", domain, interlockData.Check)
			}
		}
		checkInterval := 5000
		if interlockData.CheckInterval != 0 {
			checkInterval = interlockData.CheckInterval
			logger.Infof("using custom check interval for %s: %d", domain, checkInterval)
		}
		if len(interlockData.BackendOptions) > 0 {
			hostBackendOptions[domain] = interlockData.BackendOptions
			logger.Infof("using backend options for %s: %s", domain, strings.Join(interlockData.BackendOptions, ","))
		}
		hostAddrUrl, err := url.Parse(cnt.Engine.Addr)
		if err != nil {
			logger.Warnf("%s: unable to parse engine addr: %s", cntId, err)
			continue
		}
		host := hostAddrUrl.Host
		hostParts := strings.Split(hostAddrUrl.Host, ":")
		if len(hostParts) != 1 {
			host = hostParts[0]
		}
		if len(cnt.Ports) == 0 {
			logger.Warnf("%s: no ports exposed", cntId)
			continue
		}
		portDef := cnt.Ports[0]
		addr := fmt.Sprintf("%s:%d", host, portDef.Port)
		if interlockData.Port != 0 {
			for _, p := range cnt.Ports {
				if p.ContainerPort == interlockData.Port {
					addr = fmt.Sprintf("%s:%d", host, p.Port)
				}
			}
		}
		up := &interlock.Upstream{
			Addr:          addr,
			CheckInterval: checkInterval,
		}
		for _, alias := range interlockData.AliasDomains {
			logger.Infof("adding alias %s for %s", alias, cntId)
			proxyUpstreams[alias] = append(proxyUpstreams[alias], up)
		}
		proxyUpstreams[domain] = append(proxyUpstreams[domain], up)
		if !isKillEvent && interlockData.Warm {
			logger.Infof("warming %s: %s", cntId, addr)
			http.Get(fmt.Sprintf("http://%s", addr))
		}

	}
	for k, v := range proxyUpstreams {
		name := strings.Replace(k, ".", "_", -1)
		host := &interlock.Host{
			Name:           name,
			Domain:         k,
			Upstreams:      v,
			Check:          hostChecks[k],
			BackendOptions: hostBackendOptions[k],
		}
		logger.Infof("adding host name=%s domain=%s", host.Name, host.Domain)
		hosts = append(hosts, host)
	}
	// generate config
	cfg := &interlock.ProxyConfig{
		Hosts:  hosts,
		Config: m.config,
	}
	return cfg, nil
}

func (m *Manager) UpdateConfig(e *citadel.Event) error {
	isKillEvent := false
	if e != nil && e.Type == "kill" {
		isKillEvent = true
	}
	cfg, err := m.GenerateProxyConfig(isKillEvent)
	if err != nil {
		return err
	}
	if err := m.writeConfig(cfg); err != nil {
		return err
	}
	return nil
}

func (m *Manager) getProxyPid() (int, error) {
	f, err := ioutil.ReadFile(m.config.PidPath)
	if err != nil {
		return -1, err
	}
	buf := bytes.NewBuffer(f)
	p := buf.String()
	p = strings.TrimSpace(p)
	pid, err := strconv.Atoi(p)
	if err != nil {
		return -1, err
	}
	return pid, nil
}

func (m *Manager) Reload() error {
	args := []string{"-D", "-f", m.config.ProxyConfigPath, "-p", m.config.PidPath, "-sf"}
	if m.proxyCmd != nil {
		p, err := m.getProxyPid()
		if err != nil {
			logger.Error(err)
		}
		pid := strconv.Itoa(p)
		args = append(args, pid)
	}
	cmd := exec.Command("haproxy", args...)
	if err := cmd.Run(); err != nil {
		return err
	}
	m.proxyCmd = cmd
	logger.Info("reloaded proxy")
	return nil
}

func (m *Manager) Run() error {
	if err := m.UpdateConfig(nil); err != nil {
		return err
	}
	if err := m.cluster.Events(&EventHandler{Manager: m}); err != nil {
		return err
	}
	m.Reload()
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	signal.Notify(ch, syscall.SIGTERM)
	go func() {
		<-ch
		if m.proxyCmd != nil {
			pid, err := m.getProxyPid()
			if err != nil {
				logger.Fatal(err)
			}
			syscall.Kill(pid, syscall.SIGTERM)
		}
		os.Exit(1)
	}()

	for {
		time.Sleep(1 * time.Second)
	}
}
<commit_msg>convert request time to ms<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"text/template"
	"time"

	"github.com/citadel/citadel"
	"github.com/citadel/citadel/cluster"
	"github.com/citadel/citadel/scheduler"
	"github.com/ehazlett/interlock"
)

const (
	haproxyTmpl = `# managed by interlock
global
    {{ if .Config.SyslogAddr }}log {{ .Config.SyslogAddr }} local0
    log-send-hostname{{ end }}
    maxconn {{ .Config.MaxConn }}
    pidfile {{ .Config.PidPath }}

defaults
    mode http
    retries 3
    option redispatch
    option httplog
    option dontlognull
    timeout connect {{ .Config.ConnectTimeout }}
    timeout client {{ .Config.ClientTimeout }}
    timeout server {{ .Config.ServerTimeout }}

frontend http-default
    bind *:{{ .Config.Port }}
    monitor-uri /haproxy?monitor
    {{ if .Config.StatsUser }}stats realm Stats
    stats auth {{ .Config.StatsUser }}:{{ .Config.StatsPassword }}{{ end }}
    stats enable
    stats uri /haproxy?stats
    stats refresh 5s
    {{ range $host := .Hosts }}acl is_{{ $host.Name }} hdr_beg(host) {{ $host.Domain }}
    use_backend {{ $host.Name }} if is_{{ $host.Name }}
    {{ end }}
{{ range $host := .Hosts }}backend {{ $host.Name }}
    http-response add-header X-Request-Start t=%T.%ms
    balance roundrobin
    option forwardfor
    {{ range $option := $host.BackendOptions }}option {{ $option }}
    {{ end }}
    {{ if $host.Check }}option {{ $host.Check }}{{ end }}
    {{ range $i,$up := $host.Upstreams }}server {{ $host.Name }}_{{ $i }} {{ $up.Addr }} check inter {{ $up.CheckInterval }}
    {{ end }}
{{ end }}`
)

type (
	Manager struct {
		mux      sync.Mutex
		config   *interlock.Config
		engines  []*citadel.Engine
		cluster  *cluster.Cluster
		proxyCmd *exec.Cmd
	}
)

func NewManager(cfg *interlock.Config) (*Manager, error) {
	engines := []*citadel.Engine{}
	for _, e := range cfg.InterlockEngines {
		if err := e.Engine.Connect(nil); err != nil {
			return nil, err
		}
		engines = append(engines, e.Engine)
	}
	m := &Manager{
		config:  cfg,
		engines: engines,
	}
	if err := m.init(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *Manager) init() error {
	c, err := cluster.New(scheduler.NewResourceManager(), m.engines...)
	if err != nil {
		return err
	}
	m.cluster = c
	return nil
}

func (m *Manager) writeConfig(config *interlock.ProxyConfig) error {
	m.mux.Lock()
	defer m.mux.Unlock()
	f, err := os.OpenFile(m.config.ProxyConfigPath, os.O_WRONLY|os.O_TRUNC, 0664)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
		ff, fErr := os.Create(m.config.ProxyConfigPath)
		defer ff.Close()
		if fErr != nil {
			return fErr
		}
		f = ff
	}
	defer f.Close()
	t := template.New("haproxy")
	tmpl, err := t.Parse(haproxyTmpl)
	if err != nil {
		return err
	}
	var c bytes.Buffer
	if err := tmpl.Execute(&c, config); err != nil {
		return err
	}
	_, fErr := f.Write(c.Bytes())
	if fErr != nil {
		return fErr
	}
	f.Sync()
	return nil
}

func (m *Manager) GenerateProxyConfig(isKillEvent bool) (*interlock.ProxyConfig, error) {
	containers, err := m.cluster.ListContainers()
	if err != nil {
		return nil, err
	}
	var hosts []*interlock.Host
	proxyUpstreams := map[string][]*interlock.Upstream{}
	hostChecks := map[string]string{}
	hostBackendOptions := map[string][]string{}
	for _, cnt := range containers {
		if cnt.Image.Domainname == "" {
			continue
		}
		cntId := cnt.ID[:12]
		// load interlock data
		env := cnt.Image.Environment
		interlockData := &interlock.InterlockData{}
		if key, ok := env["INTERLOCK_DATA"]; ok {
			b := bytes.NewBufferString(key)
			if err := json.NewDecoder(b).Decode(&interlockData); err != nil {
				logger.Warnf("%s: unable to parse interlock data: %s", cntId, err)
			}
		}
		hostname := cnt.Image.Hostname
		domain := cnt.Image.Domainname
		if hostname != domain && hostname != "" {
			domain = fmt.Sprintf("%s.%s", hostname, domain)
		}
		if interlockData.Check != "" {
			if val, ok := hostChecks[domain]; ok {
				// check existing host check for different values
				if val != interlockData.Check {
					logger.Warnf("conflicting check specified for %s", domain)
				}
			} else {
				hostChecks[domain] = interlockData.Check
				logger.Infof("using custom check for %s: %s", domain, interlockData.Check)
			}
		}
		checkInterval := 5000
		if interlockData.CheckInterval != 0 {
			checkInterval = interlockData.CheckInterval
			logger.Infof("using custom check interval for %s: %d", domain, checkInterval)
		}
		if len(interlockData.BackendOptions) > 0 {
			hostBackendOptions[domain] = interlockData.BackendOptions
			logger.Infof("using backend options for %s: %s", domain, strings.Join(interlockData.BackendOptions, ","))
		}
		hostAddrUrl, err := url.Parse(cnt.Engine.Addr)
		if err != nil {
			logger.Warnf("%s: unable to parse engine addr: %s", cntId, err)
			continue
		}
		host := hostAddrUrl.Host
		hostParts := strings.Split(hostAddrUrl.Host, ":")
		if len(hostParts) != 1 {
			host = hostParts[0]
		}
		if len(cnt.Ports) == 0 {
			logger.Warnf("%s: no ports exposed", cntId)
			continue
		}
		portDef := cnt.Ports[0]
		addr := fmt.Sprintf("%s:%d", host, portDef.Port)
		if interlockData.Port != 0 {
			for _, p := range cnt.Ports {
				if p.ContainerPort == interlockData.Port {
					addr = fmt.Sprintf("%s:%d", host, p.Port)
				}
			}
		}
		up := &interlock.Upstream{
			Addr:          addr,
			CheckInterval: checkInterval,
		}
		for _, alias := range interlockData.AliasDomains {
			logger.Infof("adding alias %s for %s", alias, cntId)
			proxyUpstreams[alias] = append(proxyUpstreams[alias], up)
		}
		proxyUpstreams[domain] = append(proxyUpstreams[domain], up)
		if !isKillEvent && interlockData.Warm {
			logger.Infof("warming %s: %s", cntId, addr)
			http.Get(fmt.Sprintf("http://%s", addr))
		}

	}
	for k, v := range proxyUpstreams {
		name := strings.Replace(k, ".", "_", -1)
		host := &interlock.Host{
			Name:           name,
			Domain:         k,
			Upstreams:      v,
			Check:          hostChecks[k],
			BackendOptions: hostBackendOptions[k],
		}
		logger.Infof("adding host name=%s domain=%s", host.Name, host.Domain)
		hosts = append(hosts, host)
	}
	// generate config
	cfg := &interlock.ProxyConfig{
		Hosts:  hosts,
		Config: m.config,
	}
	return cfg, nil
}

func (m *Manager) UpdateConfig(e *citadel.Event) error {
	isKillEvent := false
	if e != nil && e.Type == "kill" {
		isKillEvent = true
	}
	cfg, err := m.GenerateProxyConfig(isKillEvent)
	if err != nil {
		return err
	}
	if err := m.writeConfig(cfg); err != nil {
		return err
	}
	return nil
}

func (m *Manager) getProxyPid() (int, error) {
	f, err := ioutil.ReadFile(m.config.PidPath)
	if err != nil {
		return -1, err
	}
	buf := bytes.NewBuffer(f)
	p := buf.String()
	p = strings.TrimSpace(p)
	pid, err := strconv.Atoi(p)
	if err != nil {
		return -1, err
	}
	return pid, nil
}

func (m *Manager) Reload() error {
	args := []string{"-D", "-f", m.config.ProxyConfigPath, "-p", m.config.PidPath, "-sf"}
	if m.proxyCmd != nil {
		p, err := m.getProxyPid()
		if err != nil {
			logger.Error(err)
		}
		pid := strconv.Itoa(p)
		args = append(args, pid)
	}
	cmd := exec.Command("haproxy", args...)
	if err := cmd.Run(); err != nil {
		return err
	}
	m.proxyCmd = cmd
	logger.Info("reloaded proxy")
	return nil
}

func (m *Manager) Run() error {
	if err := m.UpdateConfig(nil); err != nil {
		return err
	}
	if err := m.cluster.Events(&EventHandler{Manager: m}); err != nil {
		return err
	}
	m.Reload()
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	signal.Notify(ch, syscall.SIGTERM)
	go func() {
		<-ch
		if m.proxyCmd != nil {
			pid, err := m.getProxyPid()
			if err != nil {
				logger.Fatal(err)
			}
			syscall.Kill(pid, syscall.SIGTERM)
		}
		os.Exit(1)
	}()

	for {
		time.Sleep(1 * time.Second)
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	oldcmds "github.com/ipfs/go-ipfs/commands"
	lgc "github.com/ipfs/go-ipfs/commands/legacy"
	e "github.com/ipfs/go-ipfs/core/commands/e"
	corerepo "github.com/ipfs/go-ipfs/core/corerepo"
	config "github.com/ipfs/go-ipfs/repo/config"
	fsrepo "github.com/ipfs/go-ipfs/repo/fsrepo"

	cmds "gx/ipfs/QmNueRyPRQiV7PUEpnP4GgGLuK1rKQLaRW7sfPvUetYig1/go-ipfs-cmds"
	cid "gx/ipfs/QmapdYm1b22Frv3k17fqrBYTFRxwiaVJkB299Mfn33edeB/go-cid"
	cmdkit "gx/ipfs/QmdE4gMduCKCGAcczM2F5ioYDfdeKuPix138wrES1YSr7f/go-ipfs-cmdkit"
	bstore "gx/ipfs/QmdpuJBPBZ6sLPj9BQpn3Rpi38BT2cF1QMiUfyzNWeySW4/go-ipfs-blockstore"
)

type RepoVersion struct {
	Version string
}

var RepoCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Manipulate the IPFS repo.",
		ShortDescription: `
'ipfs repo' is a plumbing command used to manipulate the repo.
`,
	},

	Subcommands: map[string]*cmds.Command{
		"stat":    repoStatCmd,
		"gc":      lgc.NewCommand(repoGcCmd),
		"fsck":    lgc.NewCommand(RepoFsckCmd),
		"version": lgc.NewCommand(repoVersionCmd),
		"verify":  lgc.NewCommand(repoVerifyCmd),
	},
}

// GcResult is the result returned by "repo gc" command.
type GcResult struct {
	Key   *cid.Cid
	Error string `json:",omitempty"`
}

var repoGcCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Perform a garbage collection sweep on the repo.",
		ShortDescription: `
'ipfs repo gc' is a plumbing command that will sweep the local
set of stored objects and remove ones that are not pinned in
order to reclaim hard disk space.
`,
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("stream-errors", "Stream errors."),
		cmdkit.BoolOption("quiet", "q", "Write minimal output."),
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		streamErrors, _, _ := res.Request().Option("stream-errors").Bool()

		gcOutChan := corerepo.GarbageCollectAsync(n, req.Context())

		outChan := make(chan interface{})
		res.SetOutput(outChan)

		go func() {
			defer close(outChan)

			if streamErrors {
				errs := false
				for res := range gcOutChan {
					if res.Error != nil {
						select {
						case outChan <- &GcResult{Error: res.Error.Error()}:
						case <-req.Context().Done():
							return
						}
						errs = true
					} else {
						select {
						case outChan <- &GcResult{Key: res.KeyRemoved}:
						case <-req.Context().Done():
							return
						}
					}
				}
				if errs {
					res.SetError(fmt.Errorf("encountered errors during gc run"), cmdkit.ErrNormal)
				}
			} else {
				err := corerepo.CollectResult(req.Context(), gcOutChan, func(k *cid.Cid) {
					select {
					case outChan <- &GcResult{Key: k}:
					case <-req.Context().Done():
					}
				})
				if err != nil {
					res.SetError(err, cmdkit.ErrNormal)
				}
			}
		}()
	},
	Type: GcResult{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}

			quiet, _, err := res.Request().Option("quiet").Bool()
			if err != nil {
				return nil, err
			}

			obj, ok := v.(*GcResult)
			if !ok {
				return nil, e.TypeErr(obj, v)
			}

			if obj.Error != "" {
				fmt.Fprintf(res.Stderr(), "Error: %s\n", obj.Error)
				return nil, nil
			}

			msg := obj.Key.String() + "\n"
			if !quiet {
				msg = "removed " + msg
			}

			return bytes.NewBufferString(msg), nil
		},
	},
}

var repoStatCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Get stats for the currently used repo.",
		ShortDescription: `
'ipfs repo stat' provides information about the local set of
stored objects. It outputs:

RepoSize        int Size in bytes that the repo is currently taking.
StorageMax      string Maximum datastore size (from configuration)
NumObjects      int Number of objects in the local repo.
RepoPath        string The path to the repo being currently used.
Version         string The repo version.
`,
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("size-only", "Only report RepoSize and StorageMax."),
		cmdkit.BoolOption("human", "Output sizes in MiB."),
	},
	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) {
		n, err := GetNode(env)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		sizeOnly, _ := req.Options["size-only"].(bool)
		if sizeOnly {
			sizeStat, err := corerepo.RepoSize(req.Context, n)
			if err != nil {
				res.SetError(err, cmdkit.ErrNormal)
				return
			}
			cmds.EmitOnce(res, &corerepo.Stat{
				SizeStat: sizeStat,
			})
			return
		}

		stat, err := corerepo.RepoStat(req.Context, n)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		cmds.EmitOnce(res, &stat)
	},
	Type: &corerepo.Stat{},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w io.Writer, v interface{}) error {
			stat, ok := v.(*corerepo.Stat)
			if !ok {
				fmt.Println("adios")
				return e.TypeErr(stat, v)
			}

			wtr := tabwriter.NewWriter(w, 0, 0, 1, ' ', 0)
			defer wtr.Flush()

			human, _ := req.Options["human"].(bool)
			sizeOnly, _ := req.Options["size-only"].(bool)

			printSize := func(name string, size uint64) {
				sizeInMiB := size / (1024 * 1024)
				if human && sizeInMiB > 0 {
					fmt.Fprintf(wtr, "%s (MiB):\t%d\n", name, sizeInMiB)
				} else {
					fmt.Fprintf(wtr, "%s:\t%d\n", name, size)
				}
			}

			if !sizeOnly {
				fmt.Fprintf(wtr, "NumObjects:\t%d\n", stat.NumObjects)
			}

			printSize("RepoSize", stat.RepoSize)
			printSize("StorageMax", stat.StorageMax)

			if !sizeOnly {
				fmt.Fprintf(wtr, "RepoPath:\t%s\n", stat.RepoPath)
				fmt.Fprintf(wtr, "Version:\t%s\n", stat.Version)
			}

			return nil
		}),
	},
}

var RepoFsckCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Remove repo lockfiles.",
		ShortDescription: `
'ipfs repo fsck' is a plumbing command that will remove repo and level db
lockfiles, as well as the api file. This command can only run when no ipfs
daemons are running.
`,
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		configRoot := req.InvocContext().ConfigRoot

		dsPath, err := config.DataStorePath(configRoot)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		dsLockFile := filepath.Join(dsPath, "LOCK") // TODO: get this lockfile programmatically
		repoLockFile := filepath.Join(configRoot, fsrepo.LockFile)
		apiFile := filepath.Join(configRoot, "api") // TODO: get this programmatically

		log.Infof("Removing repo lockfile: %s", repoLockFile)
		log.Infof("Removing datastore lockfile: %s", dsLockFile)
		log.Infof("Removing api file: %s", apiFile)

		err = os.Remove(repoLockFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}
		err = os.Remove(dsLockFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}
		err = os.Remove(apiFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		res.SetOutput(&MessageOutput{"Lockfiles have been removed.\n"})
	},
	Type: MessageOutput{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: MessageTextMarshaler,
	},
}

type VerifyProgress struct {
	Msg      string
	Progress int
}

var repoVerifyCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Verify all blocks in repo are not corrupted.",
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		nd, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))
		defer close(out)

		bs := bstore.NewBlockstore(nd.Repo.Datastore())
		bs.HashOnRead(true)

		keys, err := bs.AllKeysChan(req.Context())
		if err != nil {
			log.Error(err)
			return
		}

		var fails int
		var i int
		for k := range keys {
			_, err := bs.Get(k)
			if err != nil {
				select {
				case out <- &VerifyProgress{
					Msg: fmt.Sprintf("block %s was corrupt (%s)", k, err),
				}:
				case <-req.Context().Done():
					return
				}
				fails++
			}
			i++
			select {
			case out <- &VerifyProgress{Progress: i}:
			case <-req.Context().Done():
				return
			}
		}

		if fails == 0 {
			select {
			case out <- &VerifyProgress{Msg: "verify complete, all blocks validated."}:
			case <-req.Context().Done():
				return
			}
		} else {
			res.SetError(fmt.Errorf("verify complete, some blocks were corrupt"), cmdkit.ErrNormal)
		}
	},
	Type: &VerifyProgress{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}

			obj, ok := v.(*VerifyProgress)
			if !ok {
				return nil, e.TypeErr(obj, v)
			}

			buf := new(bytes.Buffer)
			if strings.Contains(obj.Msg, "was corrupt") {
				fmt.Fprintln(os.Stdout, obj.Msg)
				return buf, nil
			}

			if obj.Msg != "" {
				if len(obj.Msg) < 20 {
					obj.Msg += "             "
				}
				fmt.Fprintln(buf, obj.Msg)
				return buf, nil
			}

			fmt.Fprintf(buf, "%d blocks processed.\r", obj.Progress)
			return buf, nil
		},
	},
}

var repoVersionCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Show the repo version.",
		ShortDescription: `
'ipfs repo version' returns the current repo version.
`,
	},

	Options: []cmdkit.Option{
		cmdkit.BoolOption("quiet", "q", "Write minimal output."),
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		res.SetOutput(&RepoVersion{
			Version: fmt.Sprint(fsrepo.RepoVersion),
		})
	},
	Type: RepoVersion{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}
			response, ok := v.(*RepoVersion)
			if !ok {
				return nil, e.TypeErr(response, v)
			}

			quiet, _, err := res.Request().Option("quiet").Bool()
			if err != nil {
				return nil, err
			}

			buf := new(bytes.Buffer)
			if quiet {
				buf = bytes.NewBufferString(fmt.Sprintf("fs-repo@%s\n", response.Version))
			} else {
				buf = bytes.NewBufferString(fmt.Sprintf("ipfs repo version fs-repo@%s\n", response.Version))
			}
			return buf, nil

		},
	},
}
<commit_msg>remove debug print statement<commit_after>package commands

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	oldcmds "github.com/ipfs/go-ipfs/commands"
	lgc "github.com/ipfs/go-ipfs/commands/legacy"
	e "github.com/ipfs/go-ipfs/core/commands/e"
	corerepo "github.com/ipfs/go-ipfs/core/corerepo"
	config "github.com/ipfs/go-ipfs/repo/config"
	fsrepo "github.com/ipfs/go-ipfs/repo/fsrepo"

	cmds "gx/ipfs/QmNueRyPRQiV7PUEpnP4GgGLuK1rKQLaRW7sfPvUetYig1/go-ipfs-cmds"
	cid "gx/ipfs/QmapdYm1b22Frv3k17fqrBYTFRxwiaVJkB299Mfn33edeB/go-cid"
	cmdkit "gx/ipfs/QmdE4gMduCKCGAcczM2F5ioYDfdeKuPix138wrES1YSr7f/go-ipfs-cmdkit"
	bstore "gx/ipfs/QmdpuJBPBZ6sLPj9BQpn3Rpi38BT2cF1QMiUfyzNWeySW4/go-ipfs-blockstore"
)

type RepoVersion struct {
	Version string
}

var RepoCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Manipulate the IPFS repo.",
		ShortDescription: `
'ipfs repo' is a plumbing command used to manipulate the repo.
`,
	},

	Subcommands: map[string]*cmds.Command{
		"stat":    repoStatCmd,
		"gc":      lgc.NewCommand(repoGcCmd),
		"fsck":    lgc.NewCommand(RepoFsckCmd),
		"version": lgc.NewCommand(repoVersionCmd),
		"verify":  lgc.NewCommand(repoVerifyCmd),
	},
}

// GcResult is the result returned by "repo gc" command.
type GcResult struct {
	Key   *cid.Cid
	Error string `json:",omitempty"`
}

var repoGcCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Perform a garbage collection sweep on the repo.",
		ShortDescription: `
'ipfs repo gc' is a plumbing command that will sweep the local
set of stored objects and remove ones that are not pinned in
order to reclaim hard disk space.
`,
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("stream-errors", "Stream errors."),
		cmdkit.BoolOption("quiet", "q", "Write minimal output."),
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		streamErrors, _, _ := res.Request().Option("stream-errors").Bool()

		gcOutChan := corerepo.GarbageCollectAsync(n, req.Context())

		outChan := make(chan interface{})
		res.SetOutput(outChan)

		go func() {
			defer close(outChan)

			if streamErrors {
				errs := false
				for res := range gcOutChan {
					if res.Error != nil {
						select {
						case outChan <- &GcResult{Error: res.Error.Error()}:
						case <-req.Context().Done():
							return
						}
						errs = true
					} else {
						select {
						case outChan <- &GcResult{Key: res.KeyRemoved}:
						case <-req.Context().Done():
							return
						}
					}
				}
				if errs {
					res.SetError(fmt.Errorf("encountered errors during gc run"), cmdkit.ErrNormal)
				}
			} else {
				err := corerepo.CollectResult(req.Context(), gcOutChan, func(k *cid.Cid) {
					select {
					case outChan <- &GcResult{Key: k}:
					case <-req.Context().Done():
					}
				})
				if err != nil {
					res.SetError(err, cmdkit.ErrNormal)
				}
			}
		}()
	},
	Type: GcResult{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}

			quiet, _, err := res.Request().Option("quiet").Bool()
			if err != nil {
				return nil, err
			}

			obj, ok := v.(*GcResult)
			if !ok {
				return nil, e.TypeErr(obj, v)
			}

			if obj.Error != "" {
				fmt.Fprintf(res.Stderr(), "Error: %s\n", obj.Error)
				return nil, nil
			}

			msg := obj.Key.String() + "\n"
			if !quiet {
				msg = "removed " + msg
			}

			return bytes.NewBufferString(msg), nil
		},
	},
}

var repoStatCmd = &cmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Get stats for the currently used repo.",
		ShortDescription: `
'ipfs repo stat' provides information about the local set of
stored objects. It outputs:

RepoSize        int Size in bytes that the repo is currently taking.
StorageMax      string Maximum datastore size (from configuration)
NumObjects      int Number of objects in the local repo.
RepoPath        string The path to the repo being currently used.
Version         string The repo version.
`,
	},
	Options: []cmdkit.Option{
		cmdkit.BoolOption("size-only", "Only report RepoSize and StorageMax."),
		cmdkit.BoolOption("human", "Output sizes in MiB."),
	},
	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) {
		n, err := GetNode(env)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		sizeOnly, _ := req.Options["size-only"].(bool)
		if sizeOnly {
			sizeStat, err := corerepo.RepoSize(req.Context, n)
			if err != nil {
				res.SetError(err, cmdkit.ErrNormal)
				return
			}
			cmds.EmitOnce(res, &corerepo.Stat{
				SizeStat: sizeStat,
			})
			return
		}

		stat, err := corerepo.RepoStat(req.Context, n)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		cmds.EmitOnce(res, &stat)
	},
	Type: &corerepo.Stat{},
	Encoders: cmds.EncoderMap{
		cmds.Text: cmds.MakeEncoder(func(req *cmds.Request, w io.Writer, v interface{}) error {
			stat, ok := v.(*corerepo.Stat)
			if !ok {
				return e.TypeErr(stat, v)
			}

			wtr := tabwriter.NewWriter(w, 0, 0, 1, ' ', 0)
			defer wtr.Flush()

			human, _ := req.Options["human"].(bool)
			sizeOnly, _ := req.Options["size-only"].(bool)

			printSize := func(name string, size uint64) {
				sizeInMiB := size / (1024 * 1024)
				if human && sizeInMiB > 0 {
					fmt.Fprintf(wtr, "%s (MiB):\t%d\n", name, sizeInMiB)
				} else {
					fmt.Fprintf(wtr, "%s:\t%d\n", name, size)
				}
			}

			if !sizeOnly {
				fmt.Fprintf(wtr, "NumObjects:\t%d\n", stat.NumObjects)
			}

			printSize("RepoSize", stat.RepoSize)
			printSize("StorageMax", stat.StorageMax)

			if !sizeOnly {
				fmt.Fprintf(wtr, "RepoPath:\t%s\n", stat.RepoPath)
				fmt.Fprintf(wtr, "Version:\t%s\n", stat.Version)
			}

			return nil
		}),
	},
}

var RepoFsckCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Remove repo lockfiles.",
		ShortDescription: `
'ipfs repo fsck' is a plumbing command that will remove repo and level db
lockfiles, as well as the api file. This command can only run when no ipfs
daemons are running.
`,
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		configRoot := req.InvocContext().ConfigRoot

		dsPath, err := config.DataStorePath(configRoot)
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		dsLockFile := filepath.Join(dsPath, "LOCK") // TODO: get this lockfile programmatically
		repoLockFile := filepath.Join(configRoot, fsrepo.LockFile)
		apiFile := filepath.Join(configRoot, "api") // TODO: get this programmatically

		log.Infof("Removing repo lockfile: %s", repoLockFile)
		log.Infof("Removing datastore lockfile: %s", dsLockFile)
		log.Infof("Removing api file: %s", apiFile)

		err = os.Remove(repoLockFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}
		err = os.Remove(dsLockFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}
		err = os.Remove(apiFile)
		if err != nil && !os.IsNotExist(err) {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		res.SetOutput(&MessageOutput{"Lockfiles have been removed.\n"})
	},
	Type: MessageOutput{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: MessageTextMarshaler,
	},
}

type VerifyProgress struct {
	Msg      string
	Progress int
}

var repoVerifyCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Verify all blocks in repo are not corrupted.",
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		nd, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmdkit.ErrNormal)
			return
		}

		out := make(chan interface{})
		res.SetOutput((<-chan interface{})(out))
		defer close(out)

		bs := bstore.NewBlockstore(nd.Repo.Datastore())
		bs.HashOnRead(true)

		keys, err := bs.AllKeysChan(req.Context())
		if err != nil {
			log.Error(err)
			return
		}

		var fails int
		var i int
		for k := range keys {
			_, err := bs.Get(k)
			if err != nil {
				select {
				case out <- &VerifyProgress{
					Msg: fmt.Sprintf("block %s was corrupt (%s)", k, err),
				}:
				case <-req.Context().Done():
					return
				}
				fails++
			}
			i++
			select {
			case out <- &VerifyProgress{Progress: i}:
			case <-req.Context().Done():
				return
			}
		}

		if fails == 0 {
			select {
			case out <- &VerifyProgress{Msg: "verify complete, all blocks validated."}:
			case <-req.Context().Done():
				return
			}
		} else {
			res.SetError(fmt.Errorf("verify complete, some blocks were corrupt"), cmdkit.ErrNormal)
		}
	},
	Type: &VerifyProgress{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}

			obj, ok := v.(*VerifyProgress)
			if !ok {
				return nil, e.TypeErr(obj, v)
			}

			buf := new(bytes.Buffer)
			if strings.Contains(obj.Msg, "was corrupt") {
				fmt.Fprintln(os.Stdout, obj.Msg)
				return buf, nil
			}

			if obj.Msg != "" {
				if len(obj.Msg) < 20 {
					obj.Msg += "             "
				}
				fmt.Fprintln(buf, obj.Msg)
				return buf, nil
			}

			fmt.Fprintf(buf, "%d blocks processed.\r", obj.Progress)
			return buf, nil
		},
	},
}

var repoVersionCmd = &oldcmds.Command{
	Helptext: cmdkit.HelpText{
		Tagline: "Show the repo version.",
		ShortDescription: `
'ipfs repo version' returns the current repo version.
`,
	},

	Options: []cmdkit.Option{
		cmdkit.BoolOption("quiet", "q", "Write minimal output."),
	},
	Run: func(req oldcmds.Request, res oldcmds.Response) {
		res.SetOutput(&RepoVersion{
			Version: fmt.Sprint(fsrepo.RepoVersion),
		})
	},
	Type: RepoVersion{},
	Marshalers: oldcmds.MarshalerMap{
		oldcmds.Text: func(res oldcmds.Response) (io.Reader, error) {
			v, err := unwrapOutput(res.Output())
			if err != nil {
				return nil, err
			}
			response, ok := v.(*RepoVersion)
			if !ok {
				return nil, e.TypeErr(response, v)
			}

			quiet, _, err := res.Request().Option("quiet").Bool()
			if err != nil {
				return nil, err
			}

			buf := new(bytes.Buffer)
			if quiet {
				buf = bytes.NewBufferString(fmt.Sprintf("fs-repo@%s\n", response.Version))
			} else {
				buf = bytes.NewBufferString(fmt.Sprintf("ipfs repo version fs-repo@%s\n", response.Version))
			}
			return buf, nil

		},
	},
}
<|endoftext|>


<commit_before>package commands

import (
	"errors"
	"io"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	uio "github.com/jbenet/go-ipfs/unixfs/io"
)

var catCmd = &cmds.Command{
	Arguments: []cmds.Argument{
		cmds.Argument{"object", cmds.ArgString, false, true},
	},
	Help: `ipfs cat <ipfs-path> - Show ipfs object data.

	Retrieves the object named by <ipfs-path> and displays the Data
	it contains.
	`,
	Run: func(res cmds.Response, req cmds.Request) {
		node := req.Context().Node
		paths := make([]string, 0, len(req.Arguments()))
		readers := make([]io.Reader, 0, len(req.Arguments()))

		for _, arg := range req.Arguments() {
			path, ok := arg.(string)
			if !ok {
				res.SetError(errors.New("cast error"), cmds.ErrNormal)
				return
			}
			paths = append(paths, path)
		}

		readers, err := cat(node, paths)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		reader := io.MultiReader(readers...)
		res.SetOutput(reader)
	},
}

func cat(node *core.IpfsNode, paths []string) ([]io.Reader, error) {
	readers := make([]io.Reader, 0, len(paths))
	for _, path := range paths {
		dagnode, err := node.Resolver.ResolvePath(path)
		if err != nil {
			return nil, err
		}
		read, err := uio.NewDagReader(dagnode, node.DAG)
		if err != nil {
			return nil, err
		}
		readers = append(readers, read)
	}
	return readers, nil
}
<commit_msg>core/commands2: Made 'cat' require object argument<commit_after>package commands

import (
	"errors"
	"io"

	cmds "github.com/jbenet/go-ipfs/commands"
	core "github.com/jbenet/go-ipfs/core"
	uio "github.com/jbenet/go-ipfs/unixfs/io"
)

var catCmd = &cmds.Command{
	Arguments: []cmds.Argument{
		cmds.Argument{"object", cmds.ArgString, true, true},
	},
	Help: `ipfs cat <object> - Show ipfs object data.

	Retrieves the object named by <object> and outputs the data
	it contains.
	`,
	Run: func(res cmds.Response, req cmds.Request) {
		node := req.Context().Node
		paths := make([]string, 0, len(req.Arguments()))
		readers := make([]io.Reader, 0, len(req.Arguments()))

		for _, arg := range req.Arguments() {
			path, ok := arg.(string)
			if !ok {
				res.SetError(errors.New("cast error"), cmds.ErrNormal)
				return
			}
			paths = append(paths, path)
		}

		readers, err := cat(node, paths)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		reader := io.MultiReader(readers...)
		res.SetOutput(reader)
	},
}

func cat(node *core.IpfsNode, paths []string) ([]io.Reader, error) {
	readers := make([]io.Reader, 0, len(paths))
	for _, path := range paths {
		dagnode, err := node.Resolver.ResolvePath(path)
		if err != nil {
			return nil, err
		}
		read, err := uio.NewDagReader(dagnode, node.DAG)
		if err != nil {
			return nil, err
		}
		readers = append(readers, read)
	}
	return readers, nil
}
<|endoftext|>


<commit_before>package binary_search_tree

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNew(t *testing.T) {
	a := assert.New(t)

	root := NewNode(100)
	tree := New(NewNode(100))
	a.Equal(tree.Root(), root)
}

func TestTree_Find(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	tree := New(NewNode(5))
	left := NewNode(3)
	tree.root.left = left
	left.right = NewNode(4)
	tree.root.right = NewNode(6)

	t.Run("normal", func(t *testing.T) {
		tests := []struct {
			input int64
			want  *Node
		}{
			{input: 4, want: NewNode(4)},
		}
		for _, test := range tests {
			got, err := tree.Find(test.input)
			r.NoError(err)
			a.Equal(test.want, got)
		}
	})

	t.Run("not found", func(t *testing.T) {
		_, err := tree.Find(100)
		r.Equal(ErrNotFound, err)
	})
}

func TestTree_Insert(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)
	tests := []struct {
		insert       int64
		want         *Node
		wantTreeFunc func() *Tree
	}{
		{
			insert: 5,
			want:   NewNode(5),
			wantTreeFunc: func() *Tree {
				return New(NewNode(5))
			},
		},
		{
			insert: 3,
			want:   NewNode(3),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				return tree
			},
		},
		{
			insert: 6,
			want:   NewNode(6),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				tree.root.right = NewNode(6)
				return tree
			},
		},
		//     5
		//  ／    ＼
		// 3       6
		//  ＼
		//    4
		{
			insert: 4,
			want:   NewNode(4),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				left := NewNode(3)
				tree.root.left = left
				left.right = NewNode(4)
				tree.root.right = NewNode(6)
				return tree
			},
		},
	}

	tree := New(nil)
	for _, test := range tests {
		node, err := tree.Insert(test.insert)
		if err != nil {
			r.NoError(err, "tree.Insert failed")
		}
		a.Equal(test.want.Value(), node.Value())
		a.Equal(test.wantTreeFunc(), tree)
	}
}

func TestTree_Insert_Exist(t *testing.T) {
	r := require.New(t)
	tree := New(NewNode(5))
	tree.root.left = NewNode(3)
	tree.root.right = NewNode(6)

	_, err := tree.Insert(6)
	r.Equal(ErrAlreadyExists, err)
}

func TestTree_Remove(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)
	tests := map[string]struct {
		setupTree        func() *Tree
		target           int64
		removedAfterTree func() *Tree
	}{
		"remove leaf": {
			//     5
			//  ／    ＼
			// 3       6
			//  ＼
			//    4
			setupTree: func() *Tree {
				tree := New(NewNode(5))
				left := NewNode(3)
				tree.root.left = left
				left.right = NewNode(4)
				tree.root.right = NewNode(6)
				return tree
			},
			target: 4,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				tree.root.right = NewNode(6)
				return tree
			},
		},
		"remove node with left child": {
			//          9
			//        ／  ＼
			//       5     14
			//     ／
			//    3
			//  ／  ＼
			// 1     4
			setupTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(5)
				tree.root.right = NewNode(14)
				left2 := NewNode(3)
				left2.left = NewNode(1)
				left2.right = NewNode(4)
				tree.root.left.left = left2
				return tree
			},
			target: 5,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(9))
				left := NewNode(3)
				left.left = NewNode(1)
				left.right = NewNode(4)
				tree.root.left = left
				tree.root.right = NewNode(14)
				return tree
			},
		},
		"remove node with right child": {
			//       9
			//     ／  ＼
			//    2     10
			//            ＼
			//             17
			//           ／  ＼
			//          13    19
			setupTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(2)
				tree.root.right = NewNode(10)
				right2 := NewNode(17)
				right2.left = NewNode(13)
				right2.right = NewNode(19)
				tree.root.right.right = right2
				return tree
			},
			target: 10,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(2)
				right := NewNode(17)
				right.left = NewNode(13)
				right.right = NewNode(19)
				tree.root.right = right
				return tree
			},
		},
		"remove node with left and right child": {
			//           20
			//         ／   ＼
			//        7      23
			//     ／   ＼     ＼
			//    4      18     29
			//  ／  ＼   ／
			// 2     5  10
			//            ＼
			//             15
			setupTree: func() *Tree {
				tree := New(NewNode(20))
				tree.root.left = NewNode(7)
				tree.root.right = NewNode(23)

				left7 := NewNode(4)
				left7.left = NewNode(2)
				left7.right = NewNode(5)
				tree.root.left.left = left7

				right7 := NewNode(18)
				right7.left = NewNode(10)
				right7.left.right = NewNode(15)
				tree.root.left.right = right7

				tree.root.right.right = NewNode(29)

				return tree
			},
			target: 7,
			//           20
			//         ／   ＼
			//        10     23
			//     ／   ＼     ＼
			//    4      18     29
			//  ／  ＼   ／
			// 2     5  15
			removedAfterTree: func() *Tree {
				tree := New(NewNode(20))
				tree.root.left = NewNode(10)
				tree.root.right = NewNode(23)

				left10 := NewNode(4)
				left10.left = NewNode(2)
				left10.right = NewNode(5)
				tree.root.left.left = left10

				right10 := NewNode(18)
				right10.left = NewNode(15)
				tree.root.left.right = right10

				tree.root.right.right = NewNode(29)

				return tree
			},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			tree := test.setupTree()
			got, err := tree.Remove(test.target)
			r.NoError(err, "target is %v", test.target)
			a.Equal(test.target, got.Value())
			a.Equal(test.removedAfterTree(), tree)
		})
	}
}
<commit_msg>Add a test for Remove<commit_after>package binary_search_tree

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNew(t *testing.T) {
	a := assert.New(t)

	root := NewNode(100)
	tree := New(NewNode(100))
	a.Equal(tree.Root(), root)
}

func TestTree_Find(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)

	tree := New(NewNode(5))
	left := NewNode(3)
	tree.root.left = left
	left.right = NewNode(4)
	tree.root.right = NewNode(6)

	t.Run("normal", func(t *testing.T) {
		tests := []struct {
			input int64
			want  *Node
		}{
			{input: 4, want: NewNode(4)},
		}
		for _, test := range tests {
			got, err := tree.Find(test.input)
			r.NoError(err)
			a.Equal(test.want, got)
		}
	})

	t.Run("not found", func(t *testing.T) {
		_, err := tree.Find(100)
		r.Equal(ErrNotFound, err)
	})
}

func TestTree_Insert(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)
	tests := []struct {
		insert       int64
		want         *Node
		wantTreeFunc func() *Tree
	}{
		{
			insert: 5,
			want:   NewNode(5),
			wantTreeFunc: func() *Tree {
				return New(NewNode(5))
			},
		},
		{
			insert: 3,
			want:   NewNode(3),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				return tree
			},
		},
		{
			insert: 6,
			want:   NewNode(6),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				tree.root.right = NewNode(6)
				return tree
			},
		},
		//     5
		//  ／    ＼
		// 3       6
		//  ＼
		//    4
		{
			insert: 4,
			want:   NewNode(4),
			wantTreeFunc: func() *Tree {
				tree := New(NewNode(5))
				left := NewNode(3)
				tree.root.left = left
				left.right = NewNode(4)
				tree.root.right = NewNode(6)
				return tree
			},
		},
	}

	tree := New(nil)
	for _, test := range tests {
		node, err := tree.Insert(test.insert)
		if err != nil {
			r.NoError(err, "tree.Insert failed")
		}
		a.Equal(test.want.Value(), node.Value())
		a.Equal(test.wantTreeFunc(), tree)
	}
}

func TestTree_Insert_Exist(t *testing.T) {
	r := require.New(t)
	tree := New(NewNode(5))
	tree.root.left = NewNode(3)
	tree.root.right = NewNode(6)

	_, err := tree.Insert(6)
	r.Equal(ErrAlreadyExists, err)
}

func TestTree_Remove(t *testing.T) {
	a := assert.New(t)
	r := require.New(t)
	tests := map[string]struct {
		setupTree        func() *Tree
		target           int64
		removedAfterTree func() *Tree
	}{
		"remove leaf": {
			//     5
			//  ／    ＼
			// 3       6
			//  ＼
			//    4
			setupTree: func() *Tree {
				tree := New(NewNode(5))
				left := NewNode(3)
				tree.root.left = left
				left.right = NewNode(4)
				tree.root.right = NewNode(6)
				return tree
			},
			target: 4,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(5))
				tree.root.left = NewNode(3)
				tree.root.right = NewNode(6)
				return tree
			},
		},
		"remove node with left child": {
			//          9
			//        ／  ＼
			//       5     14
			//     ／
			//    3
			//  ／  ＼
			// 1     4
			setupTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(5)
				tree.root.right = NewNode(14)
				left2 := NewNode(3)
				left2.left = NewNode(1)
				left2.right = NewNode(4)
				tree.root.left.left = left2
				return tree
			},
			target: 5,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(9))
				left := NewNode(3)
				left.left = NewNode(1)
				left.right = NewNode(4)
				tree.root.left = left
				tree.root.right = NewNode(14)
				return tree
			},
		},
		"remove node with right child": {
			//       9
			//     ／  ＼
			//    2     10
			//            ＼
			//             17
			//           ／  ＼
			//          13    19
			setupTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(2)
				tree.root.right = NewNode(10)
				right2 := NewNode(17)
				right2.left = NewNode(13)
				right2.right = NewNode(19)
				tree.root.right.right = right2
				return tree
			},
			target: 10,
			removedAfterTree: func() *Tree {
				tree := New(NewNode(9))
				tree.root.left = NewNode(2)
				right := NewNode(17)
				right.left = NewNode(13)
				right.right = NewNode(19)
				tree.root.right = right
				return tree
			},
		},
		"remove node with left and right child": {
			//           20
			//         ／   ＼
			//        7      23
			//     ／   ＼     ＼
			//    4      18     29
			//  ／  ＼   ／
			// 2     5  10
			//            ＼
			//             15
			setupTree: func() *Tree {
				tree := New(NewNode(20))
				tree.root.left = NewNode(7)
				tree.root.right = NewNode(23)

				left7 := NewNode(4)
				left7.left = NewNode(2)
				left7.right = NewNode(5)
				tree.root.left.left = left7

				right7 := NewNode(18)
				right7.left = NewNode(10)
				right7.left.right = NewNode(15)
				tree.root.left.right = right7

				tree.root.right.right = NewNode(29)

				return tree
			},
			target: 7,
			//           20
			//         ／   ＼
			//        10     23
			//     ／   ＼     ＼
			//    4      18     29
			//  ／  ＼   ／
			// 2     5  15
			removedAfterTree: func() *Tree {
				tree := New(NewNode(20))
				tree.root.left = NewNode(10)
				tree.root.right = NewNode(23)

				left10 := NewNode(4)
				left10.left = NewNode(2)
				left10.right = NewNode(5)
				tree.root.left.left = left10

				right10 := NewNode(18)
				right10.left = NewNode(15)
				tree.root.left.right = right10

				tree.root.right.right = NewNode(29)

				return tree
			},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			tree := test.setupTree()
			got, err := tree.Remove(test.target)
			r.NoError(err, "target is %v", test.target)
			a.Equal(test.target, got.Value())
			a.Equal(test.removedAfterTree(), tree)
		})
	}
}

func TestTree_Remove_NotFound(t *testing.T) {
	a := assert.New(t)

	tree := New(NewNode(5))
	tree.root.left = NewNode(3)
	tree.root.right = NewNode(6)
	tree.root.right.right = NewNode(10)

	_, err := tree.Remove(15)
	a.Equal(ErrNotFound, err)
}
<|endoftext|>


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

package sha3

// This file provides functions for creating instances of the SHA-3
// and SHAKE hash functions, as well as utility functions for hashing
// bytes.

import (
	"hash"
)

// New224 creates a new SHA3-224 hash.
// Its generic security strength is 224 bits against preimage attacks,
// and 112 bits against collision attacks.
func New224() hash.Hash { return &state{rate: 144, outputLen: 28, dsbyte: 0x06} }

// New256 creates a new SHA3-256 hash.
// Its generic security strength is 256 bits against preimage attacks,
// and 128 bits against collision attacks.
func New256() hash.Hash { return &state{rate: 136, outputLen: 32, dsbyte: 0x06} }

// New384 creates a new SHA3-384 hash.
// Its generic security strength is 384 bits against preimage attacks,
// and 192 bits against collision attacks.
func New384() hash.Hash { return &state{rate: 104, outputLen: 48, dsbyte: 0x06} }

// New512 creates a new SHA3-512 hash.
// Its generic security strength is 512 bits against preimage attacks,
// and 256 bits against collision attacks.
func New512() hash.Hash { return &state{rate: 72, outputLen: 64, dsbyte: 0x06} }

// Sum224 returns the SHA3-224 digest of the data.
func Sum224(data []byte) (digest [28]byte) {
	h := New224()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum256 returns the SHA3-256 digest of the data.
func Sum256(data []byte) (digest [32]byte) {
	h := New256()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum384 returns the SHA3-384 digest of the data.
func Sum384(data []byte) (digest [48]byte) {
	h := New384()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum512 returns the SHA3-512 digest of the data.
func Sum512(data []byte) (digest [64]byte) {
	h := New512()
	h.Write(data)
	h.Sum(digest[:0])
	return
}
<commit_msg>crypto/sha3: Export Keccak-256 hash function<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sha3

// This file provides functions for creating instances of the SHA-3
// and SHAKE hash functions, as well as utility functions for hashing
// bytes.

import (
	"hash"
)

// NewKeccak256 creates a new Keccak-256 hash.
func NewKeccak256() hash.Hash { return &state{rate: 136, outputLen: 32, dsbyte: 0x01} }

// New224 creates a new SHA3-224 hash.
// Its generic security strength is 224 bits against preimage attacks,
// and 112 bits against collision attacks.
func New224() hash.Hash { return &state{rate: 144, outputLen: 28, dsbyte: 0x06} }

// New256 creates a new SHA3-256 hash.
// Its generic security strength is 256 bits against preimage attacks,
// and 128 bits against collision attacks.
func New256() hash.Hash { return &state{rate: 136, outputLen: 32, dsbyte: 0x06} }

// New384 creates a new SHA3-384 hash.
// Its generic security strength is 384 bits against preimage attacks,
// and 192 bits against collision attacks.
func New384() hash.Hash { return &state{rate: 104, outputLen: 48, dsbyte: 0x06} }

// New512 creates a new SHA3-512 hash.
// Its generic security strength is 512 bits against preimage attacks,
// and 256 bits against collision attacks.
func New512() hash.Hash { return &state{rate: 72, outputLen: 64, dsbyte: 0x06} }

// Sum224 returns the SHA3-224 digest of the data.
func Sum224(data []byte) (digest [28]byte) {
	h := New224()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum256 returns the SHA3-256 digest of the data.
func Sum256(data []byte) (digest [32]byte) {
	h := New256()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum384 returns the SHA3-384 digest of the data.
func Sum384(data []byte) (digest [48]byte) {
	h := New384()
	h.Write(data)
	h.Sum(digest[:0])
	return
}

// Sum512 returns the SHA3-512 digest of the data.
func Sum512(data []byte) (digest [64]byte) {
	h := New512()
	h.Write(data)
	h.Sum(digest[:0])
	return
}
<|endoftext|>