<commit_before>// +build linux

/*
 * Copyright (C) 2018 Red Hat, Inc.
 *
 * 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 ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package runc

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/vishvananda/netns"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology"
	tp "github.com/skydive-project/skydive/topology/probes"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
)

// ProbeHandler describes a Docker topology graph that enhance the graph
type ProbeHandler struct {
	common.RWMutex
	*ns.ProbeHandler
	state          common.ServiceState
	hostNs         netns.NsHandle
	wg             sync.WaitGroup
	watcher        *fsnotify.Watcher
	paths          []string
	containers     map[string]*container
	containersLock common.RWMutex
}

type container struct {
	namespace string
	node      *graph.Node
}

func (p *ProbeHandler) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

type initProcessStart int64

// mount contains source (host machine) to destination (guest machine) mappings
type mount struct {
	Source      string `json:"source"`
	Destination string `json:"destination"`
}

type containerState struct {
	path             string
	ID               string           `json:"id"`
	InitProcessPid   int              `json:"init_process_pid"`
	InitProcessStart initProcessStart `json:"init_process_start"`
	Config           struct {
		Labels []string `json:"labels"`
		Mounts []mount  `json:"mounts"`
	} `json:"config"`
}

// UnmarshalJSON custom marshall to handle both string and int64
func (ips *initProcessStart) UnmarshalJSON(data []byte) error {
	var v interface{}
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	i, err := common.ToInt64(v)
	if err != nil {
		return err
	}

	*ips = initProcessStart(i)

	return nil
}

func (p *ProbeHandler) getLabels(raw []string) graph.Metadata {
	labels := graph.Metadata{}

	for _, label := range raw {
		kv := strings.SplitN(label, "=", 2)
		switch len(kv) {
		case 1:
			p.Ctx.Logger.Warningf("Label format should be key=value: %s", label)
		case 2:
			var value interface{}
			if err := json.Unmarshal([]byte(kv[1]), &value); err != nil {
				labels.SetFieldAndNormalize(kv[0], kv[1])
			} else {
				labels.SetFieldAndNormalize(kv[0], value)
			}
		}
	}

	return labels
}

func getCreateConfig(path string) (*CreateConfig, error) {
	body, err := ioutil.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("Unable to read create config %s: %s", path, err)
	}

	var cc CreateConfig
	if err := json.Unmarshal(body, &cc); err != nil {
		return nil, fmt.Errorf("Unable to parse create config %s: %s", path, err)
	}

	return &cc, nil
}

func getStatus(state *containerState) string {
	if state.InitProcessPid == 0 {
		return "stopped"
	}

	info, err := common.GetProcessInfo(state.InitProcessPid)
	if err != nil {
		return "stopped"
	}
	if info.Start != int64(state.InitProcessStart) || info.State == common.Zombie || info.State == common.Dead {
		return "stopped"
	}
	base := path.Base(state.path)
	if _, err := os.Stat(path.Join(base, "exec.fifo")); err == nil {
		return "created"
	}
	return "running"
}

func parseState(path string) (*containerState, error) {
	body, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("Unable to read container state %s: %s", path, err)
	}

	var state containerState
	if err := json.Unmarshal(body, &state); err != nil {
		return nil, fmt.Errorf("Unable to parse container state %s: %s", path, err)
	}

	state.path = path

	return &state, nil
}

func getHostsFromState(state *containerState) (string, error) {
	const path = "/etc/hosts"
	for _, mount := range state.Config.Mounts {
		if mount.Destination == path {
			return mount.Source, nil
		}
	}
	return "", fmt.Errorf("Unable to find binding of %s", path)
}

func (p *ProbeHandler) parseHosts(state *containerState) *Hosts {
	path, err := getHostsFromState(state)
	if err != nil {
		return nil
	}

	hosts, err := readHosts(path)
	if err != nil {
		if os.IsNotExist(err) {
			p.Ctx.Logger.Debug(err)
		} else {
			p.Ctx.Logger.Error(err)
		}

		return nil
	}

	return hosts
}

func (p *ProbeHandler) getMetadata(state *containerState) Metadata {
	m := Metadata{
		ContainerID: state.ID,
		Status:      getStatus(state),
	}

	if labels := p.getLabels(state.Config.Labels); len(labels) > 0 {
		m.Labels = labels

		if b, ok := labels["bundle"]; ok {
			cc, err := getCreateConfig(b.(string) + "/artifacts/create-config")
			if err != nil {
				p.Ctx.Logger.Error(err)
			} else {
				m.CreateConfig = cc
			}
		}
	}

	if hosts := p.parseHosts(state); hosts != nil {
		m.Hosts = hosts
	}

	return m
}

func (p *ProbeHandler) updateContainer(path string, cnt *container) error {
	state, err := parseState(path)
	if err != nil {
		return err
	}

	p.Ctx.Graph.Lock()
	p.Ctx.Graph.AddMetadata(cnt.node, "Runc", p.getMetadata(state))
	p.Ctx.Graph.Unlock()

	return nil
}

func (p *ProbeHandler) registerContainer(path string) error {
	state, err := parseState(path)
	if err != nil {
		return err
	}

	p.Lock()
	defer p.Unlock()

	nsHandle, err := netns.GetFromPid(state.InitProcessPid)
	if err != nil {
		return err
	}
	defer nsHandle.Close()

	namespace := p.containerNamespace(state.InitProcessPid)
	p.Ctx.Logger.Debugf("Register runc container %s and PID %d", state.ID, state.InitProcessPid)

	var n *graph.Node
	if p.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = p.Ctx.RootNode
	} else {
		if n, err = p.Register(namespace, state.ID); err != nil {
			return err
		}

		p.Ctx.Graph.Lock()
		tr := p.Ctx.Graph.StartMetadataTransaction(n)
		tr.AddMetadata("Runtime", "runc")

		if _, err := n.GetFieldString("Manager"); err != nil {
			tr.AddMetadata("Manager", "runc")
		}
		if err := tr.Commit(); err != nil {
			p.Ctx.Logger.Error(err)
		}
		p.Ctx.Graph.Unlock()
	}

	pid := int64(state.InitProcessPid)
	runcMetadata := p.getMetadata(state)

	p.Ctx.Graph.Lock()
	defer p.Ctx.Graph.Unlock()

	containerNode := p.Ctx.Graph.LookupFirstNode(graph.Metadata{"InitProcessPID": pid})
	if containerNode != nil {
		tr := p.Ctx.Graph.StartMetadataTransaction(containerNode)
		tr.AddMetadata("Runc", runcMetadata)
		tr.AddMetadata("Runtime", "runc")
		if err := tr.Commit(); err != nil {
			p.Ctx.Logger.Error(err)
		}
	} else {
		metadata := graph.Metadata{
			"Type":           "container",
			"Name":           state.ID,
			"Manager":        "runc",
			"Runtime":        "runc",
			"InitProcessPID": pid,
			"Runc":           runcMetadata,
		}

		if containerNode, err = p.Ctx.Graph.NewNode(graph.GenID(), metadata); err != nil {
			return err
		}
	}
	topology.AddOwnershipLink(p.Ctx.Graph, n, containerNode, nil)

	p.containersLock.Lock()
	p.containers[path] = &container{
		namespace: namespace,
		node:      containerNode,
	}
	p.containersLock.Unlock()

	return nil
}

func (p *ProbeHandler) initializeFolder(path string) {
	p.watcher.Add(path)

	files, _ := ioutil.ReadDir(path)
	for _, f := range files {
		subpath := path + "/" + f.Name()

		if p.isFolder(subpath) {
			p.initializeFolder(subpath)
		} else if p.isStatePath(subpath) {
			p.containersLock.RLock()
			c, ok := p.containers[subpath]
			p.containersLock.RUnlock()

			var err error
			if ok {
				err = p.updateContainer(subpath, c)
			} else {
				err = p.registerContainer(subpath)
			}
			if err != nil {
				p.Ctx.Logger.Error(err)
			}
		}
	}
}

func (p *ProbeHandler) initialize(path string) {
	defer p.wg.Done()

	common.Retry(func() error {
		if p.state.Load() != common.RunningState {
			return nil
		}

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

		p.initializeFolder(path)

		p.Ctx.Logger.Debugf("Probe initialized for %s", path)

		return nil
	}, math.MaxInt32, time.Second)
}

func (p *ProbeHandler) isStatePath(path string) bool {
	return strings.HasSuffix(path, "/state.json")
}

func (p *ProbeHandler) isFolder(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}

// Start the probe
func (p *ProbeHandler) start() {
	defer p.wg.Done()

	var err error
	if p.hostNs, err = netns.Get(); err != nil {
		p.Ctx.Logger.Errorf("Unable to get host namespace: %s", err)
		return
	}

	if !p.state.CompareAndSwap(common.StoppedState, common.RunningState) {
		return
	}

	for _, path := range p.paths {
		p.wg.Add(1)
		go p.initialize(path)
	}

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for p.state.Load() == common.RunningState {
		select {
		case ev := <-p.watcher.Events:
			if ev.Op&fsnotify.Create == fsnotify.Create {
				if p.isFolder(ev.Name) {
					p.watcher.Add(ev.Name)
				}
			}
			if ev.Op&fsnotify.Write == fsnotify.Write {
				if p.isStatePath(ev.Name) {
					p.containersLock.RLock()
					cnt, ok := p.containers[ev.Name]
					p.containersLock.RUnlock()

					var err error
					if ok {
						err = p.updateContainer(ev.Name, cnt)
					} else {
						err = p.registerContainer(ev.Name)
					}
					if err != nil {
						p.Ctx.Logger.Error(err)
					}
				}
			}
			if ev.Op&fsnotify.Remove == fsnotify.Remove {
				if p.isStatePath(ev.Name) {
					p.containersLock.RLock()
					cnt, ok := p.containers[ev.Name]
					p.containersLock.RUnlock()

					if ok {
						p.Unregister(cnt.namespace)

						p.containersLock.Lock()
						delete(p.containers, ev.Name)
						p.containersLock.Unlock()
					}
				} else {
					p.containersLock.Lock()
					for path, cnt := range p.containers {
						if strings.HasPrefix(path, ev.Name) {
							p.Unregister(cnt.namespace)

							delete(p.containers, path)
						}
					}
					p.containersLock.Unlock()
				}
			}
		case err := <-p.watcher.Errors:
			p.Ctx.Logger.Errorf("Error while watching runc state folder: %s", err)
		case <-ticker.C:
		}
	}
}

// Start the probe
func (p *ProbeHandler) Start() {
	p.wg.Add(1)
	go p.start()
}

// Stop the probe
func (p *ProbeHandler) Stop() {
	if !p.state.CompareAndSwap(common.RunningState, common.StoppingState) {
		return
	}
	p.wg.Wait()

	p.hostNs.Close()

	p.state.Store(common.StoppedState)
}

// Init initializes a new topology runc probe
func (p *ProbeHandler) Init(ctx tp.Context, bundle *probe.Bundle) (probe.Handler, error) {
	nsHandler := bundle.GetHandler("netns")
	if nsHandler == nil {
		return nil, errors.New("unable to find the netns handler")
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("Unable to create a new Watcher: %s", err)
	}

	paths := ctx.Config.GetStringSlice("agent.topology.runc.run_path")

	p.ProbeHandler = nsHandler.(*ns.ProbeHandler)
	p.state = common.StoppedState
	p.watcher = watcher
	p.paths = paths
	p.containers = make(map[string]*container)

	return p, nil
}
<commit_msg>runc: improve a bit error message when getting ns handler<commit_after>// +build linux

/*
 * Copyright (C) 2018 Red Hat, Inc.
 *
 * 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 ofthe 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 specificlanguage governing permissions and
 * limitations under the License.
 *
 */

package runc

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/vishvananda/netns"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/graffiti/graph"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology"
	tp "github.com/skydive-project/skydive/topology/probes"
	ns "github.com/skydive-project/skydive/topology/probes/netns"
)

// ProbeHandler describes a Docker topology graph that enhance the graph
type ProbeHandler struct {
	common.RWMutex
	*ns.ProbeHandler
	state          common.ServiceState
	hostNs         netns.NsHandle
	wg             sync.WaitGroup
	watcher        *fsnotify.Watcher
	paths          []string
	containers     map[string]*container
	containersLock common.RWMutex
}

type container struct {
	namespace string
	node      *graph.Node
}

func (p *ProbeHandler) containerNamespace(pid int) string {
	return fmt.Sprintf("/proc/%d/ns/net", pid)
}

type initProcessStart int64

// mount contains source (host machine) to destination (guest machine) mappings
type mount struct {
	Source      string `json:"source"`
	Destination string `json:"destination"`
}

type containerState struct {
	path             string
	ID               string           `json:"id"`
	InitProcessPid   int              `json:"init_process_pid"`
	InitProcessStart initProcessStart `json:"init_process_start"`
	Config           struct {
		Labels []string `json:"labels"`
		Mounts []mount  `json:"mounts"`
	} `json:"config"`
}

// UnmarshalJSON custom marshall to handle both string and int64
func (ips *initProcessStart) UnmarshalJSON(data []byte) error {
	var v interface{}
	if err := json.Unmarshal(data, &v); err != nil {
		return err
	}

	i, err := common.ToInt64(v)
	if err != nil {
		return err
	}

	*ips = initProcessStart(i)

	return nil
}

func (p *ProbeHandler) getLabels(raw []string) graph.Metadata {
	labels := graph.Metadata{}

	for _, label := range raw {
		kv := strings.SplitN(label, "=", 2)
		switch len(kv) {
		case 1:
			p.Ctx.Logger.Warningf("Label format should be key=value: %s", label)
		case 2:
			var value interface{}
			if err := json.Unmarshal([]byte(kv[1]), &value); err != nil {
				labels.SetFieldAndNormalize(kv[0], kv[1])
			} else {
				labels.SetFieldAndNormalize(kv[0], value)
			}
		}
	}

	return labels
}

func getCreateConfig(path string) (*CreateConfig, error) {
	body, err := ioutil.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("Unable to read create config %s: %s", path, err)
	}

	var cc CreateConfig
	if err := json.Unmarshal(body, &cc); err != nil {
		return nil, fmt.Errorf("Unable to parse create config %s: %s", path, err)
	}

	return &cc, nil
}

func getStatus(state *containerState) string {
	if state.InitProcessPid == 0 {
		return "stopped"
	}

	info, err := common.GetProcessInfo(state.InitProcessPid)
	if err != nil {
		return "stopped"
	}
	if info.Start != int64(state.InitProcessStart) || info.State == common.Zombie || info.State == common.Dead {
		return "stopped"
	}
	base := path.Base(state.path)
	if _, err := os.Stat(path.Join(base, "exec.fifo")); err == nil {
		return "created"
	}
	return "running"
}

func parseState(path string) (*containerState, error) {
	body, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("Unable to read container state %s: %s", path, err)
	}

	var state containerState
	if err := json.Unmarshal(body, &state); err != nil {
		return nil, fmt.Errorf("Unable to parse container state %s: %s", path, err)
	}

	state.path = path

	return &state, nil
}

func getHostsFromState(state *containerState) (string, error) {
	const path = "/etc/hosts"
	for _, mount := range state.Config.Mounts {
		if mount.Destination == path {
			return mount.Source, nil
		}
	}
	return "", fmt.Errorf("Unable to find binding of %s", path)
}

func (p *ProbeHandler) parseHosts(state *containerState) *Hosts {
	path, err := getHostsFromState(state)
	if err != nil {
		return nil
	}

	hosts, err := readHosts(path)
	if err != nil {
		if os.IsNotExist(err) {
			p.Ctx.Logger.Debug(err)
		} else {
			p.Ctx.Logger.Error(err)
		}

		return nil
	}

	return hosts
}

func (p *ProbeHandler) getMetadata(state *containerState) Metadata {
	m := Metadata{
		ContainerID: state.ID,
		Status:      getStatus(state),
	}

	if labels := p.getLabels(state.Config.Labels); len(labels) > 0 {
		m.Labels = labels

		if b, ok := labels["bundle"]; ok {
			cc, err := getCreateConfig(b.(string) + "/artifacts/create-config")
			if err != nil {
				p.Ctx.Logger.Error(err)
			} else {
				m.CreateConfig = cc
			}
		}
	}

	if hosts := p.parseHosts(state); hosts != nil {
		m.Hosts = hosts
	}

	return m
}

func (p *ProbeHandler) updateContainer(path string, cnt *container) error {
	state, err := parseState(path)
	if err != nil {
		return err
	}

	p.Ctx.Graph.Lock()
	p.Ctx.Graph.AddMetadata(cnt.node, "Runc", p.getMetadata(state))
	p.Ctx.Graph.Unlock()

	return nil
}

func (p *ProbeHandler) registerContainer(path string) error {
	state, err := parseState(path)
	if err != nil {
		return err
	}

	p.Lock()
	defer p.Unlock()

	nsHandle, err := netns.GetFromPid(state.InitProcessPid)
	if err != nil {
		return fmt.Errorf("Unable to open netns, pid %d, state file %s : %s", state.InitProcessPid, path, err)
	}
	defer nsHandle.Close()

	namespace := p.containerNamespace(state.InitProcessPid)
	p.Ctx.Logger.Debugf("Register runc container %s and PID %d", state.ID, state.InitProcessPid)

	var n *graph.Node
	if p.hostNs.Equal(nsHandle) {
		// The container is in net=host mode
		n = p.Ctx.RootNode
	} else {
		if n, err = p.Register(namespace, state.ID); err != nil {
			return err
		}

		p.Ctx.Graph.Lock()
		tr := p.Ctx.Graph.StartMetadataTransaction(n)
		tr.AddMetadata("Runtime", "runc")

		if _, err := n.GetFieldString("Manager"); err != nil {
			tr.AddMetadata("Manager", "runc")
		}
		if err := tr.Commit(); err != nil {
			p.Ctx.Logger.Error(err)
		}
		p.Ctx.Graph.Unlock()
	}

	pid := int64(state.InitProcessPid)
	runcMetadata := p.getMetadata(state)

	p.Ctx.Graph.Lock()
	defer p.Ctx.Graph.Unlock()

	containerNode := p.Ctx.Graph.LookupFirstNode(graph.Metadata{"InitProcessPID": pid})
	if containerNode != nil {
		tr := p.Ctx.Graph.StartMetadataTransaction(containerNode)
		tr.AddMetadata("Runc", runcMetadata)
		tr.AddMetadata("Runtime", "runc")
		if err := tr.Commit(); err != nil {
			p.Ctx.Logger.Error(err)
		}
	} else {
		metadata := graph.Metadata{
			"Type":           "container",
			"Name":           state.ID,
			"Manager":        "runc",
			"Runtime":        "runc",
			"InitProcessPID": pid,
			"Runc":           runcMetadata,
		}

		if containerNode, err = p.Ctx.Graph.NewNode(graph.GenID(), metadata); err != nil {
			return err
		}
	}
	topology.AddOwnershipLink(p.Ctx.Graph, n, containerNode, nil)

	p.containersLock.Lock()
	p.containers[path] = &container{
		namespace: namespace,
		node:      containerNode,
	}
	p.containersLock.Unlock()

	return nil
}

func (p *ProbeHandler) initializeFolder(path string) {
	p.watcher.Add(path)

	files, _ := ioutil.ReadDir(path)
	for _, f := range files {
		subpath := path + "/" + f.Name()

		if p.isFolder(subpath) {
			p.initializeFolder(subpath)
		} else if p.isStatePath(subpath) {
			p.containersLock.RLock()
			c, ok := p.containers[subpath]
			p.containersLock.RUnlock()

			var err error
			if ok {
				err = p.updateContainer(subpath, c)
			} else {
				err = p.registerContainer(subpath)
			}
			if err != nil {
				p.Ctx.Logger.Error(err)
			}
		}
	}
}

func (p *ProbeHandler) initialize(path string) {
	defer p.wg.Done()

	common.Retry(func() error {
		if p.state.Load() != common.RunningState {
			return nil
		}

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

		p.initializeFolder(path)

		p.Ctx.Logger.Debugf("Probe initialized for %s", path)

		return nil
	}, math.MaxInt32, time.Second)
}

func (p *ProbeHandler) isStatePath(path string) bool {
	return strings.HasSuffix(path, "/state.json")
}

func (p *ProbeHandler) isFolder(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}

// Start the probe
func (p *ProbeHandler) start() {
	defer p.wg.Done()

	var err error
	if p.hostNs, err = netns.Get(); err != nil {
		p.Ctx.Logger.Errorf("Unable to get host namespace: %s", err)
		return
	}

	if !p.state.CompareAndSwap(common.StoppedState, common.RunningState) {
		return
	}

	for _, path := range p.paths {
		p.wg.Add(1)
		go p.initialize(path)
	}

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for p.state.Load() == common.RunningState {
		select {
		case ev := <-p.watcher.Events:
			if ev.Op&fsnotify.Create == fsnotify.Create {
				if p.isFolder(ev.Name) {
					p.watcher.Add(ev.Name)
				}
			}
			if ev.Op&fsnotify.Write == fsnotify.Write {
				if p.isStatePath(ev.Name) {
					p.containersLock.RLock()
					cnt, ok := p.containers[ev.Name]
					p.containersLock.RUnlock()

					var err error
					if ok {
						err = p.updateContainer(ev.Name, cnt)
					} else {
						err = p.registerContainer(ev.Name)
					}
					if err != nil {
						p.Ctx.Logger.Error(err)
					}
				}
			}
			if ev.Op&fsnotify.Remove == fsnotify.Remove {
				if p.isStatePath(ev.Name) {
					p.containersLock.RLock()
					cnt, ok := p.containers[ev.Name]
					p.containersLock.RUnlock()

					if ok {
						p.Unregister(cnt.namespace)

						p.containersLock.Lock()
						delete(p.containers, ev.Name)
						p.containersLock.Unlock()
					}
				} else {
					p.containersLock.Lock()
					for path, cnt := range p.containers {
						if strings.HasPrefix(path, ev.Name) {
							p.Unregister(cnt.namespace)

							delete(p.containers, path)
						}
					}
					p.containersLock.Unlock()
				}
			}
		case err := <-p.watcher.Errors:
			p.Ctx.Logger.Errorf("Error while watching runc state folder: %s", err)
		case <-ticker.C:
		}
	}
}

// Start the probe
func (p *ProbeHandler) Start() {
	p.wg.Add(1)
	go p.start()
}

// Stop the probe
func (p *ProbeHandler) Stop() {
	if !p.state.CompareAndSwap(common.RunningState, common.StoppingState) {
		return
	}
	p.wg.Wait()

	p.hostNs.Close()

	p.state.Store(common.StoppedState)
}

// Init initializes a new topology runc probe
func (p *ProbeHandler) Init(ctx tp.Context, bundle *probe.Bundle) (probe.Handler, error) {
	nsHandler := bundle.GetHandler("netns")
	if nsHandler == nil {
		return nil, errors.New("unable to find the netns handler")
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, fmt.Errorf("Unable to create a new Watcher: %s", err)
	}

	paths := ctx.Config.GetStringSlice("agent.topology.runc.run_path")

	p.ProbeHandler = nsHandler.(*ns.ProbeHandler)
	p.state = common.StoppedState
	p.watcher = watcher
	p.paths = paths
	p.containers = make(map[string]*container)

	return p, nil
}
<|endoftext|>


<commit_before>package trustmanager

import (
	"path/filepath"
	"strings"
	"sync"

	"errors"
	"fmt"

	"github.com/docker/notary/pkg/passphrase"
	"github.com/endophage/gotuf/data"
)

const (
	keyExtension = "key"
)

// ErrPasswordInvalid is returned when signing fails. It could also mean the signing
// key file was corrupted, but we have no way to distinguish.
type ErrPasswordInvalid struct{}

// ErrPasswordInvalid is returned when signing fails. It could also mean the signing
// key file was corrupted, but we have no way to distinguish.
func (err ErrPasswordInvalid) Error() string {
	return "password invalid, operation has failed."
}

// KeyStore is a generic interface for private key storage
type KeyStore interface {
	LimitedFileStore

	AddKey(name, alias string, privKey data.PrivateKey) error
	GetKey(name string) (data.PrivateKey, string, error)
	ListKeys() []string
	RemoveKey(name string) error
}

type cachedKey struct {
	alias string
	key   data.PrivateKey
}

// PassphraseRetriever is a callback function that should retrieve a passphrase
// for a given named key. If it should be treated as new passphrase (e.g. with
// confirmation), createNew will be true. Attempts is passed in so that implementers
// decide how many chances to give to a human, for example.
type PassphraseRetriever func(keyId, alias string, createNew bool, attempts int) (passphrase string, giveup bool, err error)

// KeyFileStore persists and manages private keys on disk
type KeyFileStore struct {
	sync.Mutex
	SimpleFileStore
	passphrase.Retriever
	cachedKeys map[string]*cachedKey
}

// KeyMemoryStore manages private keys in memory
type KeyMemoryStore struct {
	sync.Mutex
	MemoryFileStore
	passphrase.Retriever
	cachedKeys map[string]*cachedKey
}

// NewKeyFileStore returns a new KeyFileStore creating a private directory to
// hold the keys.
func NewKeyFileStore(baseDir string, passphraseRetriever passphrase.Retriever) (*KeyFileStore, error) {
	fileStore, err := NewPrivateSimpleFileStore(baseDir, keyExtension)
	if err != nil {
		return nil, err
	}
	cachedKeys := make(map[string]*cachedKey)

	return &KeyFileStore{SimpleFileStore: *fileStore,
		Retriever:  passphraseRetriever,
		cachedKeys: cachedKeys}, nil
}

// AddKey stores the contents of a PEM-encoded private key as a PEM block
func (s *KeyFileStore) AddKey(name, alias string, privKey data.PrivateKey) error {
	s.Lock()
	defer s.Unlock()
	return addKey(s, s.Retriever, s.cachedKeys, name, alias, privKey)
}

// GetKey returns the PrivateKey given a KeyID
func (s *KeyFileStore) GetKey(name string) (data.PrivateKey, string, error) {
	s.Lock()
	defer s.Unlock()
	return getKey(s, s.Retriever, s.cachedKeys, name)
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func (s *KeyFileStore) ListKeys() []string {
	return listKeys(s)
}

// RemoveKey removes the key from the keyfilestore
func (s *KeyFileStore) RemoveKey(name string) error {
	s.Lock()
	defer s.Unlock()
	return removeKey(s, s.cachedKeys, name)
}

// NewKeyMemoryStore returns a new KeyMemoryStore which holds keys in memory
func NewKeyMemoryStore(passphraseRetriever passphrase.Retriever) *KeyMemoryStore {
	memStore := NewMemoryFileStore()
	cachedKeys := make(map[string]*cachedKey)

	return &KeyMemoryStore{MemoryFileStore: *memStore,
		Retriever:  passphraseRetriever,
		cachedKeys: cachedKeys}
}

// AddKey stores the contents of a PEM-encoded private key as a PEM block
func (s *KeyMemoryStore) AddKey(name, alias string, privKey data.PrivateKey) error {
	s.Lock()
	defer s.Unlock()
	return addKey(s, s.Retriever, s.cachedKeys, name, alias, privKey)
}

// GetKey returns the PrivateKey given a KeyID
func (s *KeyMemoryStore) GetKey(name string) (data.PrivateKey, string, error) {
	s.Lock()
	defer s.Unlock()
	return getKey(s, s.Retriever, s.cachedKeys, name)
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func (s *KeyMemoryStore) ListKeys() []string {
	return listKeys(s)
}

// RemoveKey removes the key from the keystore
func (s *KeyMemoryStore) RemoveKey(name string) error {
	s.Lock()
	defer s.Unlock()
	return removeKey(s, s.cachedKeys, name)
}

func addKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cachedKeys map[string]*cachedKey, name, alias string, privKey data.PrivateKey) error {
	pemPrivKey, err := KeyToPEM(privKey)
	if err != nil {
		return err
	}

	attempts := 0
	passphrase := ""
	giveup := false
	for {
		passphrase, giveup, err = passphraseRetriever(name, alias, true, attempts)
		if err != nil {
			attempts++
			continue
		}
		if giveup {
			return ErrPasswordInvalid{}
		}
		if attempts > 10 {
			return errors.New("maximum number of passphrase attempts exceeded")
		}
		break
	}

	if passphrase != "" {
		pemPrivKey, err = EncryptPrivateKey(privKey, passphrase)
		if err != nil {
			return err
		}
	}

	cachedKeys[name] = &cachedKey{alias: alias, key: privKey}
	return s.Add(name+"_"+alias, pemPrivKey)
}

func getKeyAlias(s LimitedFileStore, keyID string) (string, error) {
	files := s.ListFiles(true)
	name := strings.TrimSpace(strings.TrimSuffix(filepath.Base(keyID), filepath.Ext(keyID)))

	for _, file := range files {
		filename := filepath.Base(file)

		if strings.HasPrefix(filename, name) {
			aliasPlusDotKey := strings.TrimPrefix(filename, name+"_")
			retVal := strings.TrimSuffix(aliasPlusDotKey, "."+keyExtension)
			return retVal, nil
		}
	}

	return "", fmt.Errorf("keyId %s has no alias", name)
}

// GetKey returns the PrivateKey given a KeyID
func getKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cachedKeys map[string]*cachedKey, name string) (data.PrivateKey, string, error) {
	cachedKeyEntry, ok := cachedKeys[name]
	if ok {
		return cachedKeyEntry.key, cachedKeyEntry.alias, nil
	}
	keyAlias, err := getKeyAlias(s, name)
	if err != nil {
		return nil, "", err
	}

	keyBytes, err := s.Get(name + "_" + keyAlias)
	if err != nil {
		return nil, "", err
	}

	var retErr error
	// See if the key is encrypted. If its encrypted we'll fail to parse the private key
	privKey, err := ParsePEMPrivateKey(keyBytes, "")
	if err != nil {
		// We need to decrypt the key, lets get a passphrase
		for attempts := 0; ; attempts++ {
			passphrase, giveup, err := passphraseRetriever(name, string(keyAlias), false, attempts)
			// Check if the passphrase retriever got an error or if it is telling us to give up
			if giveup || err != nil {
				return nil, "", errors.New("obtaining passphrase failed")
			}
			if attempts > 10 {
				return nil, "", errors.New("maximum number of passphrase attempts exceeded")
			}

			// Try to convert PEM encoded bytes back to a PrivateKey using the passphrase
			privKey, err = ParsePEMPrivateKey(keyBytes, passphrase)
			if err != nil {
				retErr = ErrPasswordInvalid{}
			} else {
				// We managed to parse the PrivateKey. We've succeeded!
				retErr = nil
				break
			}
		}
	}
	if retErr != nil {
		return nil, "", retErr
	}
	cachedKeys[name] = &cachedKey{alias: keyAlias, key: privKey}
	return privKey, keyAlias, nil
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func listKeys(s LimitedFileStore) []string {
	var keyIDList []string

	for _, f := range s.ListFiles(false) {
		keyID := strings.TrimSpace(strings.TrimSuffix(f, filepath.Ext(f)))
		keyID = keyID[:strings.LastIndex(keyID, "_")]
		keyIDList = append(keyIDList, keyID)
	}
	return keyIDList
}

// RemoveKey removes the key from the keyfilestore
func removeKey(s LimitedFileStore, cachedKeys map[string]*cachedKey, name string) error {
	keyAlias, err := getKeyAlias(s, name)
	if err != nil {
		return err
	}

	delete(cachedKeys, name)

	return s.Remove(name + "_" + keyAlias)
}
<commit_msg>Add missing use of invalid passphrase error<commit_after>package trustmanager

import (
	"path/filepath"
	"strings"
	"sync"

	"errors"
	"fmt"

	"github.com/docker/notary/pkg/passphrase"
	"github.com/endophage/gotuf/data"
)

const (
	keyExtension = "key"
)

// ErrPasswordInvalid is returned when signing fails. It could also mean the signing
// key file was corrupted, but we have no way to distinguish.
type ErrPasswordInvalid struct{}

// ErrPasswordInvalid is returned when signing fails. It could also mean the signing
// key file was corrupted, but we have no way to distinguish.
func (err ErrPasswordInvalid) Error() string {
	return "password invalid, operation has failed."
}

// KeyStore is a generic interface for private key storage
type KeyStore interface {
	LimitedFileStore

	AddKey(name, alias string, privKey data.PrivateKey) error
	GetKey(name string) (data.PrivateKey, string, error)
	ListKeys() []string
	RemoveKey(name string) error
}

type cachedKey struct {
	alias string
	key   data.PrivateKey
}

// PassphraseRetriever is a callback function that should retrieve a passphrase
// for a given named key. If it should be treated as new passphrase (e.g. with
// confirmation), createNew will be true. Attempts is passed in so that implementers
// decide how many chances to give to a human, for example.
type PassphraseRetriever func(keyId, alias string, createNew bool, attempts int) (passphrase string, giveup bool, err error)

// KeyFileStore persists and manages private keys on disk
type KeyFileStore struct {
	sync.Mutex
	SimpleFileStore
	passphrase.Retriever
	cachedKeys map[string]*cachedKey
}

// KeyMemoryStore manages private keys in memory
type KeyMemoryStore struct {
	sync.Mutex
	MemoryFileStore
	passphrase.Retriever
	cachedKeys map[string]*cachedKey
}

// NewKeyFileStore returns a new KeyFileStore creating a private directory to
// hold the keys.
func NewKeyFileStore(baseDir string, passphraseRetriever passphrase.Retriever) (*KeyFileStore, error) {
	fileStore, err := NewPrivateSimpleFileStore(baseDir, keyExtension)
	if err != nil {
		return nil, err
	}
	cachedKeys := make(map[string]*cachedKey)

	return &KeyFileStore{SimpleFileStore: *fileStore,
		Retriever:  passphraseRetriever,
		cachedKeys: cachedKeys}, nil
}

// AddKey stores the contents of a PEM-encoded private key as a PEM block
func (s *KeyFileStore) AddKey(name, alias string, privKey data.PrivateKey) error {
	s.Lock()
	defer s.Unlock()
	return addKey(s, s.Retriever, s.cachedKeys, name, alias, privKey)
}

// GetKey returns the PrivateKey given a KeyID
func (s *KeyFileStore) GetKey(name string) (data.PrivateKey, string, error) {
	s.Lock()
	defer s.Unlock()
	return getKey(s, s.Retriever, s.cachedKeys, name)
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func (s *KeyFileStore) ListKeys() []string {
	return listKeys(s)
}

// RemoveKey removes the key from the keyfilestore
func (s *KeyFileStore) RemoveKey(name string) error {
	s.Lock()
	defer s.Unlock()
	return removeKey(s, s.cachedKeys, name)
}

// NewKeyMemoryStore returns a new KeyMemoryStore which holds keys in memory
func NewKeyMemoryStore(passphraseRetriever passphrase.Retriever) *KeyMemoryStore {
	memStore := NewMemoryFileStore()
	cachedKeys := make(map[string]*cachedKey)

	return &KeyMemoryStore{MemoryFileStore: *memStore,
		Retriever:  passphraseRetriever,
		cachedKeys: cachedKeys}
}

// AddKey stores the contents of a PEM-encoded private key as a PEM block
func (s *KeyMemoryStore) AddKey(name, alias string, privKey data.PrivateKey) error {
	s.Lock()
	defer s.Unlock()
	return addKey(s, s.Retriever, s.cachedKeys, name, alias, privKey)
}

// GetKey returns the PrivateKey given a KeyID
func (s *KeyMemoryStore) GetKey(name string) (data.PrivateKey, string, error) {
	s.Lock()
	defer s.Unlock()
	return getKey(s, s.Retriever, s.cachedKeys, name)
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func (s *KeyMemoryStore) ListKeys() []string {
	return listKeys(s)
}

// RemoveKey removes the key from the keystore
func (s *KeyMemoryStore) RemoveKey(name string) error {
	s.Lock()
	defer s.Unlock()
	return removeKey(s, s.cachedKeys, name)
}

func addKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cachedKeys map[string]*cachedKey, name, alias string, privKey data.PrivateKey) error {
	pemPrivKey, err := KeyToPEM(privKey)
	if err != nil {
		return err
	}

	attempts := 0
	passphrase := ""
	giveup := false
	for {
		passphrase, giveup, err = passphraseRetriever(name, alias, true, attempts)
		if err != nil {
			attempts++
			continue
		}
		if giveup {
			return ErrPasswordInvalid{}
		}
		if attempts > 10 {
			return errors.New("maximum number of passphrase attempts exceeded")
		}
		break
	}

	if passphrase != "" {
		pemPrivKey, err = EncryptPrivateKey(privKey, passphrase)
		if err != nil {
			return err
		}
	}

	cachedKeys[name] = &cachedKey{alias: alias, key: privKey}
	return s.Add(name+"_"+alias, pemPrivKey)
}

func getKeyAlias(s LimitedFileStore, keyID string) (string, error) {
	files := s.ListFiles(true)
	name := strings.TrimSpace(strings.TrimSuffix(filepath.Base(keyID), filepath.Ext(keyID)))

	for _, file := range files {
		filename := filepath.Base(file)

		if strings.HasPrefix(filename, name) {
			aliasPlusDotKey := strings.TrimPrefix(filename, name+"_")
			retVal := strings.TrimSuffix(aliasPlusDotKey, "."+keyExtension)
			return retVal, nil
		}
	}

	return "", fmt.Errorf("keyId %s has no alias", name)
}

// GetKey returns the PrivateKey given a KeyID
func getKey(s LimitedFileStore, passphraseRetriever passphrase.Retriever, cachedKeys map[string]*cachedKey, name string) (data.PrivateKey, string, error) {
	cachedKeyEntry, ok := cachedKeys[name]
	if ok {
		return cachedKeyEntry.key, cachedKeyEntry.alias, nil
	}
	keyAlias, err := getKeyAlias(s, name)
	if err != nil {
		return nil, "", err
	}

	keyBytes, err := s.Get(name + "_" + keyAlias)
	if err != nil {
		return nil, "", err
	}

	var retErr error
	// See if the key is encrypted. If its encrypted we'll fail to parse the private key
	privKey, err := ParsePEMPrivateKey(keyBytes, "")
	if err != nil {
		// We need to decrypt the key, lets get a passphrase
		for attempts := 0; ; attempts++ {
			passphrase, giveup, err := passphraseRetriever(name, string(keyAlias), false, attempts)
			// Check if the passphrase retriever got an error or if it is telling us to give up
			if giveup || err != nil {
				return nil, "", ErrPasswordInvalid{}
			}
			if attempts > 10 {
				return nil, "", errors.New("maximum number of passphrase attempts exceeded")
			}

			// Try to convert PEM encoded bytes back to a PrivateKey using the passphrase
			privKey, err = ParsePEMPrivateKey(keyBytes, passphrase)
			if err != nil {
				retErr = ErrPasswordInvalid{}
			} else {
				// We managed to parse the PrivateKey. We've succeeded!
				retErr = nil
				break
			}
		}
	}
	if retErr != nil {
		return nil, "", retErr
	}
	cachedKeys[name] = &cachedKey{alias: keyAlias, key: privKey}
	return privKey, keyAlias, nil
}

// ListKeys returns a list of unique PublicKeys present on the KeyFileStore.
// There might be symlinks associating Certificate IDs to Public Keys, so this
// method only returns the IDs that aren't symlinks
func listKeys(s LimitedFileStore) []string {
	var keyIDList []string

	for _, f := range s.ListFiles(false) {
		keyID := strings.TrimSpace(strings.TrimSuffix(f, filepath.Ext(f)))
		keyID = keyID[:strings.LastIndex(keyID, "_")]
		keyIDList = append(keyIDList, keyID)
	}
	return keyIDList
}

// RemoveKey removes the key from the keyfilestore
func removeKey(s LimitedFileStore, cachedKeys map[string]*cachedKey, name string) error {
	keyAlias, err := getKeyAlias(s, name)
	if err != nil {
		return err
	}

	delete(cachedKeys, name)

	return s.Remove(name + "_" + keyAlias)
}
<|endoftext|>


<commit_before>package proxyconfig

import (
	"errors"
	"fmt"
	"koding/tools/config"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type ProxyMessage struct {
	Action   string `json:"action"`
	Name     string `json:"name"`
	Key      string `json:"key"`
	Host     string `json:"host"`
	HostData string `json:"hostdata"`
	Uuid     string `json:"uuid"`
}

type ProxyResponse struct {
	Action string `json:"action"`
	Uuid   string `json:"uuid"`
}

type KeyData struct {
	Key          string
	Host         string
	HostData     string
	CurrentIndex int
}

func NewKeyData(key, host, hostdata string, currentindex int) *KeyData {
	return &KeyData{
		Key:          key,
		Host:         host,
		HostData:     hostdata,
		CurrentIndex: currentindex,
	}
}

type KeyRoutingTable struct {
	Keys map[string][]KeyData `json:"keys"`
}

func NewKeyRoutingTable() *KeyRoutingTable {
	return &KeyRoutingTable{
		Keys: make(map[string][]KeyData),
	}
}

type DomainRoutingTable struct {
	Domain map[string]string `json:"domain"`
}

func NewDomainRoutingTable() *DomainRoutingTable {
	return &DomainRoutingTable{
		Domain: make(map[string]string),
	}
}

type ProxyConfiguration struct {
	Session    *mgo.Session
	Collection *mgo.Collection
}

type Proxy struct {
	Services           map[string]KeyRoutingTable
	DomainRoutingTable DomainRoutingTable
	Uuid               string
}

func NewProxy(uuid string) *Proxy {
	return &Proxy{
		Services:           make(map[string]KeyRoutingTable),
		DomainRoutingTable: *NewDomainRoutingTable(),
		Uuid:               uuid,
	}
}

func Connect() (*ProxyConfiguration, error) {
	host := config.Current.Kontrold.Mongo.Host
	session, err := mgo.Dial(host)
	if err != nil {
		return nil, err
	}

	session.SetMode(mgo.Strong, true)

	col := session.DB("kontrol").C("proxies")

	pr := &ProxyConfiguration{
		Session:    session,
		Collection: col,
	}

	return pr, nil
}

// Base INSERT/CREATE crud action
func (p *ProxyConfiguration) Add(proxy Proxy) error {
	err := p.Collection.Insert(proxy)
	if err != nil {
		return err
	}
	return nil
}
func (p *ProxyConfiguration) AddProxy(uuid string) error {
	err := p.HasUuid(uuid)
	if err == nil {
		return fmt.Errorf("registering not possible uuid '%s' uuid exists.", uuid)
	}

	proxy := *NewProxy(uuid)
	err = p.Add(proxy)
	if err != nil {
		return err
	}
	return nil
}

// Base DELETE crud action
func (p *ProxyConfiguration) Delete(uuid string) error {
	err := p.Collection.Remove(bson.M{"uuid": uuid})
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) DeleteProxy(uuid string) error {
	err := p.HasUuid(uuid)
	if err != nil {
		return fmt.Errorf("deleting not possible '%s'", err)
	}

	err = p.Delete(uuid)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) DeleteKey(name, key, host, hostdata, uuid string) error {
	proxy, err := p.GetProxy(uuid)
	if err != nil {
		return fmt.Errorf("deleting key not possible '%s'", err)
	}

	_, ok := proxy.Services[name]
	if !ok {
		return errors.New("service name is wrong. deletin key is not possible")

	}

	keyRoutingTable := proxy.Services[name]
	delete(keyRoutingTable.Keys, key)

	proxy.Services[name] = keyRoutingTable
	err = p.UpdateProxy(proxy)
	if err != nil {
		return err
	}

	return nil
}

func (p *ProxyConfiguration) UpdateProxy(proxy Proxy) error {
	err := p.Collection.Update(bson.M{"uuid": proxy.Uuid}, proxy)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) AddKey(name, key, host, hostdata, uuid string) error {
	proxy, err := p.GetProxy(uuid)
	if err != nil {
		return fmt.Errorf("adding key not possible '%s'", err)
	}

	_, ok := proxy.Services[name]
	if !ok {
		proxy.Services[name] = *NewKeyRoutingTable()
	}
	keyRoutingTable := proxy.Services[name]

	if len(keyRoutingTable.Keys) == 0 { // empty routing table, add it
		keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))
		err = p.UpdateProxy(proxy)
		if err != nil {
			return err
		}
		return nil
	}

	_, ok = keyRoutingTable.Keys[key] // new key, add it
	if !ok {
		keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))
		err = p.UpdateProxy(proxy)
		if err != nil {
			return err
		}
		return nil
	}

	// check for existing hostnames, if exist abort
	for _, value := range keyRoutingTable.Keys[key] {
		if value.Host == host {
			return nil
		}
	}

	keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))

	proxy.Services[name] = keyRoutingTable
	err = p.UpdateProxy(proxy)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) HasUuid(uuid string) error {
	result := Proxy{}
	err := p.Collection.Find(bson.M{"uuid": uuid}).One(&result)
	if err != nil {
		return fmt.Errorf("no proxy with the uuid %s exist.", uuid)
	}
	return nil
}

func (p *ProxyConfiguration) GetProxy(uuid string) (Proxy, error) {
	result := Proxy{}
	err := p.Collection.Find(bson.M{"uuid": uuid}).One(&result)
	if err != nil {
		return result, fmt.Errorf("no proxy with the uuid %s exist.", uuid)
	}

	return result, nil

}
<commit_msg>Removing unused DomainRoutingTable<commit_after>package proxyconfig

import (
	"errors"
	"fmt"
	"koding/tools/config"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

type ProxyMessage struct {
	Action   string `json:"action"`
	Name     string `json:"name"`
	Key      string `json:"key"`
	Host     string `json:"host"`
	HostData string `json:"hostdata"`
	Uuid     string `json:"uuid"`
}

type ProxyResponse struct {
	Action string `json:"action"`
	Uuid   string `json:"uuid"`
}

type KeyData struct {
	Key          string
	Host         string
	HostData     string
	CurrentIndex int
}

func NewKeyData(key, host, hostdata string, currentindex int) *KeyData {
	return &KeyData{
		Key:          key,
		Host:         host,
		HostData:     hostdata,
		CurrentIndex: currentindex,
	}
}

type KeyRoutingTable struct {
	Keys map[string][]KeyData `json:"keys"`
}

func NewKeyRoutingTable() *KeyRoutingTable {
	return &KeyRoutingTable{
		Keys: make(map[string][]KeyData),
	}
}

type Proxy struct {
	Services map[string]KeyRoutingTable
	Uuid     string
}

func NewProxy(uuid string) *Proxy {
	return &Proxy{
		Services: make(map[string]KeyRoutingTable),
		Uuid:     uuid,
	}
}

type ProxyConfiguration struct {
	Session    *mgo.Session
	Collection *mgo.Collection
}

func Connect() (*ProxyConfiguration, error) {
	host := config.Current.Kontrold.Mongo.Host
	session, err := mgo.Dial(host)
	if err != nil {
		return nil, err
	}

	session.SetMode(mgo.Strong, true)

	col := session.DB("kontrol").C("proxies")

	pr := &ProxyConfiguration{
		Session:    session,
		Collection: col,
	}

	return pr, nil
}

// Base INSERT/CREATE crud action
func (p *ProxyConfiguration) Add(proxy Proxy) error {
	err := p.Collection.Insert(proxy)
	if err != nil {
		return err
	}
	return nil
}
func (p *ProxyConfiguration) AddProxy(uuid string) error {
	err := p.HasUuid(uuid)
	if err == nil {
		return fmt.Errorf("registering not possible uuid '%s' uuid exists.", uuid)
	}

	proxy := *NewProxy(uuid)
	err = p.Add(proxy)
	if err != nil {
		return err
	}
	return nil
}

// Base DELETE crud action
func (p *ProxyConfiguration) Delete(uuid string) error {
	err := p.Collection.Remove(bson.M{"uuid": uuid})
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) DeleteProxy(uuid string) error {
	err := p.HasUuid(uuid)
	if err != nil {
		return fmt.Errorf("deleting not possible '%s'", err)
	}

	err = p.Delete(uuid)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) DeleteKey(name, key, host, hostdata, uuid string) error {
	proxy, err := p.GetProxy(uuid)
	if err != nil {
		return fmt.Errorf("deleting key not possible '%s'", err)
	}

	_, ok := proxy.Services[name]
	if !ok {
		return errors.New("service name is wrong. deletin key is not possible")

	}

	keyRoutingTable := proxy.Services[name]
	delete(keyRoutingTable.Keys, key)

	proxy.Services[name] = keyRoutingTable
	err = p.UpdateProxy(proxy)
	if err != nil {
		return err
	}

	return nil
}

func (p *ProxyConfiguration) UpdateProxy(proxy Proxy) error {
	err := p.Collection.Update(bson.M{"uuid": proxy.Uuid}, proxy)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) AddKey(name, key, host, hostdata, uuid string) error {
	proxy, err := p.GetProxy(uuid)
	if err != nil {
		return fmt.Errorf("adding key not possible '%s'", err)
	}

	_, ok := proxy.Services[name]
	if !ok {
		proxy.Services[name] = *NewKeyRoutingTable()
	}
	keyRoutingTable := proxy.Services[name]

	if len(keyRoutingTable.Keys) == 0 { // empty routing table, add it
		keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))
		err = p.UpdateProxy(proxy)
		if err != nil {
			return err
		}
		return nil
	}

	_, ok = keyRoutingTable.Keys[key] // new key, add it
	if !ok {
		keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))
		err = p.UpdateProxy(proxy)
		if err != nil {
			return err
		}
		return nil
	}

	// check for existing hostnames, if exist abort
	for _, value := range keyRoutingTable.Keys[key] {
		if value.Host == host {
			return nil
		}
	}

	keyRoutingTable.Keys[key] = append(keyRoutingTable.Keys[key], *NewKeyData(key, host, hostdata, 0))

	proxy.Services[name] = keyRoutingTable
	err = p.UpdateProxy(proxy)
	if err != nil {
		return err
	}
	return nil
}

func (p *ProxyConfiguration) HasUuid(uuid string) error {
	result := Proxy{}
	err := p.Collection.Find(bson.M{"uuid": uuid}).One(&result)
	if err != nil {
		return fmt.Errorf("no proxy with the uuid %s exist.", uuid)
	}
	return nil
}

func (p *ProxyConfiguration) GetProxy(uuid string) (Proxy, error) {
	result := Proxy{}
	err := p.Collection.Find(bson.M{"uuid": uuid}).One(&result)
	if err != nil {
		return result, fmt.Errorf("no proxy with the uuid %s exist.", uuid)
	}

	return result, nil

}
<|endoftext|>


<commit_before>package openstack

import (
	"bytes"
	"fmt"
	"launchpad.net/goose/nova"
	"launchpad.net/goose/swift"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/trivial"
	"net/http"
	"os"
)

type VirtualFile struct {
	bytes.Reader
}

var _ http.File = (*VirtualFile)(nil)

func (f *VirtualFile) Close() error {
	return nil
}

func (f *VirtualFile) Readdir(count int) ([]os.FileInfo, error) {
	return nil, nil
}

func (f *VirtualFile) Stat() (os.FileInfo, error) {
	return nil, fmt.Errorf("Can't stat VirtualFile")
}

func init() {
	http.DefaultTransport.(*http.Transport).RegisterProtocol("file", http.NewFileTransport(http.Dir("testdata")))
}

var origMetadataHost = metadataHost

func UseTestMetadata(local bool) {
	if local {
		metadataHost = "file:"
	} else {
		metadataHost = origMetadataHost
	}
}

var origMetadataJSON = metadataJSON

func UseMetadataJSON(path string) {
	if path != "" {
		metadataJSON = path
	} else {
		metadataJSON = origMetadataJSON
	}
}

var originalShortAttempt = shortAttempt
var originalLongAttempt = longAttempt

// ShortTimeouts sets the timeouts to a short period as we
// know that the testing server doesn't get better with time,
// and this reduces the test time from 30s to 3s.
func ShortTimeouts(short bool) {
	if short {
		shortAttempt = trivial.AttemptStrategy{
			Total: 0.25e9,
			Delay: 0.01e9,
		}
		longAttempt = shortAttempt
	} else {
		shortAttempt = originalShortAttempt
		longAttempt = originalLongAttempt
	}
}

var ShortAttempt = &shortAttempt

func DeleteStorageContent(s environs.Storage) error {
	return s.(*storage).deleteAll()
}

// WritablePublicStorage returns a Storage instance which is authorised to write to the PublicStorage bucket.
// It is used by tests which need to upload files.
func WritablePublicStorage(e environs.Environ) environs.Storage {
	ecfg := e.(*environ).ecfg()
	authModeCfg := AuthMode(ecfg.authMode())
	writablePublicStorage := &storage{
		containerName: ecfg.publicBucket(),
		swift:         swift.New(e.(*environ).client(ecfg, authModeCfg)),
	}

	// Ensure the container exists.
	err := writablePublicStorage.makeContainer(ecfg.publicBucket(), swift.PublicRead)
	if err != nil {
		panic(fmt.Errorf("cannot create writable public container: %v", err))
	}
	return writablePublicStorage
}
func InstanceAddress(addresses map[string][]nova.IPAddress) (string, error) {
	return instanceAddress(addresses)
}

func FindInstanceSpec(e environs.Environ, series, arch, flavor string) (imageId, flavorId string, err error) {
	env := e.(*environ)
	spec, err := findInstanceSpec(env, &instanceConstraint{
		series: series,
		arch:   arch,
		region: env.ecfg().region(),
		flavor: flavor,
	})
	if err == nil {
		imageId = spec.imageId
		flavorId = spec.flavorId
	}
	return
}

func SetUseFloatingIP(e environs.Environ, val bool) {
	env := e.(*environ)
	env.ecfg().attrs["use-floating-ip"] = val
}

func DefaultInstanceType(e environs.Environ) string {
	ecfg := e.(*environ).ecfg()
	return ecfg.defaultInstanceType()
}

// ImageDetails specify parameters used to start a test machine for the live tests.
type ImageDetails struct {
	Flavor  string
	ImageId string
}

type BootstrapState struct {
	StateInstances []state.InstanceId
}

func LoadState(e environs.Environ) (*BootstrapState, error) {
	s, err := e.(*environ).loadState()
	if err != nil {
		return nil, err
	}
	return &BootstrapState{s.StateInstances}, nil
}
<commit_msg>Move the code out of export test, start prepping for using internal data.<commit_after>package openstack

import (
	"fmt"
	"launchpad.net/goose/nova"
	"launchpad.net/goose/swift"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/jujutest"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/trivial"
	"net/http"
)

func init() {
	http.DefaultTransport.(*http.Transport).RegisterProtocol("file", http.NewFileTransport(http.Dir("testdata")))
}

var origMetadataHost = metadataHost

var metadataContent = `{"uuid": "d8e02d56-2648-49a3-bf97-6be8f1204f38",` +
	`"availability_zone": "nova", "hostname": "test.novalocal", ` +
	`"launch_index": 0, "meta": {"priority": "low", "role": "webserver"}, ` +
	`"public_keys": {"mykey": "ssh-rsa fake-key\n"}, "name": "test"}`

var metadataTestingBase = []jujutest.FileContent{
	{"latest/meta-data/instance-id", "i-000abc"},
	{"latest/meta-data/local-ipv4", "203.1.1.2"},
	{"latest/meta-data/public-ipv4", "10.1.1.2"},
	{"latest/openstack/2012-08-10/meta_data.json", metadataContent},
}

func UseTestMetadata(local bool) {
	if local {
		metadataHost = "file:"
	} else {
		metadataHost = origMetadataHost
	}
}

var origMetadataJSON = metadataJSON

func UseMetadataJSON(path string) {
	if path != "" {
		metadataJSON = path
	} else {
		metadataJSON = origMetadataJSON
	}
}

var originalShortAttempt = shortAttempt
var originalLongAttempt = longAttempt

// ShortTimeouts sets the timeouts to a short period as we
// know that the testing server doesn't get better with time,
// and this reduces the test time from 30s to 3s.
func ShortTimeouts(short bool) {
	if short {
		shortAttempt = trivial.AttemptStrategy{
			Total: 0.25e9,
			Delay: 0.01e9,
		}
		longAttempt = shortAttempt
	} else {
		shortAttempt = originalShortAttempt
		longAttempt = originalLongAttempt
	}
}

var ShortAttempt = &shortAttempt

func DeleteStorageContent(s environs.Storage) error {
	return s.(*storage).deleteAll()
}

// WritablePublicStorage returns a Storage instance which is authorised to write to the PublicStorage bucket.
// It is used by tests which need to upload files.
func WritablePublicStorage(e environs.Environ) environs.Storage {
	ecfg := e.(*environ).ecfg()
	authModeCfg := AuthMode(ecfg.authMode())
	writablePublicStorage := &storage{
		containerName: ecfg.publicBucket(),
		swift:         swift.New(e.(*environ).client(ecfg, authModeCfg)),
	}

	// Ensure the container exists.
	err := writablePublicStorage.makeContainer(ecfg.publicBucket(), swift.PublicRead)
	if err != nil {
		panic(fmt.Errorf("cannot create writable public container: %v", err))
	}
	return writablePublicStorage
}
func InstanceAddress(addresses map[string][]nova.IPAddress) (string, error) {
	return instanceAddress(addresses)
}

func FindInstanceSpec(e environs.Environ, series, arch, flavor string) (imageId, flavorId string, err error) {
	env := e.(*environ)
	spec, err := findInstanceSpec(env, &instanceConstraint{
		series: series,
		arch:   arch,
		region: env.ecfg().region(),
		flavor: flavor,
	})
	if err == nil {
		imageId = spec.imageId
		flavorId = spec.flavorId
	}
	return
}

func SetUseFloatingIP(e environs.Environ, val bool) {
	env := e.(*environ)
	env.ecfg().attrs["use-floating-ip"] = val
}

func DefaultInstanceType(e environs.Environ) string {
	ecfg := e.(*environ).ecfg()
	return ecfg.defaultInstanceType()
}

// ImageDetails specify parameters used to start a test machine for the live tests.
type ImageDetails struct {
	Flavor  string
	ImageId string
}

type BootstrapState struct {
	StateInstances []state.InstanceId
}

func LoadState(e environs.Environ) (*BootstrapState, error) {
	s, err := e.(*environ).loadState()
	if err != nil {
		return nil, err
	}
	return &BootstrapState{s.StateInstances}, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The Chromium 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 main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_driver/go/td"
)

func main() {
	var (
		projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
		taskId    = flag.String("task_id", "", "ID of this task.")
		bot       = flag.String("bot", "", "Name of the task.")
		output    = flag.String("o", "", "Dump JSON step data to the given file, or stdout if -.")
		local     = flag.Bool("local", true, "Running locally (else on the bots)?")

		resources = flag.String("resources", "resources", "Passed to fm -i.")
		script    = flag.String("script", "", "File (or - for stdin) with one job per line.")
	)
	ctx := td.StartRun(projectId, taskId, bot, output, local)
	defer td.EndRun(ctx)

	actualStderr := os.Stderr
	if *local {
		// Task Driver echoes every exec.Run() stdout and stderr to the console,
		// which makes it hard to find failures (especially stdout).  Send them to /dev/null.
		devnull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
		if err != nil {
			td.Fatal(ctx, err)
		}
		os.Stdout = devnull
		os.Stderr = devnull
	}

	if flag.NArg() < 1 {
		td.Fatalf(ctx, "Please pass an fm binary.")
	}
	fm := flag.Arg(0)

	// Run `fm <flag>` to find the names of all linked GMs or tests.
	query := func(flag string) []string {
		stdout := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout}
		cmd.Args = append(cmd.Args, "-i", *resources)
		cmd.Args = append(cmd.Args, flag)
		if err := exec.Run(ctx, cmd); err != nil {
			td.Fatal(ctx, err)
		}

		lines := []string{}
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
		return lines
	}
	gms := query("--listGMs")
	tests := query("--listTests")

	type Work struct {
		Sources []string // Passed to FM -s: names of gms/tests, paths to image files, .skps, etc.
		Flags   []string // Other flags to pass to FM: --ct 565, --msaa 16, etc.
	}
	todo := []Work{}

	// Parse a job like "gms b=cpu ct=8888" into Work{Sources=<all GMs>, Flags={-b,cpu,--ct,8888}}.
	parse := func(job []string) (w Work) {
		for _, token := range job {
			// Everything after # is a comment.
			if strings.HasPrefix(token, "#") {
				break
			}

			// Treat "gm" or "gms" as a shortcut for all known GMs.
			if token == "gm" || token == "gms" {
				w.Sources = append(w.Sources, gms...)
				continue
			}
			// Same for tests.
			if token == "test" || token == "tests" {
				w.Sources = append(w.Sources, tests...)
				continue
			}

			// Is this a flag to pass through to FM?
			if parts := strings.Split(token, "="); len(parts) == 2 {
				f := "-"
				if len(parts[0]) > 1 {
					f += "-"
				}
				f += parts[0]

				w.Flags = append(w.Flags, f, parts[1])
				continue
			}

			// Anything else must be the name of a source for FM to run.
			w.Sources = append(w.Sources, token)
		}
		return
	}

	// Parse one job from the command line, handy for ad hoc local runs.
	todo = append(todo, parse(flag.Args()[1:]))

	// Any number of jobs can come from -script.
	if *script != "" {
		file := os.Stdin
		if *script != "-" {
			file, err := os.Open(*script)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer file.Close()
		}
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			todo = append(todo, parse(strings.Fields(scanner.Text())))
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
	}

	// If we're a bot (or acting as if we are one), add its work too.
	if *bot != "" {
		parts := strings.Split(*bot, "-")
		OS := parts[1]

		// For no reason but as a demo, skip GM aarectmodes and test GoodHash.
		filter := func(in []string, test func(string) bool) (out []string) {
			for _, s := range in {
				if test(s) {
					out = append(out, s)
				}
			}
			return
		}
		if OS == "Debian10" {
			gms = filter(gms, func(s string) bool { return s != "aarectmodes" })
			tests = filter(tests, func(s string) bool { return s != "GoodHash" })
		}

		// You could use parse() here if you like, but it's just as easy to make Work{} directly.
		work := func(sources []string, flags string) {
			todo = append(todo, Work{sources, strings.Fields(flags)})
		}
		work(tests, "-b cpu")
		work(gms, "-b cpu")
		work(gms, "-b cpu --skvm")
	}

	// We'll try to spread our work roughly evenly over a number of worker goroutines.
	// We can't combine Work with different Flags, but we can do the opposite,
	// splitting a single Work into smaller Work units with the same Flags,
	// even all the way down to a single Source.  So we'll optimistically run
	// batches of Sources together, but if a batch fails or crashes, we'll
	// split it up and re-run one at a time to find the precise failures.
	var failures int32 = 0
	wg := &sync.WaitGroup{}
	worker := func(queue chan Work) {
		for w := range queue {
			stdout := &bytes.Buffer{}
			stderr := &bytes.Buffer{}
			cmd := &exec.Command{Name: fm, Stdout: stdout, Stderr: stderr}
			cmd.Args = append(cmd.Args, "-i", *resources)
			cmd.Args = append(cmd.Args, "-s")
			cmd.Args = append(cmd.Args, w.Sources...)
			cmd.Args = append(cmd.Args, w.Flags...)
			if err := exec.Run(ctx, cmd); err != nil {
				if len(w.Sources) == 1 {
					// If a source ran alone and failed, that's just a failure.
					atomic.AddInt32(&failures, 1)
					td.FailStep(ctx, err)
					if *local {
						lines := []string{}
						scanner := bufio.NewScanner(stderr)
						for scanner.Scan() {
							lines = append(lines, scanner.Text())
						}
						if err := scanner.Err(); err != nil {
							td.Fatal(ctx, err)
						}

						fmt.Fprintf(actualStderr, "%v %v #failed:\n\t%v\n",
							cmd.Name,
							strings.Join(cmd.Args, " "),
							strings.Join(lines, "\n\t"))
					}
				} else {
					// If a batch fails, retry each individually.
					for _, source := range w.Sources {
						// Requeuing Work from the workers makes sizing the chan buffer tricky:
						// we don't ever want this `queue <-` to block a worker on a full buffer.
						wg.Add(1)
						queue <- Work{[]string{source}, w.Flags}
					}
				}
			}
			wg.Done()
		}
	}

	workers := runtime.NumCPU()
	queue := make(chan Work, 1<<20) // Huge buffer to avoid having to be smart about requeuing.
	for i := 0; i < workers; i++ {
		go worker(queue)
	}

	for _, w := range todo {
		if len(w.Sources) == 0 {
			continue // A blank or commented job line from -script or the command line.
		}

		// Shuffle the sources randomly as a cheap way to approximate evenly expensive batches.
		// (Intentionally not rand.Seed()'d to stay deterministically reproducible.)
		rand.Shuffle(len(w.Sources), func(i, j int) {
			w.Sources[i], w.Sources[j] = w.Sources[j], w.Sources[i]
		})

		// Round batch sizes up so there's at least one source per batch.
		batch := (len(w.Sources) + workers - 1) / workers
		util.ChunkIter(len(w.Sources), batch, func(start, end int) error {
			wg.Add(1)
			queue <- Work{w.Sources[start:end], w.Flags}
			return nil
		})
	}
	wg.Wait()

	if failures > 0 {
		if *local {
			// td.Fatalf() would work fine, but barfs up a panic that we don't need to see.
			fmt.Fprintf(actualStderr, "%v runs of %v failed after retries.\n", failures, fm)
			os.Exit(1)
		} else {
			td.Fatalf(ctx, "%v runs of %v failed after retries.", failures, fm)
		}
	}
}
<commit_msg>fetch known hashes<commit_after>// Copyright 2020 The Chromium 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 main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_driver/go/td"
)

func main() {
	var (
		projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
		taskId    = flag.String("task_id", "", "ID of this task.")
		bot       = flag.String("bot", "", "Name of the task.")
		output    = flag.String("o", "", "Dump JSON step data to the given file, or stdout if -.")
		local     = flag.Bool("local", true, "Running locally (else on the bots)?")

		resources = flag.String("resources", "resources", "Passed to fm -i.")
		script    = flag.String("script", "", "File (or - for stdin) with one job per line.")
	)
	ctx := td.StartRun(projectId, taskId, bot, output, local)
	defer td.EndRun(ctx)

	actualStdout := os.Stdout
	actualStderr := os.Stderr
	if *local {
		// Task Driver echoes every exec.Run() stdout and stderr to the console,
		// which makes it hard to find failures (especially stdout).  Send them to /dev/null.
		devnull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
		if err != nil {
			td.Fatal(ctx, err)
		}
		os.Stdout = devnull
		os.Stderr = devnull
	}

	if flag.NArg() < 1 {
		td.Fatalf(ctx, "Please pass an fm binary.")
	}
	fm := flag.Arg(0)

	// Run `fm <flag>` to find the names of all linked GMs or tests.
	query := func(flag string) []string {
		stdout := &bytes.Buffer{}
		cmd := &exec.Command{Name: fm, Stdout: stdout}
		cmd.Args = append(cmd.Args, "-i", *resources)
		cmd.Args = append(cmd.Args, flag)
		if err := exec.Run(ctx, cmd); err != nil {
			td.Fatal(ctx, err)
		}

		lines := []string{}
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
		return lines
	}
	gms := query("--listGMs")
	tests := query("--listTests")

	// Query Gold for all known hashes when running as a bot.
	known := map[string]bool{
		"0832f708a97acc6da385446384647a8f": true, // MD5 of passing unit test.
	}
	if *bot != "" {
		func() {
			url := "https://storage.googleapis.com/skia-infra-gm/hash_files/gold-prod-hashes.txt"
			resp, err := http.Get(url)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer resp.Body.Close()

			scanner := bufio.NewScanner(resp.Body)
			for scanner.Scan() {
				known[scanner.Text()] = true
			}
			if err := scanner.Err(); err != nil {
				td.Fatal(ctx, err)
			}

			fmt.Fprintf(actualStdout, "Gold knew %v unique hashes.\n", len(known))
		}()
	}

	type Work struct {
		Sources []string // Passed to FM -s: names of gms/tests, paths to image files, .skps, etc.
		Flags   []string // Other flags to pass to FM: --ct 565, --msaa 16, etc.
	}
	todo := []Work{}

	// Parse a job like "gms b=cpu ct=8888" into Work{Sources=<all GMs>, Flags={-b,cpu,--ct,8888}}.
	parse := func(job []string) (w Work) {
		for _, token := range job {
			// Everything after # is a comment.
			if strings.HasPrefix(token, "#") {
				break
			}

			// Treat "gm" or "gms" as a shortcut for all known GMs.
			if token == "gm" || token == "gms" {
				w.Sources = append(w.Sources, gms...)
				continue
			}
			// Same for tests.
			if token == "test" || token == "tests" {
				w.Sources = append(w.Sources, tests...)
				continue
			}

			// Is this a flag to pass through to FM?
			if parts := strings.Split(token, "="); len(parts) == 2 {
				f := "-"
				if len(parts[0]) > 1 {
					f += "-"
				}
				f += parts[0]

				w.Flags = append(w.Flags, f, parts[1])
				continue
			}

			// Anything else must be the name of a source for FM to run.
			w.Sources = append(w.Sources, token)
		}
		return
	}

	// Parse one job from the command line, handy for ad hoc local runs.
	todo = append(todo, parse(flag.Args()[1:]))

	// Any number of jobs can come from -script.
	if *script != "" {
		file := os.Stdin
		if *script != "-" {
			file, err := os.Open(*script)
			if err != nil {
				td.Fatal(ctx, err)
			}
			defer file.Close()
		}
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			todo = append(todo, parse(strings.Fields(scanner.Text())))
		}
		if err := scanner.Err(); err != nil {
			td.Fatal(ctx, err)
		}
	}

	// If we're a bot (or acting as if we are one), add its work too.
	if *bot != "" {
		parts := strings.Split(*bot, "-")
		OS := parts[1]

		// For no reason but as a demo, skip GM aarectmodes and test GoodHash.
		filter := func(in []string, test func(string) bool) (out []string) {
			for _, s := range in {
				if test(s) {
					out = append(out, s)
				}
			}
			return
		}
		if OS == "Debian10" {
			gms = filter(gms, func(s string) bool { return s != "aarectmodes" })
			tests = filter(tests, func(s string) bool { return s != "GoodHash" })
		}

		// You could use parse() here if you like, but it's just as easy to make Work{} directly.
		work := func(sources []string, flags string) {
			todo = append(todo, Work{sources, strings.Fields(flags)})
		}
		work(tests, "-b cpu")
		work(gms, "-b cpu")
		work(gms, "-b cpu --skvm")
	}

	// We'll try to spread our work roughly evenly over a number of worker goroutines.
	// We can't combine Work with different Flags, but we can do the opposite,
	// splitting a single Work into smaller Work units with the same Flags,
	// even all the way down to a single Source.  So we'll optimistically run
	// batches of Sources together, but if a batch fails or crashes, we'll
	// split it up and re-run one at a time to find the precise failures.
	var failures int32 = 0
	wg := &sync.WaitGroup{}
	worker := func(queue chan Work) {
		for w := range queue {
			stdout := &bytes.Buffer{}
			stderr := &bytes.Buffer{}
			cmd := &exec.Command{Name: fm, Stdout: stdout, Stderr: stderr}
			cmd.Args = append(cmd.Args, "-i", *resources)
			cmd.Args = append(cmd.Args, "-s")
			cmd.Args = append(cmd.Args, w.Sources...)
			cmd.Args = append(cmd.Args, w.Flags...)
			// TODO: when len(w.Sources) == 1, add -w ... to cmd.Args to write a .png for upload.

			// On cmd failure or unknown hash, we'll split the Work batch up into individual reruns.
			requeue := func() {
				// Requeuing Work from the workers is what makes sizing the chan buffer tricky:
				// we don't ever want these `queue <-` to block a worker because of a full buffer.
				for _, source := range w.Sources {
					wg.Add(1)
					queue <- Work{[]string{source}, w.Flags}
				}
			}

			if err := exec.Run(ctx, cmd); err != nil {
				if len(w.Sources) == 1 {
					// If a source ran alone and failed, that's just a failure.
					atomic.AddInt32(&failures, 1)
					td.FailStep(ctx, err)
					if *local {
						lines := []string{}
						scanner := bufio.NewScanner(stderr)
						for scanner.Scan() {
							lines = append(lines, scanner.Text())
						}
						if err := scanner.Err(); err != nil {
							td.Fatal(ctx, err)
						}

						fmt.Fprintf(actualStderr, "%v %v #failed:\n\t%v\n",
							cmd.Name,
							strings.Join(cmd.Args, " "),
							strings.Join(lines, "\n\t"))
					}
				} else {
					// If a batch of sources failed, break up the batch to isolate the failures.
					requeue()
				}
			} else {
				// FM completed successfully.  Scan stdout for any unknown hash.
				unknown := func() string {
					if *bot != "" { // The map known[] is only filled when *bot != "".
						scanner := bufio.NewScanner(stdout)
						for scanner.Scan() {
							if parts := strings.Fields(scanner.Text()); len(parts) == 3 {
								md5 := parts[1]
								if !known[md5] {
									return md5
								}
							}
						}
						if err := scanner.Err(); err != nil {
							td.Fatal(ctx, err)
						}
					}
					return ""
				}()

				if unknown != "" {
					if len(w.Sources) == 1 {
						// TODO upload .png with goldctl.
						fmt.Fprintf(actualStdout, "%v %v #%v\n",
							cmd.Name,
							strings.Join(cmd.Args, " "),
							unknown)
					} else {
						// Split the batch to run individually and TODO, write .pngs.
						requeue()
					}
				}
			}
			wg.Done()
		}
	}

	workers := runtime.NumCPU()
	queue := make(chan Work, 1<<20) // Huge buffer to avoid having to be smart about requeuing.
	for i := 0; i < workers; i++ {
		go worker(queue)
	}

	for _, w := range todo {
		if len(w.Sources) == 0 {
			continue // A blank or commented job line from -script or the command line.
		}

		// Shuffle the sources randomly as a cheap way to approximate evenly expensive batches.
		// (Intentionally not rand.Seed()'d to stay deterministically reproducible.)
		rand.Shuffle(len(w.Sources), func(i, j int) {
			w.Sources[i], w.Sources[j] = w.Sources[j], w.Sources[i]
		})

		// Round batch sizes up so there's at least one source per batch.
		batch := (len(w.Sources) + workers - 1) / workers
		util.ChunkIter(len(w.Sources), batch, func(start, end int) error {
			wg.Add(1)
			queue <- Work{w.Sources[start:end], w.Flags}
			return nil
		})
	}
	wg.Wait()

	if failures > 0 {
		if *local {
			// td.Fatalf() would work fine, but barfs up a panic that we don't need to see.
			fmt.Fprintf(actualStderr, "%v runs of %v failed after retries.\n", failures, fm)
			os.Exit(1)
		} else {
			td.Fatalf(ctx, "%v runs of %v failed after retries.", failures, fm)
		}
	}
}
<|endoftext|>


<commit_before>// call somatic variants with multiple tumor samples associated with 1 normal. This uses
// the genotype likelihood approach from bcbio and speedseq but allows for multiple tumor
// samples. Anything that meets the criteria will have a "PASS" or "." FILTER and a list
// of tumor samples with the somatic variant in INFO["SOMATIC"]
// Copied from bcbio and speedseq.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/brentp/vcfgo"
	"github.com/brentp/xopen"
)

func getTumorLOD(gls []float64) float64 {
	lod := -1.0
	if len(gls) < 3 {
		return lod
	}
	for _, gl := range gls[1:] {
		t := gl - gls[0]
		if t > lod {
			lod = t
		}
	}
	return lod
}

func getNormalLOD(gls []float64, thresh float64) float64 {
	lod := 1e7
	if len(gls) < 3 {
		return thresh
	}
	for _, gl := range gls[1:] {
		t := gls[0] - gl
		if t < lod {
			lod = t
		}
	}
	return lod
}

func somaticLOD(normalGLs []float64, tumorGLs []float64, thresh float64) bool {
	tumorLOD := getTumorLOD(tumorGLs)
	normalLOD := getNormalLOD(normalGLs, thresh)
	return tumorLOD >= thresh && normalLOD >= thresh
}

func getFreq(refCount int, altCounts []int) float64 {
	ac := 0
	for _, a := range altCounts {
		ac += a
	}
	return float64(ac) / float64(refCount+ac)
}

var THRESH_RATIO = 2.7

func somaticFreqs(normal *vcfgo.SampleGenotype, tumor *vcfgo.SampleGenotype) bool {
	var nFreq, tFreq float64
	nrd, err := normal.RefDepth()
	if err != nil {
		nFreq = 0.0
	} else {
		nads, err := normal.AltDepths()
		if err != nil {
			nFreq = 0.0
		} else {
			nFreq = getFreq(nrd, nads)
		}
	}
	trd, err := tumor.RefDepth()
	if err != nil {
		tFreq = 0.0
	} else {
		tads, err := tumor.AltDepths()
		if err != nil {
			tFreq = 0.0
		} else {
			tFreq = getFreq(trd, tads)
		}
	}
	return nFreq <= 0.001 || nFreq <= tFreq/THRESH_RATIO

}

func Somatics(v *vcfgo.Variant, normalIdx int) []string {
	thresh := 3.5

	normal := v.Samples[normalIdx]
	somatics := make([]string, 0)
	for i, tumor := range v.Samples {
		if i == normalIdx {
			continue
		}
		if !somaticFreqs(normal, tumor) {
			continue
		}
		if somaticLOD(normal.GL, tumor.GL, thresh) {
			somatics = append(somatics, v.Header.SampleNames[i])
		}
	}
	return somatics
}

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func main() {
	index := flag.Int("index", 0, "0-based index of the normal sample")
	flag.Parse()
	vcfs := flag.Args()
	if len(vcfs) != 1 {
		fmt.Printf("---------------- call somatic variants present in any tumor sample ----------------\n")
		fmt.Printf("----- uses the method from bcbio and speedseq but for multiple tumor samples ------\n")
		fmt.Printf("%s -index 0 normal-and-tumors.vcf.gz\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(1)
	}

	fhr, err := xopen.Ropen(vcfs[0])
	check(err)

	rdr, err := vcfgo.NewReader(fhr, false)
	check(err)

	fhw, err := xopen.Wopen("-")
	check(err)

	hdr := rdr.Header
	hdr.Filters["NOT_SOMATIC"] = "not somatic between normal and any tumor"

	hdr.Infos["SOMATIC"] = &vcfgo.Info{Id: "SOMATIC", Number: "1", Type: "String", Description: "Tumor samples with somatic event"}

	wtr, err := vcfgo.NewWriter(fhw, hdr)
	check(err)

	for v := rdr.Read(); v != nil; v = rdr.Read() {

		somatics := Somatics(v, *index)
		if len(somatics) > 0 {
			v.Info.Add("SOMATIC", strings.Join(somatics, "|"))
		} else {
			if v.Filter == "." || v.Filter == "PASS" {
				v.Filter = ""
			}
			if v.Filter != "" {
				v.Filter += ";"
			}
			v.Filter += "NOT_SOMATIC"
		}
		wtr.WriteVariant(v)
	}
}
<commit_msg>close (and flush)<commit_after>// call somatic variants with multiple tumor samples associated with 1 normal. This uses
// the genotype likelihood approach from bcbio and speedseq but allows for multiple tumor
// samples. Anything that meets the criteria will have a "PASS" or "." FILTER and a list
// of tumor samples with the somatic variant in INFO["SOMATIC"]
// Copied from bcbio and speedseq.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/brentp/vcfgo"
	"github.com/brentp/xopen"
)

func getTumorLOD(gls []float64) float64 {
	lod := -1.0
	if len(gls) < 3 {
		return lod
	}
	for _, gl := range gls[1:] {
		t := gl - gls[0]
		if t > lod {
			lod = t
		}
	}
	return lod
}

func getNormalLOD(gls []float64, thresh float64) float64 {
	lod := 1e7
	if len(gls) < 3 {
		return thresh
	}
	for _, gl := range gls[1:] {
		t := gls[0] - gl
		if t < lod {
			lod = t
		}
	}
	return lod
}

func somaticLOD(normalGLs []float64, tumorGLs []float64, thresh float64) bool {
	tumorLOD := getTumorLOD(tumorGLs)
	normalLOD := getNormalLOD(normalGLs, thresh)
	return tumorLOD >= thresh && normalLOD >= thresh
}

func getFreq(refCount int, altCounts []int) float64 {
	ac := 0
	for _, a := range altCounts {
		ac += a
	}
	return float64(ac) / float64(refCount+ac)
}

var THRESH_RATIO = 2.7

func somaticFreqs(normal *vcfgo.SampleGenotype, tumor *vcfgo.SampleGenotype) bool {
	var nFreq, tFreq float64
	nrd, err := normal.RefDepth()
	if err != nil {
		nFreq = 0.0
	} else {
		nads, err := normal.AltDepths()
		if err != nil {
			nFreq = 0.0
		} else {
			nFreq = getFreq(nrd, nads)
		}
	}
	trd, err := tumor.RefDepth()
	if err != nil {
		tFreq = 0.0
	} else {
		tads, err := tumor.AltDepths()
		if err != nil {
			tFreq = 0.0
		} else {
			tFreq = getFreq(trd, tads)
		}
	}
	return nFreq <= 0.001 || nFreq <= tFreq/THRESH_RATIO

}

func Somatics(v *vcfgo.Variant, normalIdx int) []string {
	thresh := 3.5

	normal := v.Samples[normalIdx]
	somatics := make([]string, 0)
	for i, tumor := range v.Samples {
		if i == normalIdx {
			continue
		}
		if !somaticFreqs(normal, tumor) {
			continue
		}
		if somaticLOD(normal.GL, tumor.GL, thresh) {
			somatics = append(somatics, v.Header.SampleNames[i])
		}
	}
	return somatics
}

func check(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func main() {
	index := flag.Int("index", 0, "0-based index of the normal sample")
	flag.Parse()
	vcfs := flag.Args()
	if len(vcfs) != 1 {
		fmt.Printf("---------------- call somatic variants present in any tumor sample ----------------\n")
		fmt.Printf("----- uses the method from bcbio and speedseq but for multiple tumor samples ------\n")
		fmt.Printf("%s -index 0 normal-and-tumors.vcf.gz\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(1)
	}

	fhr, err := xopen.Ropen(vcfs[0])
	check(err)

	rdr, err := vcfgo.NewReader(fhr, false)
	check(err)

	fhw, err := xopen.Wopen("-")
	check(err)

	hdr := rdr.Header
	hdr.Filters["NOT_SOMATIC"] = "not somatic between normal and any tumor"

	hdr.Infos["SOMATIC"] = &vcfgo.Info{Id: "SOMATIC", Number: "1", Type: "String", Description: "Tumor samples with somatic event"}

	wtr, err := vcfgo.NewWriter(fhw, hdr)
	check(err)

	log.Printf("using %s as the normal sample\n", hdr.SampleNames[*index])

	for v := rdr.Read(); v != nil; v = rdr.Read() {

		somatics := Somatics(v, *index)
		if len(somatics) > 0 {
			v.Info.Add("SOMATIC", strings.Join(somatics, "|"))
		} else {
			if v.Filter == "." || v.Filter == "PASS" {
				v.Filter = ""
			}
			if v.Filter != "" {
				v.Filter += ";"
			}
			v.Filter += "NOT_SOMATIC"
		}
		wtr.WriteVariant(v)
	}
	fhw.Close()
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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 containersource

import (
	"context"
	"fmt"
	"strings"

	"go.uber.org/zap"
	"k8s.io/client-go/rest"

	"github.com/knative/eventing-sources/pkg/apis/sources/v1alpha1"
	"github.com/knative/eventing-sources/pkg/controller/containersource/resources"
	"github.com/knative/eventing-sources/pkg/controller/sinks"
	"github.com/knative/pkg/logging"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/tools/record"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

type reconciler struct {
	client        client.Client
	scheme        *runtime.Scheme
	dynamicClient dynamic.Interface
	recorder      record.EventRecorder
}

// Reconcile compares the actual state with the desired, and attempts to
// converge the two.
func (r *reconciler) Reconcile(ctx context.Context, object runtime.Object) (runtime.Object, error) {
	logger := logging.FromContext(ctx)

	source, ok := object.(*v1alpha1.ContainerSource)
	if !ok {
		logger.Errorf("could not find container source %v\n", object)
		return object, nil
	}

	// See if the source has been deleted
	accessor, err := meta.Accessor(source)
	if err != nil {
		logger.Warnf("Failed to get metadata accessor: %s", zap.Error(err))
		return object, err
	}
	// No need to reconcile if the source has been marked for deletion.
	deletionTimestamp := accessor.GetDeletionTimestamp()
	if deletionTimestamp != nil {
		return object, nil
	}

	source.Status.InitializeConditions()

	args := &resources.ContainerArguments{
		Name:               source.Name,
		Namespace:          source.Namespace,
		Image:              source.Spec.Image,
		Args:               source.Spec.Args,
		Env:                source.Spec.Env,
		ServiceAccountName: source.Spec.ServiceAccountName,
	}

	err = r.setSinkURIArg(source, args)
	if err != nil {
		return source, err
	}

	deploy, err := r.getDeployment(ctx, source)
	if err != nil {
		if errors.IsNotFound(err) {
			deploy, err = r.createDeployment(ctx, source, nil, args)
			if err != nil {
				r.recorder.Eventf(source, corev1.EventTypeNormal, "DeploymentBlocked", "waiting for %v", err)
				return object, err
			}
			r.recorder.Eventf(source, corev1.EventTypeNormal, "Deployed", "Created deployment %q", deploy.Name)
			source.Status.MarkDeploying("Deploying", "Created deployment %s", args.Name)
		} else {
			return source, err
		}
	} else {
		if deploy.Status.ReadyReplicas > 0 {
			source.Status.MarkDeployed()
		}
	}

	return source, nil
}

func (r *reconciler) setSinkURIArg(source *v1alpha1.ContainerSource, args *resources.ContainerArguments) error {
	if uri, ok := sinkArg(source); ok {
		args.SinkInArgs = true
		source.Status.MarkSink(uri)
		return nil
	}

	if source.Spec.Sink == nil {
		source.Status.MarkNoSink("Missing", "")
		return fmt.Errorf("Sink missing from spec")
	}

	uri, err := sinks.GetSinkURI(r.dynamicClient, source.Spec.Sink, source.Namespace)
	if err != nil {
		source.Status.MarkNoSink("NotFound", "")
		return err
	}
	source.Status.MarkSink(uri)
	args.Sink = uri

	return nil
}

func sinkArg(source *v1alpha1.ContainerSource) (string, bool) {
	for _, a := range source.Spec.Args {
		if strings.HasPrefix(a, "--sink=") {
			return strings.Replace(a, "--sink=", "", -1), true
		}
	}
	return "", false
}

func (r *reconciler) getDeployment(ctx context.Context, source *v1alpha1.ContainerSource) (*appsv1.Deployment, error) {
	logger := logging.FromContext(ctx)

	list := &appsv1.DeploymentList{}
	err := r.client.List(
		ctx,
		&client.ListOptions{
			Namespace:     source.Namespace,
			LabelSelector: labels.Everything(),
			// TODO this is here because the fake client needs it.
			// Remove this when it's no longer needed.
			Raw: &metav1.ListOptions{
				TypeMeta: metav1.TypeMeta{
					APIVersion: appsv1.SchemeGroupVersion.String(),
					Kind:       "Deployment",
				},
			},
		},
		list)
	if err != nil {
		logger.Errorf("Unable to list deployments: %v", zap.Error(err))
		return nil, err
	}
	for _, c := range list.Items {
		if metav1.IsControlledBy(&c, source) {
			return &c, nil
		}
	}
	return nil, errors.NewNotFound(schema.GroupResource{}, "")
}

func (r *reconciler) createDeployment(ctx context.Context, source *v1alpha1.ContainerSource, org *appsv1.Deployment, args *resources.ContainerArguments) (*appsv1.Deployment, error) {
	deployment := resources.MakeDeployment(org, args)

	if err := controllerutil.SetControllerReference(source, deployment, r.scheme); err != nil {
		return nil, err
	}

	if err := r.client.Create(ctx, deployment); err != nil {
		return nil, err
	}
	return deployment, nil
}

func (r *reconciler) InjectClient(c client.Client) error {
	r.client = c
	return nil
}

func (r *reconciler) InjectConfig(c *rest.Config) error {
	var err error
	r.dynamicClient, err = dynamic.NewForConfig(c)
	return err
}
<commit_msg>Update Deployment spec on ContainerSource change (#84)<commit_after>/*
Copyright 2018 The Knative 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 containersource

import (
	"context"
	"fmt"
	"strings"

	"go.uber.org/zap"
	"k8s.io/client-go/rest"

	"github.com/knative/eventing-sources/pkg/apis/sources/v1alpha1"
	"github.com/knative/eventing-sources/pkg/controller/containersource/resources"
	"github.com/knative/eventing-sources/pkg/controller/sinks"
	"github.com/knative/pkg/logging"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/equality"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/tools/record"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

type reconciler struct {
	client        client.Client
	scheme        *runtime.Scheme
	dynamicClient dynamic.Interface
	recorder      record.EventRecorder
}

// Reconcile compares the actual state with the desired, and attempts to
// converge the two.
func (r *reconciler) Reconcile(ctx context.Context, object runtime.Object) (runtime.Object, error) {
	logger := logging.FromContext(ctx)

	source, ok := object.(*v1alpha1.ContainerSource)
	if !ok {
		logger.Errorf("could not find container source %v\n", object)
		return object, nil
	}

	// See if the source has been deleted
	accessor, err := meta.Accessor(source)
	if err != nil {
		logger.Warnf("Failed to get metadata accessor: %s", zap.Error(err))
		return object, err
	}
	// No need to reconcile if the source has been marked for deletion.
	deletionTimestamp := accessor.GetDeletionTimestamp()
	if deletionTimestamp != nil {
		return object, nil
	}

	source.Status.InitializeConditions()

	args := &resources.ContainerArguments{
		Name:               source.Name,
		Namespace:          source.Namespace,
		Image:              source.Spec.Image,
		Args:               source.Spec.Args,
		Env:                source.Spec.Env,
		ServiceAccountName: source.Spec.ServiceAccountName,
	}

	err = r.setSinkURIArg(source, args)
	if err != nil {
		return source, err
	}

	deploy, err := r.getDeployment(ctx, source)
	if err != nil {
		if errors.IsNotFound(err) {
			deploy, err = r.createDeployment(ctx, source, nil, args)
			if err != nil {
				r.recorder.Eventf(source, corev1.EventTypeNormal, "DeploymentBlocked", "waiting for %v", err)
				return object, err
			}
			r.recorder.Eventf(source, corev1.EventTypeNormal, "Deployed", "Created deployment %q", deploy.Name)
			source.Status.MarkDeploying("Deploying", "Created deployment %s", args.Name)
			// Since the Deployment has just been created, there's nothing more
			// to do until it gets a status. This ContainerSource will be reconciled
			// again when the Deployment is updated.
			return object, nil
		}
		return object, err
	}

	// Update Deployment spec if it's changed
	expected := resources.MakeDeployment(nil, args)
	if !equality.Semantic.DeepEqual(deploy.Spec, expected.Spec) {
		deploy.Spec = expected.Spec
		if r.client.Update(ctx, deploy); err != nil {
			return object, err
		}
	}

	// Update source status
	if deploy.Status.ReadyReplicas > 0 {
		source.Status.MarkDeployed()
	}

	return source, nil
}

func (r *reconciler) setSinkURIArg(source *v1alpha1.ContainerSource, args *resources.ContainerArguments) error {
	if uri, ok := sinkArg(source); ok {
		args.SinkInArgs = true
		source.Status.MarkSink(uri)
		return nil
	}

	if source.Spec.Sink == nil {
		source.Status.MarkNoSink("Missing", "")
		return fmt.Errorf("Sink missing from spec")
	}

	uri, err := sinks.GetSinkURI(r.dynamicClient, source.Spec.Sink, source.Namespace)
	if err != nil {
		source.Status.MarkNoSink("NotFound", "")
		return err
	}
	source.Status.MarkSink(uri)
	args.Sink = uri

	return nil
}

func sinkArg(source *v1alpha1.ContainerSource) (string, bool) {
	for _, a := range source.Spec.Args {
		if strings.HasPrefix(a, "--sink=") {
			return strings.Replace(a, "--sink=", "", -1), true
		}
	}
	return "", false
}

func (r *reconciler) getDeployment(ctx context.Context, source *v1alpha1.ContainerSource) (*appsv1.Deployment, error) {
	logger := logging.FromContext(ctx)

	list := &appsv1.DeploymentList{}
	err := r.client.List(
		ctx,
		&client.ListOptions{
			Namespace:     source.Namespace,
			LabelSelector: labels.Everything(),
			// TODO this is here because the fake client needs it.
			// Remove this when it's no longer needed.
			Raw: &metav1.ListOptions{
				TypeMeta: metav1.TypeMeta{
					APIVersion: appsv1.SchemeGroupVersion.String(),
					Kind:       "Deployment",
				},
			},
		},
		list)
	if err != nil {
		logger.Errorf("Unable to list deployments: %v", zap.Error(err))
		return nil, err
	}
	for _, c := range list.Items {
		if metav1.IsControlledBy(&c, source) {
			return &c, nil
		}
	}
	return nil, errors.NewNotFound(schema.GroupResource{}, "")
}

func (r *reconciler) createDeployment(ctx context.Context, source *v1alpha1.ContainerSource, org *appsv1.Deployment, args *resources.ContainerArguments) (*appsv1.Deployment, error) {
	deployment := resources.MakeDeployment(org, args)

	if err := controllerutil.SetControllerReference(source, deployment, r.scheme); err != nil {
		return nil, err
	}

	if err := r.client.Create(ctx, deployment); err != nil {
		return nil, err
	}
	return deployment, nil
}

func (r *reconciler) InjectClient(c client.Client) error {
	r.client = c
	return nil
}

func (r *reconciler) InjectConfig(c *rest.Config) error {
	var err error
	r.dynamicClient, err = dynamic.NewForConfig(c)
	return err
}
<|endoftext|>


<commit_before>package chroot

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"path/filepath"
)

// StepCopyFiles copies some files from the host into the chroot environment.
//
// Produces:
//   copy_files_cleanup CleanupFunc - A function to clean up the copied files
//   early.
type StepCopyFiles struct {
	files []string
}

func (s *StepCopyFiles) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	mountPath := state.Get("mount_path").(string)
	ui := state.Get("ui").(packer.Ui)
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)
	stderr := new(bytes.Buffer)

	s.files = make([]string, 0, len(config.CopyFiles))
	if len(config.CopyFiles) > 0 {
		ui.Say("Copying files from host to chroot...")
		for _, path := range config.CopyFiles {
			ui.Message(path)
			chrootPath := filepath.Join(mountPath, path)
			log.Printf("Copying '%s' to '%s'", path, chrootPath)

			cmdText, err := wrappedCommand(fmt.Sprintf("cp %s %s", path, chrootPath))
			if err != nil {
				err := fmt.Errorf("Error building copy command: %s", err)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

			stderr.Reset()
			cmd := ShellCommand(cmdText)
			cmd.Stderr = stderr
			if err := cmd.Run(); err != nil {
				err := fmt.Errorf(
					"Error copying file: %s\nnStderr: %s", err, stderr.String())
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

			s.files = append(s.files, chrootPath)
		}
	}

	state.Put("copy_files_cleanup", s)
	return multistep.ActionContinue
}

func (s *StepCopyFiles) Cleanup(state multistep.StateBag) {
	ui := state.Get("ui").(packer.Ui)
	if err := s.CleanupFunc(state); err != nil {
		ui.Error(err.Error())
	}
}

func (s *StepCopyFiles) CleanupFunc(state multistep.StateBag) error {
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)
	if s.files != nil {
		for _, file := range s.files {
			log.Printf("Removing: %s", file)
			localCmdText, err := wrappedCommand(fmt.Sprintf("rm -f %s", file))
			if err != nil {
				return err
			}

			localCmd := ShellCommand(localCmdText)
			if err := localCmd.Run(); err != nil {
				return err
			}
		}
	}

	s.files = nil
	return nil
}
<commit_msg>builder/amazon/chroot:<commit_after>package chroot

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"path/filepath"
)

// StepCopyFiles copies some files from the host into the chroot environment.
//
// Produces:
//   copy_files_cleanup CleanupFunc - A function to clean up the copied files
//   early.
type StepCopyFiles struct {
	files []string
}

func (s *StepCopyFiles) Run(state multistep.StateBag) multistep.StepAction {
	config := state.Get("config").(*Config)
	mountPath := state.Get("mount_path").(string)
	ui := state.Get("ui").(packer.Ui)
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)
	stderr := new(bytes.Buffer)

	s.files = make([]string, 0, len(config.CopyFiles))
	if len(config.CopyFiles) > 0 {
		ui.Say("Copying files from host to chroot...")
		for _, path := range config.CopyFiles {
			ui.Message(path)
			chrootPath := filepath.Join(mountPath, path)
			log.Printf("Copying '%s' to '%s'", path, chrootPath)

			cmdText, err := wrappedCommand(fmt.Sprintf("cp --remove-destination %s %s", path, chrootPath))
			if err != nil {
				err := fmt.Errorf("Error building copy command: %s", err)
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

			stderr.Reset()
			cmd := ShellCommand(cmdText)
			cmd.Stderr = stderr
			if err := cmd.Run(); err != nil {
				err := fmt.Errorf(
					"Error copying file: %s\nnStderr: %s", err, stderr.String())
				state.Put("error", err)
				ui.Error(err.Error())
				return multistep.ActionHalt
			}

			s.files = append(s.files, chrootPath)
		}
	}

	state.Put("copy_files_cleanup", s)
	return multistep.ActionContinue
}

func (s *StepCopyFiles) Cleanup(state multistep.StateBag) {
	ui := state.Get("ui").(packer.Ui)
	if err := s.CleanupFunc(state); err != nil {
		ui.Error(err.Error())
	}
}

func (s *StepCopyFiles) CleanupFunc(state multistep.StateBag) error {
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)
	if s.files != nil {
		for _, file := range s.files {
			log.Printf("Removing: %s", file)
			localCmdText, err := wrappedCommand(fmt.Sprintf("rm -f %s", file))
			if err != nil {
				return err
			}

			localCmd := ShellCommand(localCmdText)
			if err := localCmd.Run(); err != nil {
				return err
			}
		}
	}

	s.files = nil
	return nil
}
<|endoftext|>