<commit_before>package physical

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"

	log "github.com/mgutz/logxi/v1"

	"github.com/hashicorp/vault/helper/jsonutil"
)

// FileBackend is a physical backend that stores data on disk
// at a given file path. It can be used for durable single server
// situations, or to develop locally where durability is not critical.
//
// WARNING: the file backend implementation is currently extremely unsafe
// and non-performant. It is meant mostly for local testing and development.
// It can be improved in the future.
type FileBackend struct {
	Path   string
	l      sync.Mutex
	logger log.Logger
}

// newFileBackend constructs a Filebackend using the given directory
func newFileBackend(conf map[string]string, logger log.Logger) (Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("'path' must be set")
	}

	return &FileBackend{
		Path:   path,
		logger: logger,
	}, nil
}

func (b *FileBackend) Delete(k string) error {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	fullPath := filepath.Join(path, key)

	// If the path doesn't exist return success; this is in line with Vault's
	// expected behavior and we don't want to check for an empty directory if
	// we couldn't even find the path in the first place.
	err := os.Remove(fullPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		} else {
			return err
		}
	}

	// Check for the directory being empty and remove if so, with another
	// additional guard for the path not existing
	dir, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		} else {
			return err
		}
	}

	list, err := dir.Readdir(1)
	dir.Close()
	if err != nil && err != io.EOF {
		return err
	}

	// If we have no entries, it's an empty directory; remove it
	if err == io.EOF || list == nil || len(list) == 0 {
		err = os.Remove(path)
		if err != nil {
			return err
		}
	}

	return nil
}

func (b *FileBackend) Get(k string) (*Entry, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	var entry Entry
	if err := jsonutil.DecodeJSONFromReader(f, &entry); err != nil {
		return nil, err
	}

	return &entry, nil
}

func (b *FileBackend) Put(entry *Entry) error {
	path, key := b.path(entry.Key)

	b.l.Lock()
	defer b.l.Unlock()

	// Make the parent tree
	if err := os.MkdirAll(path, 0755); err != nil {
		return err
	}

	// JSON encode the entry and write it
	f, err := os.OpenFile(
		filepath.Join(path, key),
		os.O_CREATE|os.O_TRUNC|os.O_WRONLY,
		0600)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(entry)
}

func (b *FileBackend) List(prefix string) ([]string, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path := b.Path
	if prefix != "" {
		path = filepath.Join(path, prefix)
	}

	// Read the directory contents
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	names, err := f.Readdirnames(-1)
	if err != nil {
		return nil, err
	}

	for i, name := range names {
		if name[0] == '_' {
			names[i] = name[1:]
		} else {
			names[i] = name + "/"
		}
	}

	return names, nil
}

func (b *FileBackend) path(k string) (string, string) {
	path := filepath.Join(b.Path, k)
	key := filepath.Base(path)
	path = filepath.Dir(path)
	return path, "_" + key
}
<commit_msg>Fix file backend so that it properly removes nested secrets.<commit_after>package physical

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	log "github.com/mgutz/logxi/v1"

	"github.com/hashicorp/vault/helper/jsonutil"
)

// FileBackend is a physical backend that stores data on disk
// at a given file path. It can be used for durable single server
// situations, or to develop locally where durability is not critical.
//
// WARNING: the file backend implementation is currently extremely unsafe
// and non-performant. It is meant mostly for local testing and development.
// It can be improved in the future.
type FileBackend struct {
	Path   string
	l      sync.Mutex
	logger log.Logger
}

// newFileBackend constructs a Filebackend using the given directory
func newFileBackend(conf map[string]string, logger log.Logger) (Backend, error) {
	path, ok := conf["path"]
	if !ok {
		return nil, fmt.Errorf("'path' must be set")
	}

	return &FileBackend{
		Path:   path,
		logger: logger,
	}, nil
}

func (b *FileBackend) Delete(path string) error {
	b.l.Lock()
	defer b.l.Unlock()

	basePath, key := b.path(path)
	fullPath := filepath.Join(basePath, key)

	err := os.Remove(fullPath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("Failed to remove `%s`: %v", fullPath, err)
	}

	err = b.cleanupLogicalPath(path)

	return err
}

// cleanupLogicalPath is used to remove all empty nodes, begining with deepest
// one, aborting on first non-empty one, up to top-level node.
func (b *FileBackend) cleanupLogicalPath(path string) error {
	nodes := strings.Split(path, "/")
	for i := len(nodes) - 1; i > 0; i-- {
		fullPath := b.Path + "/" + strings.Join(nodes[:i], "/")

		dir, err := os.Open(fullPath)
		if err != nil {
			if os.IsNotExist(err) {
				return nil
			} else {
				return err
			}
		}

		list, err := dir.Readdir(1)
		dir.Close()
		if err != nil && err != io.EOF {
			return err
		}

		// If we have no entries, it's an empty directory; remove it
		if err == io.EOF || list == nil || len(list) == 0 {
			err = os.Remove(fullPath)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (b *FileBackend) Get(k string) (*Entry, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path, key := b.path(k)
	path = filepath.Join(path, key)

	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	var entry Entry
	if err := jsonutil.DecodeJSONFromReader(f, &entry); err != nil {
		return nil, err
	}

	return &entry, nil
}

func (b *FileBackend) Put(entry *Entry) error {
	path, key := b.path(entry.Key)

	b.l.Lock()
	defer b.l.Unlock()

	// Make the parent tree
	if err := os.MkdirAll(path, 0755); err != nil {
		return err
	}

	// JSON encode the entry and write it
	f, err := os.OpenFile(
		filepath.Join(path, key),
		os.O_CREATE|os.O_TRUNC|os.O_WRONLY,
		0600)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(entry)
}

func (b *FileBackend) List(prefix string) ([]string, error) {
	b.l.Lock()
	defer b.l.Unlock()

	path := b.Path
	if prefix != "" {
		path = filepath.Join(path, prefix)
	}

	// Read the directory contents
	f, err := os.Open(path)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}

		return nil, err
	}
	defer f.Close()

	names, err := f.Readdirnames(-1)
	if err != nil {
		return nil, err
	}

	for i, name := range names {
		if name[0] == '_' {
			names[i] = name[1:]
		} else {
			names[i] = name + "/"
		}
	}

	return names, nil
}

func (b *FileBackend) path(k string) (string, string) {
	path := filepath.Join(b.Path, k)
	key := filepath.Base(path)
	path = filepath.Dir(path)
	return path, "_" + key
}
<|endoftext|>


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

package bpf

import (
	"fmt"
	"os/exec"
	"path"
	"sync"

	log "github.com/Sirupsen/logrus"
	"golang.org/x/sys/unix"
)

var (
	// Path to where bpffs is mounted (default: /sys/fs/bpf)
	mapRoot = "/sys/fs/bpf"

	// Prefix for all maps (default: tc/globals)
	mapPrefix = "tc/globals"

	// Set to true on first get request to detect misorder
	lockedDown = false
	once       sync.Once
)

func lockDown() {
	lockedDown = true
}

func SetMapRoot(path string) {
	if lockedDown {
		panic("SetMapRoot() call after MapRoot was read")
	}
	mapRoot = path

	mountMutex.Lock()
	for _, m := range delayedOpens {
		m.OpenOrCreate()
	}

	mounted = true
	delayedOpens = []*Map{}
	mountMutex.Unlock()
}

func GetMapRoot() string {
	once.Do(lockDown)
	return mapRoot
}

func SetMapPrefix(path string) {
	if lockedDown {
		panic("SetMapPrefix() call after MapPrefix was read")
	}
	mapPrefix = path
}

func GetMapPrefix() string {
	once.Do(lockDown)
	return mapPrefix
}

func MapPrefixPath() string {
	once.Do(lockDown)
	return path.Join(mapRoot, mapPrefix)
}

func MapPath(name string) string {
	once.Do(lockDown)
	return path.Join(mapRoot, mapPrefix, name)
}

var (
	mountOnce    sync.Once
	mountMutex   sync.Mutex
	delayedOpens = []*Map{}
	mounted      bool
)

// OpenAfterMount schedules a map to be opened/created after the BPF filesystem
// has been mounted. If the filesystem is already mounted, the map is
// opened/created immediately.
func OpenAfterMount(m *Map) error {
	mountMutex.Lock()
	defer mountMutex.Unlock()

	if mounted {
		_, err := m.OpenOrCreate()
		return err
	}

	delayedOpens = append(delayedOpens, m)
	return nil
}

//isBpffs check if the path is a valid bpf filesystem
func isBpffs(path string) bool {
	// This is the value of the BPF Filesystem. If is into the container the
	// mountpoint doesn't provide enough information. Defined on uapi/linux/magic.h
	magic := uint32(0xCAFE4A11)
	var fsdata unix.Statfs_t
	if err := unix.Statfs(path, &fsdata); err != nil {
		log.Errorf("%s is not mounted", path)
		return false
	}
	return int32(magic) == int32(fsdata.Type)
}

func mountFS() error {
	// Mount BPF Map directory if not already done
	args := []string{"-q", mapRoot}
	_, err := exec.Command("mountpoint", args...).CombinedOutput()
	if err != nil {
		args = []string{"bpffs", mapRoot, "-t", "bpf"}
		out, err := exec.Command("mount", args...).CombinedOutput()
		if err != nil {
			return fmt.Errorf("Command execution failed: %s\n%s", err, out)
		}
	}
	if !isBpffs(mapRoot) {
		log.Fatalf("BPF: '%s' is not mounted as BPF filesystem.", mapRoot)
	}
	mountMutex.Lock()
	for _, m := range delayedOpens {
		m.OpenOrCreate()
	}

	mounted = true
	delayedOpens = []*Map{}
	mountMutex.Unlock()

	return nil
}

// MountFS mounts the BPF filesystem and then opens/creates all maps which have
// previously been scheduled to be opened/created
func MountFS() {
	mountOnce.Do(func() {
		if err := mountFS(); err != nil {
			log.WithError(err).Fatalf("Unable to mount BPF filesystem")
		}
	})
}
<commit_msg>Fix #1257 Refuse to start if bpffs filesystem is mounted multiple times (#1473)<commit_after>// Copyright 2016-2017 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bpf

import (
	"bytes"
	"fmt"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"

	log "github.com/Sirupsen/logrus"
	"golang.org/x/sys/unix"
)

var (
	// Path to where bpffs is mounted (default: /sys/fs/bpf)
	mapRoot = "/sys/fs/bpf"

	// Prefix for all maps (default: tc/globals)
	mapPrefix = "tc/globals"

	// Set to true on first get request to detect misorder
	lockedDown = false
	once       sync.Once
)

func lockDown() {
	lockedDown = true
}

func SetMapRoot(path string) {
	if lockedDown {
		panic("SetMapRoot() call after MapRoot was read")
	}
	mapRoot = path

	mountMutex.Lock()
	for _, m := range delayedOpens {
		m.OpenOrCreate()
	}

	mounted = true
	delayedOpens = []*Map{}
	mountMutex.Unlock()
}

func GetMapRoot() string {
	once.Do(lockDown)
	return mapRoot
}

func SetMapPrefix(path string) {
	if lockedDown {
		panic("SetMapPrefix() call after MapPrefix was read")
	}
	mapPrefix = path
}

func GetMapPrefix() string {
	once.Do(lockDown)
	return mapPrefix
}

func MapPrefixPath() string {
	once.Do(lockDown)
	return path.Join(mapRoot, mapPrefix)
}

func MapPath(name string) string {
	once.Do(lockDown)
	return path.Join(mapRoot, mapPrefix, name)
}

var (
	mountOnce    sync.Once
	mountMutex   sync.Mutex
	delayedOpens = []*Map{}
	mounted      bool
)

// OpenAfterMount schedules a map to be opened/created after the BPF filesystem
// has been mounted. If the filesystem is already mounted, the map is
// opened/created immediately.
func OpenAfterMount(m *Map) error {
	mountMutex.Lock()
	defer mountMutex.Unlock()

	if mounted {
		_, err := m.OpenOrCreate()
		return err
	}

	delayedOpens = append(delayedOpens, m)
	return nil
}

//isBpffs check if the path is a valid bpf filesystem
func isBpffs(path string) bool {
	// This is the value of the BPF Filesystem. If is into the container the
	// mountpoint doesn't provide enough information. Defined on uapi/linux/magic.h
	magic := uint32(0xCAFE4A11)
	var fsdata unix.Statfs_t
	if err := unix.Statfs(path, &fsdata); err != nil {
		log.Errorf("%s is not mounted", path)
		return false
	}
	return int32(magic) == int32(fsdata.Type)
}

func mountCmdPipe(cmds []*exec.Cmd) (mountCmdOutput, mountCmdStandardError []byte, mountCmdError error) {

	// We need atleast one command to pipe.
	if len(cmds) < 1 {
		return nil, nil, nil
	}

	// Total output of commands.
	var output bytes.Buffer
	var stderr bytes.Buffer

	lastCmd := len(cmds) - 1
	for i, cmd := range cmds[:lastCmd] {
		var err error
		// We need to connect every command's stdin to the previous command's stdout
		if cmds[i+1].Stdin, err = cmd.StdoutPipe(); err != nil {
			return nil, nil, err
		}
		// We need to connect each command's stderr to a buffer
		cmd.Stderr = &stderr
	}

	// Connect the output and error for the last command
	cmds[lastCmd].Stdout, cmds[lastCmd].Stderr = &output, &stderr

	// Let's start each command
	for _, cmd := range cmds {
		if err := cmd.Start(); err != nil {
			return output.Bytes(), stderr.Bytes(), err
		}
	}

	// We wait for each command to complete
	for _, cmd := range cmds {
		if err := cmd.Wait(); err != nil {
			return output.Bytes(), stderr.Bytes(), err
		}
	}

	// Return the output and the standard error
	return output.Bytes(), stderr.Bytes(), nil
}

func mountFS() error {
	// Mount BPF Map directory if not already done
	args := []string{"-q", mapRoot}
	_, err := exec.Command("mountpoint", args...).CombinedOutput()
	if err != nil {
		args = []string{"bpffs", mapRoot, "-t", "bpf"}
		out, err := exec.Command("mount", args...).CombinedOutput()
		if err != nil {
			return fmt.Errorf("command execution failed: %s\n%s", err, out)
		}
	} else { // Already mounted. We need to fail if mounted multiple times.

		// Execute the following command to find multiple bpffs mount points
		// % mount | grep "<mapRoot> " | wc -l | cut -f1 -d' '
		newmapRoot := mapRoot + " " // Append space to ignore /sys/fs/bpf/xdp and /sys/fs/bpf/ip mountpoints.
		cmds := []*exec.Cmd{
			exec.Command("mount"),
			exec.Command("grep", newmapRoot),
			exec.Command("wc", "-l"),
			exec.Command("cut", "-f1", "-d "),
		}

		output, stderr, _ := mountCmdPipe(cmds)

		if len(stderr) > 0 {
			return fmt.Errorf("command execution failed: %s", stderr)
		}

		// Strip the newline character at the end.
		parts := strings.Split(string(output), "\n")

		// Convert the string to integer
		num, err := strconv.ParseInt(parts[0], 10, 32)
		if err != nil {
			return fmt.Errorf("command execution failed: %s", err)
		}

		if num > 1 {
			return fmt.Errorf("multiple mount points detected at %s", mapRoot)
		}

	}
	if !isBpffs(mapRoot) {
		log.Fatalf("BPF: '%s' is not mounted as BPF filesystem.", mapRoot)
	}
	mountMutex.Lock()
	for _, m := range delayedOpens {
		m.OpenOrCreate()
	}

	mounted = true
	delayedOpens = []*Map{}
	mountMutex.Unlock()

	return nil
}

// MountFS mounts the BPF filesystem and then opens/creates all maps which have
// previously been scheduled to be opened/created
func MountFS() {
	mountOnce.Do(func() {
		if err := mountFS(); err != nil {
			log.WithError(err).Fatalf("Unable to mount BPF filesystem")
		}
	})
}
<|endoftext|>


<commit_before>package crank

import (
	"fmt"
	"net/rpc"
	"syscall"
	"time"
)

type API struct {
	m *Manager
}

func NewRPCServer(m *Manager) *rpc.Server {
	server := rpc.NewServer()
	api := &API{m}
	err := server.RegisterName("crank", api)
	if err != nil {
		panic(err) // Coding error
	}
	return server
}

// Used by other query structs
type ProcessQuery struct {
	Starting bool
	Ready    bool
	Stopping bool
	Pid      int
}

type processFilter func(*Process) *Process

// START

type StartQuery struct {
	Command      []string
	StartTimeout int
	StopTimeout  int
	Wait         bool
}

type StartReply struct {
}

func (self *API) Start(query *StartQuery, reply *StartReply) error {
	// FIXME: concurrency access
	config := self.m.config.clone()

	if len(query.Command) > 0 {
		config.Command = query.Command
	}

	if query.StartTimeout > 0 {
		config.StartTimeout = time.Duration(query.StartTimeout) * time.Second
	}

	if query.StopTimeout > 0 {
		config.StopTimeout = time.Duration(query.StopTimeout) * time.Second
	}

	// TODO: support the query.Wait flag

	self.m.Start(config)

	return nil
}

// PS

type PsQuery struct {
	ProcessQuery
}

type PsReply struct {
	PS []*ProcessInfo
}

type ProcessInfo struct {
	Pid   int
	State string
	Usage *syscall.Rusage
	Err   error
}

func (pi *ProcessInfo) String() string {
	since := func(tv syscall.Timeval) time.Duration {
		return time.Duration(tv.Nano())
	}

	if pi.Err != nil {
		return fmt.Sprintf("%d %s\n", pi.Pid, pi.State, pi.Err)
	} else {
		usage := pi.Usage
		return fmt.Sprintf("%d %s %v %v %v\n", pi.Pid, pi.State, since(usage.Utime), since(usage.Stime), ByteCount(usage.Maxrss))
	}
}

func (self *API) Ps(query *PsQuery, reply *PsReply) error {
	ss := self.m.childs

	if query.Pid > 0 {
		ss = ss.choose(func(p *Process, _ ProcessState) bool {
			return p.Pid() == query.Pid
		})
	}

	if query.Starting || query.Ready || query.Stopping {
		ss = ss.choose(func(p *Process, state ProcessState) bool {
			if query.Starting && (state == PROCESS_STARTING) {
				return true
			}
			if query.Ready && (state == PROCESS_READY) {
				return true
			}
			if query.Stopping && (state == PROCESS_STOPPING) {
				return true
			}
			return false
		})
	}

	reply.PS = make([]*ProcessInfo, 0, ss.len())
	for s, state := range ss {
		usage, err := s.Usage()
		reply.PS = append(reply.PS, &ProcessInfo{s.Pid(), state.String(), usage, err})
	}

	fmt.Println(query, reply)
	return nil
}

type KillQuery struct {
	ProcessQuery
	Signal string
	Wait   bool
}

type KillReply struct {
}

func (self *API) Kill(query *KillQuery, reply *KillReply) (err error) {
	var sig syscall.Signal
	if query.Signal == "" {
		sig = syscall.SIGTERM
	} else {
		if sig, err = str2signal(query.Signal); err != nil {
			return
		}
	}

	var ss processSet
	if query.Starting || query.Ready || query.Stopping || query.Pid > 0 {
		ss = self.m.childs
	} else {
		// Empty set
		ss = EmptyProcessSet
	}

	if query.Starting || query.Ready || query.Stopping {
		ss = ss.choose(func(p *Process, state ProcessState) bool {
			if query.Starting && (state == PROCESS_STARTING) {
				return true
			}
			if query.Ready && (state == PROCESS_READY) {
				return true
			}
			if query.Stopping && (state == PROCESS_STOPPING) {
				return true
			}
			return false
		})
	}

	if query.Pid > 0 {
		ss = ss.choose(func(p *Process, _ ProcessState) bool {
			return p.Pid() == query.Pid
		})
	}

	ss.each(func(s *Process) {
		s.Signal(sig)
	})

	fmt.Println(query, reply)
	return
}
<commit_msg>Show the processes' command on `crankctl ps`<commit_after>package crank

import (
	"fmt"
	"net/rpc"
	"syscall"
	"time"
)

type API struct {
	m *Manager
}

func NewRPCServer(m *Manager) *rpc.Server {
	server := rpc.NewServer()
	api := &API{m}
	err := server.RegisterName("crank", api)
	if err != nil {
		panic(err) // Coding error
	}
	return server
}

// Used by other query structs
type ProcessQuery struct {
	Starting bool
	Ready    bool
	Stopping bool
	Pid      int
}

type processFilter func(*Process) *Process

// START

type StartQuery struct {
	Command      []string
	StartTimeout int
	StopTimeout  int
	Wait         bool
}

type StartReply struct {
}

func (self *API) Start(query *StartQuery, reply *StartReply) error {
	// FIXME: concurrency access
	config := self.m.config.clone()

	if len(query.Command) > 0 {
		config.Command = query.Command
	}

	if query.StartTimeout > 0 {
		config.StartTimeout = time.Duration(query.StartTimeout) * time.Second
	}

	if query.StopTimeout > 0 {
		config.StopTimeout = time.Duration(query.StopTimeout) * time.Second
	}

	// TODO: support the query.Wait flag

	self.m.Start(config)

	return nil
}

// PS

type PsQuery struct {
	ProcessQuery
}

type PsReply struct {
	PS []*ProcessInfo
}

type ProcessInfo struct {
	Pid     int
	State   string
	Command []string
	Usage   *syscall.Rusage
	Err     error
}

func (pi *ProcessInfo) String() string {
	since := func(tv syscall.Timeval) time.Duration {
		return time.Duration(tv.Nano())
	}

	if pi.Err != nil {
		return fmt.Sprintf("%d %s %v %v\n", pi.Pid, pi.State, pi.Command, pi.Err)
	} else {
		usage := pi.Usage
		return fmt.Sprintf("%d %s %v %v %v %v\n", pi.Pid, pi.State, pi.Command, since(usage.Utime), since(usage.Stime), ByteCount(usage.Maxrss))
	}
}

func (self *API) Ps(query *PsQuery, reply *PsReply) error {
	ss := self.m.childs

	if query.Pid > 0 {
		ss = ss.choose(func(p *Process, _ ProcessState) bool {
			return p.Pid() == query.Pid
		})
	}

	if query.Starting || query.Ready || query.Stopping {
		ss = ss.choose(func(p *Process, state ProcessState) bool {
			if query.Starting && (state == PROCESS_STARTING) {
				return true
			}
			if query.Ready && (state == PROCESS_READY) {
				return true
			}
			if query.Stopping && (state == PROCESS_STOPPING) {
				return true
			}
			return false
		})
	}

	reply.PS = make([]*ProcessInfo, 0, ss.len())
	for s, state := range ss {
		usage, err := s.Usage()
		reply.PS = append(reply.PS, &ProcessInfo{s.Pid(), state.String(), s.config.Command, usage, err})
	}

	fmt.Println(query, reply)
	return nil
}

type KillQuery struct {
	ProcessQuery
	Signal string
	Wait   bool
}

type KillReply struct {
}

func (self *API) Kill(query *KillQuery, reply *KillReply) (err error) {
	var sig syscall.Signal
	if query.Signal == "" {
		sig = syscall.SIGTERM
	} else {
		if sig, err = str2signal(query.Signal); err != nil {
			return
		}
	}

	var ss processSet
	if query.Starting || query.Ready || query.Stopping || query.Pid > 0 {
		ss = self.m.childs
	} else {
		// Empty set
		ss = EmptyProcessSet
	}

	if query.Starting || query.Ready || query.Stopping {
		ss = ss.choose(func(p *Process, state ProcessState) bool {
			if query.Starting && (state == PROCESS_STARTING) {
				return true
			}
			if query.Ready && (state == PROCESS_READY) {
				return true
			}
			if query.Stopping && (state == PROCESS_STOPPING) {
				return true
			}
			return false
		})
	}

	if query.Pid > 0 {
		ss = ss.choose(func(p *Process, _ ProcessState) bool {
			return p.Pid() == query.Pid
		})
	}

	ss.each(func(s *Process) {
		s.Signal(sig)
	})

	fmt.Println(query, reply)
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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

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

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

package helm // import "k8s.io/helm/pkg/helm"

import (
	"fmt"
	"sync"

	"k8s.io/helm/pkg/proto/hapi/chart"
	"k8s.io/helm/pkg/proto/hapi/release"
	rls "k8s.io/helm/pkg/proto/hapi/services"
	"k8s.io/helm/pkg/proto/hapi/version"
)

// FakeClient implements Interface
type FakeClient struct {
	Rels      []*release.Release
	Responses map[string]release.TestRun_Status
	Err       error
}

var _ Interface = &FakeClient{}
var _ Interface = (*FakeClient)(nil)

// ListReleases lists the current releases
func (c *FakeClient) ListReleases(opts ...ReleaseListOption) (*rls.ListReleasesResponse, error) {
	resp := &rls.ListReleasesResponse{
		Count:    int64(len(c.Rels)),
		Releases: c.Rels,
	}
	return resp, c.Err
}

// InstallRelease returns a response with the first Release on the fake release client
func (c *FakeClient) InstallRelease(chStr, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
	return &rls.InstallReleaseResponse{
		Release: c.Rels[0],
	}, nil
}

// InstallReleaseFromChart returns a response with the first Release on the fake release client
func (c *FakeClient) InstallReleaseFromChart(chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
	return &rls.InstallReleaseResponse{
		Release: c.Rels[0],
	}, nil
}

// DeleteRelease returns nil, nil
func (c *FakeClient) DeleteRelease(rlsName string, opts ...DeleteOption) (*rls.UninstallReleaseResponse, error) {
	return nil, nil
}

// UpdateRelease returns nil, nil
func (c *FakeClient) UpdateRelease(rlsName string, chStr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
	return nil, nil
}

// GetVersion returns a fake version
func (c *FakeClient) GetVersion(opts ...VersionOption) (*rls.GetVersionResponse, error) {
	return &rls.GetVersionResponse{
		Version: &version.Version{
			SemVer: "1.2.3-fakeclient+testonly",
		},
	}, nil
}

// UpdateReleaseFromChart returns nil, nil
func (c *FakeClient) UpdateReleaseFromChart(rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
	return nil, nil
}

// RollbackRelease returns nil, nil
func (c *FakeClient) RollbackRelease(rlsName string, opts ...RollbackOption) (*rls.RollbackReleaseResponse, error) {
	return nil, nil
}

// ReleaseStatus returns a release status response with info from the matching release name in the fake release client.
func (c *FakeClient) ReleaseStatus(rlsName string, opts ...StatusOption) (*rls.GetReleaseStatusResponse, error) {
	if len(c.Rels) > 0 {
		for _, release := range c.Rels {
			if release.Name == rlsName {
				return &rls.GetReleaseStatusResponse{
					Name:      release.Name,
					Info:      release.Info,
					Namespace: release.Namespace,
				}, nil
			}
		}
	}
	return nil, fmt.Errorf("No such release: %s", rlsName)
}

// ReleaseContent returns the configuration for the first release in the fake release client
func (c *FakeClient) ReleaseContent(rlsName string, opts ...ContentOption) (resp *rls.GetReleaseContentResponse, err error) {
	if len(c.Rels) > 0 {
		for _, release := range c.Rels {
			if release.Name == rlsName {
				resp = &rls.GetReleaseContentResponse{
					Release: release,
				}
			}
		}
	}
	return resp, c.Err
}

// ReleaseHistory returns a release's revision history.
func (c *FakeClient) ReleaseHistory(rlsName string, opts ...HistoryOption) (*rls.GetHistoryResponse, error) {
	return &rls.GetHistoryResponse{Releases: c.Rels}, c.Err
}

// RunReleaseTest executes a pre-defined tests on a release
func (c *FakeClient) RunReleaseTest(rlsName string, opts ...ReleaseTestOption) (<-chan *rls.TestReleaseResponse, <-chan error) {

	results := make(chan *rls.TestReleaseResponse)
	errc := make(chan error, 1)

	go func() {
		var wg sync.WaitGroup
		for m, s := range c.Responses {
			wg.Add(1)

			go func(msg string, status release.TestRun_Status) {
				defer wg.Done()
				results <- &rls.TestReleaseResponse{Msg: msg, Status: status}
			}(m, s)
		}

		wg.Wait()
		close(results)
		close(errc)
	}()

	return results, errc
}

// Option returns the fake release client
func (c *FakeClient) Option(opt ...Option) Interface {
	return c
}
<commit_msg>Update comment for accuracy<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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

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

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

package helm // import "k8s.io/helm/pkg/helm"

import (
	"fmt"
	"sync"

	"k8s.io/helm/pkg/proto/hapi/chart"
	"k8s.io/helm/pkg/proto/hapi/release"
	rls "k8s.io/helm/pkg/proto/hapi/services"
	"k8s.io/helm/pkg/proto/hapi/version"
)

// FakeClient implements Interface
type FakeClient struct {
	Rels      []*release.Release
	Responses map[string]release.TestRun_Status
	Err       error
}

var _ Interface = &FakeClient{}
var _ Interface = (*FakeClient)(nil)

// ListReleases lists the current releases
func (c *FakeClient) ListReleases(opts ...ReleaseListOption) (*rls.ListReleasesResponse, error) {
	resp := &rls.ListReleasesResponse{
		Count:    int64(len(c.Rels)),
		Releases: c.Rels,
	}
	return resp, c.Err
}

// InstallRelease returns a response with the first Release on the fake release client
func (c *FakeClient) InstallRelease(chStr, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
	return &rls.InstallReleaseResponse{
		Release: c.Rels[0],
	}, nil
}

// InstallReleaseFromChart returns a response with the first Release on the fake release client
func (c *FakeClient) InstallReleaseFromChart(chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
	return &rls.InstallReleaseResponse{
		Release: c.Rels[0],
	}, nil
}

// DeleteRelease returns nil, nil
func (c *FakeClient) DeleteRelease(rlsName string, opts ...DeleteOption) (*rls.UninstallReleaseResponse, error) {
	return nil, nil
}

// UpdateRelease returns nil, nil
func (c *FakeClient) UpdateRelease(rlsName string, chStr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
	return nil, nil
}

// GetVersion returns a fake version
func (c *FakeClient) GetVersion(opts ...VersionOption) (*rls.GetVersionResponse, error) {
	return &rls.GetVersionResponse{
		Version: &version.Version{
			SemVer: "1.2.3-fakeclient+testonly",
		},
	}, nil
}

// UpdateReleaseFromChart returns nil, nil
func (c *FakeClient) UpdateReleaseFromChart(rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
	return nil, nil
}

// RollbackRelease returns nil, nil
func (c *FakeClient) RollbackRelease(rlsName string, opts ...RollbackOption) (*rls.RollbackReleaseResponse, error) {
	return nil, nil
}

// ReleaseStatus returns a release status response with info from the matching release name in the fake release client.
func (c *FakeClient) ReleaseStatus(rlsName string, opts ...StatusOption) (*rls.GetReleaseStatusResponse, error) {
	if len(c.Rels) > 0 {
		for _, release := range c.Rels {
			if release.Name == rlsName {
				return &rls.GetReleaseStatusResponse{
					Name:      release.Name,
					Info:      release.Info,
					Namespace: release.Namespace,
				}, nil
			}
		}
	}
	return nil, fmt.Errorf("No such release: %s", rlsName)
}

// ReleaseContent returns the configuration for the matching release name in the fake release client.
func (c *FakeClient) ReleaseContent(rlsName string, opts ...ContentOption) (resp *rls.GetReleaseContentResponse, err error) {
	if len(c.Rels) > 0 {
		for _, release := range c.Rels {
			if release.Name == rlsName {
				resp = &rls.GetReleaseContentResponse{
					Release: release,
				}
			}
		}
	}
	return resp, c.Err
}

// ReleaseHistory returns a release's revision history.
func (c *FakeClient) ReleaseHistory(rlsName string, opts ...HistoryOption) (*rls.GetHistoryResponse, error) {
	return &rls.GetHistoryResponse{Releases: c.Rels}, c.Err
}

// RunReleaseTest executes a pre-defined tests on a release
func (c *FakeClient) RunReleaseTest(rlsName string, opts ...ReleaseTestOption) (<-chan *rls.TestReleaseResponse, <-chan error) {

	results := make(chan *rls.TestReleaseResponse)
	errc := make(chan error, 1)

	go func() {
		var wg sync.WaitGroup
		for m, s := range c.Responses {
			wg.Add(1)

			go func(msg string, status release.TestRun_Status) {
				defer wg.Done()
				results <- &rls.TestReleaseResponse{Msg: msg, Status: status}
			}(m, s)
		}

		wg.Wait()
		close(results)
		close(errc)
	}()

	return results, errc
}

// Option returns the fake release client
func (c *FakeClient) Option(opt ...Option) Interface {
	return c
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package model

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"text/template"

	"github.com/golang/glog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/model/iam"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awstasks"
)

// IAMModelBuilder configures IAM objects
type IAMModelBuilder struct {
	*KopsModelContext

	Lifecycle *fi.Lifecycle
}

var _ fi.ModelBuilder = &IAMModelBuilder{}

const RolePolicyTemplate = `{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": { "Service": "{{ IAMServiceEC2 }}"},
      "Action": "sts:AssumeRole"
    }
  ]
}`

func (b *IAMModelBuilder) Build(c *fi.ModelBuilderContext) error {
	// Collect managed Instance Group roles
	managedRoles := make(map[kops.InstanceGroupRole]bool)

	// Collect Instance Profile ARNs and their associated Instance Group roles
	sharedProfileARNsToIGRole := make(map[string]kops.InstanceGroupRole)
	for _, ig := range b.InstanceGroups {
		if ig.Spec.IAM != nil && ig.Spec.IAM.Profile != nil {
			specProfile := fi.StringValue(ig.Spec.IAM.Profile)
			if matchingRole, ok := sharedProfileARNsToIGRole[specProfile]; ok {
				if matchingRole != ig.Spec.Role {
					return fmt.Errorf("Found IAM instance profile assigned to multiple Instance Group roles %v and %v: %v",
						ig.Spec.Role, sharedProfileARNsToIGRole[specProfile], specProfile)
				}
			} else {
				sharedProfileARNsToIGRole[specProfile] = ig.Spec.Role
			}
		} else {
			managedRoles[ig.Spec.Role] = true
		}
	}

	// Generate IAM tasks for each shared role
	for profileARN, igRole := range sharedProfileARNsToIGRole {
		iamName, err := findCustomAuthNameFromArn(profileARN)
		if err != nil {
			return fmt.Errorf("unable to parse instance profile name from arn %q: %v", profileARN, err)
		}
		err = b.buildIAMTasks(igRole, iamName, c, true)
	}

	// Generate IAM tasks for each managed role
	for igRole := range managedRoles {
		iamName := b.IAMName(igRole)
		err := b.buildIAMTasks(igRole, iamName, c, false)
		if err != nil {
			return err
		}
	}

	return nil
}

func (b *IAMModelBuilder) buildIAMTasks(igRole kops.InstanceGroupRole, iamName string, c *fi.ModelBuilderContext, shared bool) error {
	{ // To minimize diff for easier code review
		var iamRole *awstasks.IAMRole
		{
			rolePolicy, err := b.buildAWSIAMRolePolicy()
			if err != nil {
				return err
			}

			iamRole = &awstasks.IAMRole{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				RolePolicyDocument: fi.WrapResource(rolePolicy),
				ExportWithID:       s(strings.ToLower(string(igRole)) + "s"),
			}
			c.AddTask(iamRole)

		}

		{
			iamPolicy := &iam.PolicyResource{
				Builder: &iam.PolicyBuilder{
					Cluster: b.Cluster,
					Role:    igRole,
					Region:  b.Region,
				},
			}

			// This is slightly tricky; we need to know the hosted zone id,
			// but we might be creating the hosted zone dynamically.

			// TODO: I don't love this technique for finding the task by name & modifying it
			dnsZoneTask, found := c.Tasks["DNSZone/"+b.NameForDNSZone()]
			if found {
				iamPolicy.DNSZone = dnsZoneTask.(*awstasks.DNSZone)
			} else {
				glog.V(2).Infof("Task %q not found; won't set route53 permissions in IAM", "DNSZone/"+b.NameForDNSZone())
			}

			t := &awstasks.IAMRolePolicy{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				Role:           iamRole,
				PolicyDocument: iamPolicy,
			}
			c.AddTask(t)
		}

		var iamInstanceProfile *awstasks.IAMInstanceProfile
		{
			iamInstanceProfile = &awstasks.IAMInstanceProfile{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,
				Shared:    fi.Bool(shared),
			}
			c.AddTask(iamInstanceProfile)
		}

		{
			iamInstanceProfileRole := &awstasks.IAMInstanceProfileRole{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				InstanceProfile: iamInstanceProfile,
				Role:            iamRole,
			}
			c.AddTask(iamInstanceProfileRole)
		}

		// Generate additional policies if needed, and attach to existing role
		{
			additionalPolicy := ""
			if b.Cluster.Spec.AdditionalPolicies != nil {
				roleAsString := reflect.ValueOf(igRole).String()
				additionalPolicies := *(b.Cluster.Spec.AdditionalPolicies)

				additionalPolicy = additionalPolicies[strings.ToLower(roleAsString)]
			}

			additionalPolicyName := "additional." + iamName

			t := &awstasks.IAMRolePolicy{
				Name:      s(additionalPolicyName),
				Lifecycle: b.Lifecycle,

				Role: iamRole,
			}

			if additionalPolicy != "" {
				p := &iam.Policy{
					Version: iam.PolicyDefaultVersion,
				}

				statements := make([]*iam.Statement, 0)
				json.Unmarshal([]byte(additionalPolicy), &statements)
				p.Statement = append(p.Statement, statements...)

				policy, err := p.AsJSON()
				if err != nil {
					return fmt.Errorf("error building IAM policy: %v", err)
				}

				t.PolicyDocument = fi.WrapResource(fi.NewStringResource(policy))
			} else {
				t.PolicyDocument = fi.WrapResource(fi.NewStringResource(""))
			}

			c.AddTask(t)
		}
	}
	return nil
}

// buildAWSIAMRolePolicy produces the AWS IAM role policy for the given role
func (b *IAMModelBuilder) buildAWSIAMRolePolicy() (fi.Resource, error) {
	functions := template.FuncMap{
		"IAMServiceEC2": func() string {
			// IAMServiceEC2 returns the name of the IAM service for EC2 in the current region
			// it is ec2.amazonaws.com everywhere but in cn-north, where it is ec2.amazonaws.com.cn
			switch b.Region {
			case "cn-north-1":
				return "ec2.amazonaws.com.cn"
			case "cn-northwest-1":
				return "ec2.amazonaws.com.cn"
			default:
				return "ec2.amazonaws.com"
			}
		},
	}

	templateResource, err := NewTemplateResource("AWSIAMRolePolicy", RolePolicyTemplate, functions, nil)
	if err != nil {
		return nil, err
	}
	return templateResource, nil
}
<commit_msg>Check errors when parsing JSON on IAM policies<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package model

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strings"
	"text/template"

	"github.com/golang/glog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/model/iam"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awstasks"
)

// IAMModelBuilder configures IAM objects
type IAMModelBuilder struct {
	*KopsModelContext

	Lifecycle *fi.Lifecycle
}

var _ fi.ModelBuilder = &IAMModelBuilder{}

const RolePolicyTemplate = `{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": { "Service": "{{ IAMServiceEC2 }}"},
      "Action": "sts:AssumeRole"
    }
  ]
}`

func (b *IAMModelBuilder) Build(c *fi.ModelBuilderContext) error {
	// Collect managed Instance Group roles
	managedRoles := make(map[kops.InstanceGroupRole]bool)

	// Collect Instance Profile ARNs and their associated Instance Group roles
	sharedProfileARNsToIGRole := make(map[string]kops.InstanceGroupRole)
	for _, ig := range b.InstanceGroups {
		if ig.Spec.IAM != nil && ig.Spec.IAM.Profile != nil {
			specProfile := fi.StringValue(ig.Spec.IAM.Profile)
			if matchingRole, ok := sharedProfileARNsToIGRole[specProfile]; ok {
				if matchingRole != ig.Spec.Role {
					return fmt.Errorf("Found IAM instance profile assigned to multiple Instance Group roles %v and %v: %v",
						ig.Spec.Role, sharedProfileARNsToIGRole[specProfile], specProfile)
				}
			} else {
				sharedProfileARNsToIGRole[specProfile] = ig.Spec.Role
			}
		} else {
			managedRoles[ig.Spec.Role] = true
		}
	}

	// Generate IAM tasks for each shared role
	for profileARN, igRole := range sharedProfileARNsToIGRole {
		iamName, err := findCustomAuthNameFromArn(profileARN)
		if err != nil {
			return fmt.Errorf("unable to parse instance profile name from arn %q: %v", profileARN, err)
		}
		err = b.buildIAMTasks(igRole, iamName, c, true)
	}

	// Generate IAM tasks for each managed role
	for igRole := range managedRoles {
		iamName := b.IAMName(igRole)
		err := b.buildIAMTasks(igRole, iamName, c, false)
		if err != nil {
			return err
		}
	}

	return nil
}

func (b *IAMModelBuilder) buildIAMTasks(igRole kops.InstanceGroupRole, iamName string, c *fi.ModelBuilderContext, shared bool) error {
	{ // To minimize diff for easier code review
		var iamRole *awstasks.IAMRole
		{
			rolePolicy, err := b.buildAWSIAMRolePolicy()
			if err != nil {
				return err
			}

			iamRole = &awstasks.IAMRole{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				RolePolicyDocument: fi.WrapResource(rolePolicy),
				ExportWithID:       s(strings.ToLower(string(igRole)) + "s"),
			}
			c.AddTask(iamRole)

		}

		{
			iamPolicy := &iam.PolicyResource{
				Builder: &iam.PolicyBuilder{
					Cluster: b.Cluster,
					Role:    igRole,
					Region:  b.Region,
				},
			}

			// This is slightly tricky; we need to know the hosted zone id,
			// but we might be creating the hosted zone dynamically.

			// TODO: I don't love this technique for finding the task by name & modifying it
			dnsZoneTask, found := c.Tasks["DNSZone/"+b.NameForDNSZone()]
			if found {
				iamPolicy.DNSZone = dnsZoneTask.(*awstasks.DNSZone)
			} else {
				glog.V(2).Infof("Task %q not found; won't set route53 permissions in IAM", "DNSZone/"+b.NameForDNSZone())
			}

			t := &awstasks.IAMRolePolicy{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				Role:           iamRole,
				PolicyDocument: iamPolicy,
			}
			c.AddTask(t)
		}

		var iamInstanceProfile *awstasks.IAMInstanceProfile
		{
			iamInstanceProfile = &awstasks.IAMInstanceProfile{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,
				Shared:    fi.Bool(shared),
			}
			c.AddTask(iamInstanceProfile)
		}

		{
			iamInstanceProfileRole := &awstasks.IAMInstanceProfileRole{
				Name:      s(iamName),
				Lifecycle: b.Lifecycle,

				InstanceProfile: iamInstanceProfile,
				Role:            iamRole,
			}
			c.AddTask(iamInstanceProfileRole)
		}

		// Generate additional policies if needed, and attach to existing role
		{
			additionalPolicy := ""
			if b.Cluster.Spec.AdditionalPolicies != nil {
				roleAsString := reflect.ValueOf(igRole).String()
				additionalPolicies := *(b.Cluster.Spec.AdditionalPolicies)

				additionalPolicy = additionalPolicies[strings.ToLower(roleAsString)]
			}

			additionalPolicyName := "additional." + iamName

			t := &awstasks.IAMRolePolicy{
				Name:      s(additionalPolicyName),
				Lifecycle: b.Lifecycle,

				Role: iamRole,
			}

			if additionalPolicy != "" {
				p := &iam.Policy{
					Version: iam.PolicyDefaultVersion,
				}

				statements := make([]*iam.Statement, 0)
				if err := json.Unmarshal([]byte(additionalPolicy), &statements); err != nil {
					return fmt.Errorf("additionalPolicy %q is invalid: %v", strings.ToLower(string(igRole)), err)
				}

				p.Statement = append(p.Statement, statements...)

				policy, err := p.AsJSON()
				if err != nil {
					return fmt.Errorf("error building IAM policy: %v", err)
				}

				t.PolicyDocument = fi.WrapResource(fi.NewStringResource(policy))
			} else {
				t.PolicyDocument = fi.WrapResource(fi.NewStringResource(""))
			}

			c.AddTask(t)
		}
	}
	return nil
}

// buildAWSIAMRolePolicy produces the AWS IAM role policy for the given role
func (b *IAMModelBuilder) buildAWSIAMRolePolicy() (fi.Resource, error) {
	functions := template.FuncMap{
		"IAMServiceEC2": func() string {
			// IAMServiceEC2 returns the name of the IAM service for EC2 in the current region
			// it is ec2.amazonaws.com everywhere but in cn-north, where it is ec2.amazonaws.com.cn
			switch b.Region {
			case "cn-north-1":
				return "ec2.amazonaws.com.cn"
			case "cn-northwest-1":
				return "ec2.amazonaws.com.cn"
			default:
				return "ec2.amazonaws.com"
			}
		},
	}

	templateResource, err := NewTemplateResource("AWSIAMRolePolicy", RolePolicyTemplate, functions, nil)
	if err != nil {
		return nil, err
	}
	return templateResource, nil
}
<|endoftext|>


<commit_before>package task

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/golang/glog"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/zk"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
)

var (
	ErrBadConfig = errors.New("bad-config")
	ErrStopped   = errors.New("stopped")
)

type Runtime struct {
	Task

	zk zk.ZK

	status chan []byte
	stdout chan []byte
	stderr chan []byte
	stdin  chan []byte

	options interface{}
	done    bool
	ready   bool
	lock    sync.Mutex
	error   error

	stdoutBuff *bytes.Buffer
}

func (this *Task) Validate() error {
	switch {
	case !this.Info.Valid():
		return ErrBadConfig
	case !this.Success.Valid():
		return ErrBadConfig
	case !this.Error.Valid():
		return ErrBadConfig
	case this.Exec != nil:
		_, err := exec.LookPath(this.Exec.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Task) Init(zkc zk.ZK, options ...interface{}) (*Runtime, error) {
	if err := this.Validate(); err != nil {
		return nil, err
	}

	task := Runtime{
		Task: *this,
		zk:   zkc,
	}
	if len(options) > 0 {
		task.options = options[0]
	}
	task.status = make(chan []byte)

	if task.Task.Stdout != nil {
		task.stdout = make(chan []byte)
	}

	if task.Task.Stderr != nil {
		task.stderr = make(chan []byte)
	}

	now := time.Now()
	task.Stat.Started = &now

	if task.zk != nil {
		err := zk.CreateOrSet(task.zk, task.Info, task.Stat)
		if err != nil {
			return nil, err
		}
	}
	return &task, nil
}

func (this *Runtime) Stop() {
	this.lock.Lock()
	defer this.lock.Unlock()

	if this.done {
		return
	}

	if this.stdout != nil {
		this.stdout <- nil
	}
	if this.stderr != nil {
		this.stderr <- nil
	}
	this.Log("Stop")
	this.status <- nil

	this.done = true
}

func (this *Runtime) Stdin() io.Reader {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetReader(*this.Task.Stdin, c)
	} else {
		return nil
	}
}

func (this *Runtime) PublishStdin() io.Writer {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stdin, c)
	} else {
		glog.Warningln("Error getting stdin.", "Topic=", *this.Task.Stdin, "Err=", err)
		return nil
	}
}

func (this *Runtime) CaptureStdout() {
	this.stdoutBuff = new(bytes.Buffer)
}

func (this *Runtime) GetCapturedStdout() []byte {
	if this.stdoutBuff != nil {
		return this.stdoutBuff.Bytes()
	}
	return nil
}

func (this *Runtime) Stdout() io.Writer {
	var stdout io.Writer = os.Stdout
	if this.Task.Stdout != nil {
		if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
			stdout = pubsub.GetWriter(*this.Task.Stdout, c)
		} else {
			glog.Warningln("Error getting stdout.", "Topic=", *this.Task.Stdout, "Err=", err)
			return nil
		}
	}

	if this.stdoutBuff != nil {
		stdout = io.MultiWriter(stdout, this.stdoutBuff)
	}
	return stdout
}

func (this *Runtime) Stderr() io.Writer {
	if this.Task.Stderr == nil {
		return os.Stderr
	}
	if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stderr, c)
	} else {
		glog.Warningln("Error getting stderr.", "Topic=", *this.Task.Stderr, "Err=", err)
		return nil
	}
}

func (this *Runtime) Log(m ...string) {
	if this.done {
		return
	}
	source := ""
	_, file, line, ok := runtime.Caller(1)
	if ok {
		source = fmt.Sprintf("%s:%d", file, line)
	}

	s := strings.Join(m, " ")
	this.status <- []byte(s)
	glog.Infoln(source, m)
}

func (this *Runtime) Running() bool {
	return !this.done
}

func (this *Runtime) Start() (chan error, error) {
	_, _, err := this.start_streams()
	if err != nil {
		return nil, err
	}

	err = this.block_on_triggers()
	if err == zk.ErrTimeout {
		return nil, err
	}

	// Run the actual task
	if this.Task.Exec != nil {
		return this.exec()
	}
	return nil, nil
}

func (this *Runtime) block_on_triggers() error {
	if this.Start == nil {
		return nil
	}

	if this.Trigger == nil {
		return nil
	}

	// TODO - take into account ordering of cron vs registry.

	if this.Trigger.Registry != nil {
		trigger := zk.NewConditions(*this.Trigger.Registry, this.zk)
		// So now just block until the condition is true
		return trigger.Wait()
	}

	return nil
}

func (this *Runtime) exec() (chan error, error) {
	cmd := exec.Command(this.Exec.Path, this.Exec.Args...)
	cmd.Dir = this.Exec.Dir
	cmd.Env = this.Exec.Env

	// if this.Task.Stdin != nil {
	// 	stdin, err := cmd.StdinPipe()
	// 	if err != nil {
	// 		return nil, err
	// 	}
	// 	go func() {
	// 		io.Copy(io.MultiWriter(stdin, os.Stderr), this.Stdin())
	// 	}()
	// }

	if this.Task.Stdin != nil {
		sub, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options)
		if err != nil {
			return nil, err
		}
		stdin, err := sub.Subscribe(*this.Task.Stdin)
		if err != nil {
			return nil, err
		}
		wr, err := cmd.StdinPipe()
		if err != nil {
			return nil, err
		}
		go func() {
			// We need to do some special processing of input so that we can
			// terminate a session. Otherwise, this will just loop forever
			// because the pubsub topic will not go away -- even if it's a unique topic.
			for {
				m := <-stdin
				fmt.Printf(">> %s", string(m))
				switch {
				case bytes.Equal(m, []byte("@bye")):
					wr.Close()
					return
				default:
					wr.Write(m) // Need newline for shell to interpret
				}
			}
		}()
	}
	cmd.Stdout = this.Stdout()
	cmd.Stderr = this.Stderr()

	process_done := make(chan error)
	go func() {
		cmd.Start()

		// Wait for cmd to complete even if we have no more stdout/stderr
		if err := cmd.Wait(); err != nil {
			this.Error(err.Error())
			process_done <- err
			return
		}

		ps := cmd.ProcessState
		if ps == nil {
			this.Error(ErrCommandUnknown.Error())
			process_done <- ErrCommandUnknown
			return
		}

		glog.Infoln("Process pid=", ps.Pid(), "Exited=", ps.Exited(), "Success=", ps.Success())

		if !ps.Success() {
			this.Error(ErrExecFailed.Error())
			process_done <- ErrExecFailed
			return
		} else {
			this.Success(this.GetCapturedStdout())
			process_done <- nil
			return
		}
	}()

	return process_done, nil
}

func (this *Runtime) start_streams() (stdout, stderr chan<- []byte, err error) {
	this.lock.Lock()
	defer func() {
		this.error = err
		this.lock.Unlock()
	}()

	if this.error != nil {
		return nil, nil, this.error
	}

	if this.ready {
		return this.stdout, this.stderr, nil
	}

	if this.done {
		return nil, nil, ErrStopped
	}
	go func() {
		for {
			m := <-this.status
			if m == nil {
				break
			}
			if c, err := this.Task.Status.Broker().PubSub(this.Id, this.options); err == nil {
				c.Publish(this.Task.Status, m)
			} else {
				glog.Warningln("Cannot publish:", this.Task.Status.String(), "Err=", err)
			}
		}
	}()
	if this.stdout != nil {
		go func() {
			for {
				m := <-this.stdout
				if m == nil {
					break
				}
				if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stdout, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stdout.String(), "Err=", err)
				}

			}
		}()
		this.Log("Sending stdout to", this.Task.Stdout.Path())
	}
	if this.stderr != nil {
		go func() {
			for {
				m := <-this.stderr
				if m == nil {
					break
				}
				if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stderr, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stderr.String(), "Err=", err)
				}
			}
		}()
		this.Log("Sending stderr to", this.Task.Stderr.Path())
	}
	this.ready = true
	return this.stdout, this.stderr, nil
}

func (this *Runtime) Success(output interface{}) error {
	if this.done {
		return ErrStopped
	}

	switch output.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Success, output.([]byte))
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err := zk.CreateOrSetBytes(this.zk, *this.Task.Output, output.([]byte))
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Success, output.(string))
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err = zk.CreateOrSetString(this.zk, *this.Task.Output, output.(string))
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	default:
		value, err := json.Marshal(output)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Success, value)
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err = zk.CreateOrSetBytes(this.zk, *this.Task.Output, value)
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	}
	this.Log("Success", "Result written to", this.Task.Success.Path())

	now := time.Now()
	this.Stat.Success = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Success", "Completed")
	this.Stop()
	return nil
}

func (this *Runtime) Error(error interface{}) error {
	if this.done {
		return ErrStopped
	}
	switch error.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Error, error.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Error, error.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(error)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Error, value)
		if err != nil {
			return err
		}
	}
	this.Log("Error", "Error written to", this.Task.Error.Path())

	now := time.Now()
	this.Stat.Error = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Error", "Stop")
	this.Stop()
	return nil
}
<commit_msg>Minor state checks<commit_after>package task

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/golang/glog"
	"github.com/qorio/maestro/pkg/pubsub"
	"github.com/qorio/maestro/pkg/zk"
	"io"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"
)

var (
	ErrBadConfig = errors.New("bad-config")
	ErrStopped   = errors.New("stopped")
	ErrTimeout   = errors.New("timeout")
)

type Runtime struct {
	Task

	zk zk.ZK

	status chan []byte
	stdout chan []byte
	stderr chan []byte
	stdin  chan []byte

	options interface{}
	done    bool
	ready   bool
	lock    sync.Mutex
	error   error

	stdoutBuff *bytes.Buffer
}

func (this *Task) Copy() (*Task, error) {
	var buff bytes.Buffer
	enc := gob.NewEncoder(&buff)
	dec := gob.NewDecoder(&buff)
	err := enc.Encode(this)
	if err != nil {
		return nil, err
	}
	copy := new(Task)
	err = dec.Decode(copy)
	if err != nil {
		return nil, err
	}
	return copy, nil
}

func (this *Task) Validate() error {
	switch {
	case !this.Info.Valid():
		return ErrBadConfig
	case !this.Success.Valid():
		return ErrBadConfig
	case !this.Error.Valid():
		return ErrBadConfig
	case this.Exec != nil:
		_, err := exec.LookPath(this.Exec.Path)
		if err != nil {
			return err
		}
	}
	return nil
}

func (this *Task) Init(zkc zk.ZK, options ...interface{}) (*Runtime, error) {
	if err := this.Validate(); err != nil {
		return nil, err
	}

	task := Runtime{
		Task: *this,
		zk:   zkc,
	}
	if len(options) > 0 {
		task.options = options[0]
	}
	task.status = make(chan []byte)

	if task.Task.Stdout != nil {
		task.stdout = make(chan []byte)
	}

	if task.Task.Stderr != nil {
		task.stderr = make(chan []byte)
	}

	now := time.Now()
	task.Stat.Started = &now

	if task.zk != nil {
		err := zk.CreateOrSet(task.zk, task.Info, task.Stat)
		if err != nil {
			return nil, err
		}
	}
	return &task, nil
}

func (this *Runtime) Stop() {
	this.lock.Lock()
	defer this.lock.Unlock()

	if this.done {
		return
	}

	if this.stdout != nil {
		this.stdout <- nil
	}
	if this.stderr != nil {
		this.stderr <- nil
	}
	this.Log("Stop")
	this.status <- nil

	this.done = true
}

func (this *Runtime) Stdin() io.Reader {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetReader(*this.Task.Stdin, c)
	} else {
		return nil
	}
}

func (this *Runtime) PublishStdin() io.Writer {
	if this.Task.Stdin == nil {
		return os.Stdin
	}
	if c, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stdin, c)
	} else {
		glog.Warningln("Error getting stdin.", "Topic=", *this.Task.Stdin, "Err=", err)
		return nil
	}
}

func (this *Runtime) CaptureStdout() {
	this.stdoutBuff = new(bytes.Buffer)
}

func (this *Runtime) GetCapturedStdout() []byte {
	if this.stdoutBuff != nil {
		return this.stdoutBuff.Bytes()
	}
	return nil
}

func (this *Runtime) Stdout() io.Writer {
	var stdout io.Writer = os.Stdout
	if this.Task.Stdout != nil {
		if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
			stdout = pubsub.GetWriter(*this.Task.Stdout, c)
		} else {
			glog.Warningln("Error getting stdout.", "Topic=", *this.Task.Stdout, "Err=", err)
			return nil
		}
	}

	if this.stdoutBuff != nil {
		stdout = io.MultiWriter(stdout, this.stdoutBuff)
	}
	return stdout
}

func (this *Runtime) Stderr() io.Writer {
	if this.Task.Stderr == nil {
		return os.Stderr
	}
	if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
		return pubsub.GetWriter(*this.Task.Stderr, c)
	} else {
		glog.Warningln("Error getting stderr.", "Topic=", *this.Task.Stderr, "Err=", err)
		return nil
	}
}

func (this *Runtime) Log(m ...string) {
	if this.done {
		return
	}
	source := ""
	_, file, line, ok := runtime.Caller(1)
	if ok {
		source = fmt.Sprintf("%s:%d", file, line)
	}

	s := strings.Join(m, " ")
	this.status <- []byte(s)
	glog.Infoln(source, m)
}

func (this *Runtime) Running() bool {
	return !this.done
}

func (this *Runtime) ApplyEnvAndFuncs(env map[string]interface{}, funcs map[string]interface{}) error {
	if this.Task.Exec == nil {
		return nil
	}

	this.lock.Lock()
	defer this.lock.Unlock()

	applied, err := this.Task.Exec.ApplySubstitutions(env, funcs)
	if err != nil {
		return err
	}
	this.Task.Exec = applied
	return nil
}

func (this *Runtime) Start() (chan error, error) {
	_, _, err := this.start_streams()
	if err != nil {
		return nil, err
	}

	err = this.block_on_triggers()
	if err == zk.ErrTimeout {
		return nil, ErrTimeout
	}

	// Run the actual task
	if this.Task.Exec != nil {
		return this.exec()
	}
	return nil, nil
}

func (this *Runtime) block_on_triggers() error {
	if this.Start == nil {
		return nil
	}

	if this.Trigger == nil {
		return nil
	}

	// TODO - take into account ordering of cron vs registry.

	if this.Trigger.Registry != nil {
		trigger := zk.NewConditions(*this.Trigger.Registry, this.zk)
		// So now just block until the condition is true
		return trigger.Wait()
	}

	return nil
}

func (this *Runtime) exec() (chan error, error) {
	cmd := exec.Command(this.Exec.Path, this.Exec.Args...)
	cmd.Dir = this.Exec.Dir
	cmd.Env = this.Exec.Env

	if this.Task.Stdin != nil {
		sub, err := this.Task.Stdin.Broker().PubSub(this.Id, this.options)
		if err != nil {
			return nil, err
		}
		stdin, err := sub.Subscribe(*this.Task.Stdin)
		if err != nil {
			return nil, err
		}
		wr, err := cmd.StdinPipe()
		if err != nil {
			return nil, err
		}
		go func() {
			// We need to do some special processing of input so that we can
			// terminate a session. Otherwise, this will just loop forever
			// because the pubsub topic will not go away -- even if it's a unique topic.
			for {
				m := <-stdin
				fmt.Printf(">> %s", string(m))
				switch {
				case bytes.Equal(m, []byte("@bye")):
					wr.Close()
					return
				default:
					wr.Write(m) // Need newline for shell to interpret
				}
			}
		}()
	}
	cmd.Stdout = this.Stdout()
	cmd.Stderr = this.Stderr()

	process_done := make(chan error)
	go func() {
		cmd.Start()

		// Wait for cmd to complete even if we have no more stdout/stderr
		if err := cmd.Wait(); err != nil {
			this.Error(err.Error())
			process_done <- err
			return
		}

		ps := cmd.ProcessState
		if ps == nil {
			this.Error(ErrCommandUnknown.Error())
			process_done <- ErrCommandUnknown
			return
		}

		glog.Infoln("Process pid=", ps.Pid(), "Exited=", ps.Exited(), "Success=", ps.Success())

		if !ps.Success() {
			this.Error(ErrExecFailed.Error())
			process_done <- ErrExecFailed
			return
		} else {
			this.Success(this.GetCapturedStdout())
			process_done <- nil
			return
		}
	}()

	return process_done, nil
}

func (this *Runtime) start_streams() (stdout, stderr chan<- []byte, err error) {
	this.lock.Lock()
	defer func() {
		this.error = err
		this.lock.Unlock()
	}()

	if this.error != nil {
		return nil, nil, this.error
	}

	if this.ready {
		return this.stdout, this.stderr, nil
	}

	if this.done {
		return nil, nil, ErrStopped
	}
	go func() {
		for {
			m := <-this.status
			if m == nil {
				break
			}
			if c, err := this.Task.Status.Broker().PubSub(this.Id, this.options); err == nil {
				c.Publish(this.Task.Status, m)
			} else {
				glog.Warningln("Cannot publish:", this.Task.Status.String(), "Err=", err)
			}
		}
	}()
	if this.stdout != nil {
		go func() {
			for {
				m := <-this.stdout
				if m == nil {
					break
				}
				if c, err := this.Task.Stdout.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stdout, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stdout.String(), "Err=", err)
				}

			}
		}()
		this.Log("Sending stdout to", this.Task.Stdout.Path())
	}
	if this.stderr != nil {
		go func() {
			for {
				m := <-this.stderr
				if m == nil {
					break
				}
				if c, err := this.Task.Stderr.Broker().PubSub(this.Id, this.options); err == nil {
					c.Publish(*this.Task.Stderr, m)
				} else {
					glog.Warningln("Cannot publish:", this.Task.Stderr.String(), "Err=", err)
				}
			}
		}()
		this.Log("Sending stderr to", this.Task.Stderr.Path())
	}
	this.ready = true
	return this.stdout, this.stderr, nil
}

func (this *Runtime) Success(output interface{}) error {
	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		this.Stop()
		return nil
	}

	if this.done {
		return ErrStopped
	}

	switch output.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Success, output.([]byte))
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err := zk.CreateOrSetBytes(this.zk, *this.Task.Output, output.([]byte))
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Success, output.(string))
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err = zk.CreateOrSetString(this.zk, *this.Task.Output, output.(string))
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	default:
		value, err := json.Marshal(output)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Success, value)
		if err != nil {
			return err
		}
		// copy the data over
		if this.Task.Output != nil {
			err = zk.CreateOrSetBytes(this.zk, *this.Task.Output, value)
			if err != nil {
				return err
			}
			this.Log("Success", "Result written to", this.Task.Output.Path())
		}
	}
	this.Log("Success", "Result written to", this.Task.Success.Path())

	now := time.Now()
	this.Stat.Success = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Success", "Completed")
	this.Stop()
	return nil
}

func (this *Runtime) Error(error interface{}) error {
	if this.zk == nil {
		glog.Infoln("Not connected to zk.  Output not recorded")
		this.Stop()
		return nil
	}

	if this.done {
		return ErrStopped
	}
	switch error.(type) {
	case []byte:
		err := zk.CreateOrSetBytes(this.zk, this.Task.Error, error.([]byte))
		if err != nil {
			return err
		}
	case string:
		err := zk.CreateOrSetString(this.zk, this.Task.Error, error.(string))
		if err != nil {
			return err
		}
	default:
		value, err := json.Marshal(error)
		if err != nil {
			return err
		}
		err = zk.CreateOrSetBytes(this.zk, this.Task.Error, value)
		if err != nil {
			return err
		}
	}
	this.Log("Error", "Error written to", this.Task.Error.Path())

	now := time.Now()
	this.Stat.Error = &now
	err := zk.CreateOrSet(this.zk, this.Info, this.Stat)
	if err != nil {
		return err
	}

	this.Log("Error", "Stop")
	this.Stop()
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	addr         = flag.String("addr", "localhost:7702", "The address on which netprint listens.")
	tcp          = flag.Bool("tcp", false, "Accept raw TCP requests instead of HTTP.")
	udp          = flag.Bool("udp", false, "Accept raw UDP packets instead of HTTP.")
	delayFlag    = flag.String("delay", "0ms", "How long to delay before responding (HTTP only).")
	responseCode = flag.Int("response-code", http.StatusOK, "Response code for HTTP requests.")
	responseText = flag.String("response-text", "", "Response body for HTTP requests.")
	mode         = modeHTTP
	delay        time.Duration
	mut          = &sync.Mutex{}
)

func fatal(args ...interface{}) {
	fmt.Println(args...)
	os.Exit(1)
}

func fatalf(format string, args ...interface{}) {
	fmt.Printf(format, args...)
	os.Exit(1)
}

type modeType int

const (
	modeHTTP modeType = iota
	modeTCP
	modeUDP
)

// copyRecordNewline is like io.Copy, but says whether the copied data ended in a newline.
func copyRecordNewline(dst io.Writer, src io.Reader) (written int64, err error, endingNewline bool) {
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
			endingNewline = buf[nr-1] == '\n'
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err, endingNewline
}

func handleHTTP(w http.ResponseWriter, r *http.Request) {
	// It's pretty dumb to synchronize all the HTTP handlers when net/http is doing such a good job of
	// multiplexing requests onto goroutines for me, but this approach is simpler than constructing a
	// non-concurrent HTTP server.
	mut.Lock()
	defer mut.Unlock()

	fmt.Printf(">>>>> Request: %s\n", r.URL)
	written, err, endingNewline := copyRecordNewline(os.Stdout, r.Body)
	if err != nil {
		return
	}
	if written == 0 {
		fmt.Println("(Empty body.)")
	} else if !endingNewline {
		fmt.Println()
	}

	time.Sleep(delay)

	w.WriteHeader(*responseCode)
	w.Write([]byte(*responseText))
}

func runHTTP() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/", handleHTTP)
	fmt.Println("Now accepting HTTP requests on", *addr)
	return http.ListenAndServe(*addr, mux)
}

func handleTCP(conn net.Conn) {
	fmt.Printf(">>>>> %s connected.\n", conn.RemoteAddr())
	written, err, endingNewline := copyRecordNewline(os.Stdout, conn)
	if err != nil {
		return
	}
	if written == 0 {
		fmt.Println("(No data transmitted.)")
	} else if !endingNewline {
		fmt.Println()
	}
	fmt.Printf(">>>>> %s disconnected.\n", conn.RemoteAddr())
}

func runTCP() error {
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		return err
	}
	fmt.Println("Now accepting raw TCP requests on", *addr)
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		handleTCP(conn)
	}
}

func handleUDP(conn *net.UDPConn) {
	buf := make([]byte, 10*1024) // 10KB buffer to handle pretty damn big UDP packets
	for {
		n, addr, err := conn.ReadFromUDP(buf)
		if err != nil {
			return
		}
		fmt.Printf(">>>>> Received a packet from %s\n", addr)
		if n == 0 {
			fmt.Println("(No data transmitted.)")
		} else {
			os.Stdout.Write(buf[:n])
			if buf[n-1] != '\n' {
				fmt.Println()
			}
		}
	}
}

func runUDP() error {
	u, err := net.ResolveUDPAddr("udp", *addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", u)
	if err != nil {
		return err
	}
	fmt.Println("Now accepting raw UDP requests on", *addr)
	handleUDP(conn)
	return nil
}

func main() {
	flag.Parse()
	if *tcp {
		if *udp {
			fatal("Cannot specify both -tcp and -udp.")
		}
		mode = modeTCP
	}
	if *udp {
		mode = modeUDP
	}

	if *delayFlag != "" {
		var err error
		delay, err = time.ParseDuration(*delayFlag)
		if err != nil {
			fatalf("Bad -delay value: %s", err)
		}
	}

	setFlags := make(map[string]bool)
	flag.Visit(func(f *flag.Flag) { setFlags[f.Name] = true })

	for _, f := range []string{"response-code", "response-text", "delay"} {
		if setFlags[f] && mode != modeHTTP {
			fatalf("Cannot specify -%s except in HTTP mode.", f)
		}
	}
	if setFlags["response-code"] && (*responseCode < 100 || *responseCode >= 600) {
		fatal("Invalid HTTP response code:", *responseCode)
	}
	switch mode {
	case modeHTTP:
		fatal(runHTTP())
	case modeTCP:
		fatal(runTCP())
	case modeUDP:
		fatal(runUDP())
	}
}
<commit_msg>Use a duration flag<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	addr         = flag.String("addr", "localhost:7702", "The address on which netprint listens.")
	tcp          = flag.Bool("tcp", false, "Accept raw TCP requests instead of HTTP.")
	udp          = flag.Bool("udp", false, "Accept raw UDP packets instead of HTTP.")
	delay        = flag.Duration("delay", 0, "How long to delay before responding (HTTP only).")
	responseCode = flag.Int("response-code", http.StatusOK, "Response code for HTTP requests.")
	responseText = flag.String("response-text", "", "Response body for HTTP requests.")
	mode         = modeHTTP
	mut          = &sync.Mutex{}
)

func fatal(args ...interface{}) {
	fmt.Println(args...)
	os.Exit(1)
}

func fatalf(format string, args ...interface{}) {
	fmt.Printf(format, args...)
	os.Exit(1)
}

type modeType int

const (
	modeHTTP modeType = iota
	modeTCP
	modeUDP
)

// copyRecordNewline is like io.Copy, but says whether the copied data ended in a newline.
func copyRecordNewline(dst io.Writer, src io.Reader) (written int64, err error, endingNewline bool) {
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
			endingNewline = buf[nr-1] == '\n'
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err, endingNewline
}

func handleHTTP(w http.ResponseWriter, r *http.Request) {
	// It's pretty dumb to synchronize all the HTTP handlers when net/http is doing such a good job of
	// multiplexing requests onto goroutines for me, but this approach is simpler than constructing a
	// non-concurrent HTTP server.
	mut.Lock()
	defer mut.Unlock()

	fmt.Printf(">>>>> Request: %s\n", r.URL)
	written, err, endingNewline := copyRecordNewline(os.Stdout, r.Body)
	if err != nil {
		return
	}
	if written == 0 {
		fmt.Println("(Empty body.)")
	} else if !endingNewline {
		fmt.Println()
	}

	time.Sleep(*delay)

	w.WriteHeader(*responseCode)
	w.Write([]byte(*responseText))
}

func runHTTP() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/", handleHTTP)
	fmt.Println("Now accepting HTTP requests on", *addr)
	return http.ListenAndServe(*addr, mux)
}

func handleTCP(conn net.Conn) {
	fmt.Printf(">>>>> %s connected.\n", conn.RemoteAddr())
	written, err, endingNewline := copyRecordNewline(os.Stdout, conn)
	if err != nil {
		return
	}
	if written == 0 {
		fmt.Println("(No data transmitted.)")
	} else if !endingNewline {
		fmt.Println()
	}
	fmt.Printf(">>>>> %s disconnected.\n", conn.RemoteAddr())
}

func runTCP() error {
	l, err := net.Listen("tcp", *addr)
	if err != nil {
		return err
	}
	fmt.Println("Now accepting raw TCP requests on", *addr)
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		handleTCP(conn)
	}
}

func handleUDP(conn *net.UDPConn) {
	buf := make([]byte, 10*1024) // 10KB buffer to handle pretty damn big UDP packets
	for {
		n, addr, err := conn.ReadFromUDP(buf)
		if err != nil {
			return
		}
		fmt.Printf(">>>>> Received a packet from %s\n", addr)
		if n == 0 {
			fmt.Println("(No data transmitted.)")
		} else {
			os.Stdout.Write(buf[:n])
			if buf[n-1] != '\n' {
				fmt.Println()
			}
		}
	}
}

func runUDP() error {
	u, err := net.ResolveUDPAddr("udp", *addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", u)
	if err != nil {
		return err
	}
	fmt.Println("Now accepting raw UDP requests on", *addr)
	handleUDP(conn)
	return nil
}

func main() {
	flag.Parse()
	if *tcp {
		if *udp {
			fatal("Cannot specify both -tcp and -udp.")
		}
		mode = modeTCP
	}
	if *udp {
		mode = modeUDP
	}

	setFlags := make(map[string]bool)
	flag.Visit(func(f *flag.Flag) { setFlags[f.Name] = true })

	for _, f := range []string{"response-code", "response-text", "delay"} {
		if setFlags[f] && mode != modeHTTP {
			fatalf("Cannot specify -%s except in HTTP mode.", f)
		}
	}
	if setFlags["response-code"] && (*responseCode < 100 || *responseCode >= 600) {
		fatal("Invalid HTTP response code:", *responseCode)
	}
	switch mode {
	case modeHTTP:
		fatal(runHTTP())
	case modeTCP:
		fatal(runTCP())
	case modeUDP:
		fatal(runUDP())
	}
}
<|endoftext|>


<commit_before>package commands

import (
	boshlog "bosh/logger"
	boshsys "bosh/system"
	fakesys "bosh/system/fakes"
	"github.com/stretchr/testify/assert"
	"os"
	"path/filepath"
	"testing"
)

func TestCompressFilesInDir(t *testing.T) {
	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	srcDir := fixtureSrcDir(t)
	tgzName, err := dc.CompressFilesInDir(srcDir)
	assert.NoError(t, err)

	defer os.Remove(tgzName)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	_, _, err = cmdRunner.RunCommand("tar", "--no-same-owner", "-xzpf", tgzName, "-C", dstDir)
	assert.NoError(t, err)

	tarDirStat, err := os.Stat(dstDir)
	assert.NoError(t, err)
	assert.Equal(t, os.FileMode(0755), tarDirStat.Mode().Perm())

	// regular files
	content, err := fs.ReadFile(dstDir + "/app.stdout.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is app stdout")

	content, err = fs.ReadFile(dstDir + "/app.stderr.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is app stderr")

	// file in a directory
	content, err = fs.ReadFile(dstDir + "/other_logs/other_app.stdout.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is other app stdout")
}

func TestDecompressFileToDir(t *testing.T) {
	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	err := dc.DecompressFileToDir(fixtureSrcTgz(t), dstDir)
	assert.NoError(t, err)

	// regular files
	content, err := fs.ReadFile(dstDir + "/not-nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "not-nested-file")

	// nested directory with a file
	content, err = fs.ReadFile(dstDir + "/dir/nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "nested-file")

	// nested directory with a file inside another directory
	content, err = fs.ReadFile(dstDir + "/dir/nested-dir/double-nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "double-nested-file")

	// directory without a file (empty)
	content, err = fs.ReadFile(dstDir + "/empty-dir")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "is a directory")

	// nested directory without a file (empty) inside another directory
	content, err = fs.ReadFile(dstDir + "/dir/empty-nested-dir")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "is a directory")
}

func TestDecompressFileToDirReturnsError(t *testing.T) {
	nonExistentDstDir := filepath.Join(os.TempDir(), "TestDecompressFileToDirReturnsError")

	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	// propagates errors raised when untarring
	err := dc.DecompressFileToDir(fixtureSrcTgz(t), nonExistentDstDir)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), nonExistentDstDir) // path in error
}

func TestDecompressFileToDirUsesNoSameOwnerOption(t *testing.T) {
	fs, _ := getCompressorDependencies()
	cmdRunner := fakesys.NewFakeCmdRunner()
	dc := NewTarballCompressor(cmdRunner, fs)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	tarballPath := fixtureSrcTgz(t)
	err := dc.DecompressFileToDir(tarballPath, dstDir)
	assert.NoError(t, err)

	assert.Equal(t, 1, len(cmdRunner.RunCommands))
	assert.Equal(t, []string{
		"tar", "--no-same-owner",
		"-xzvf", tarballPath,
		"-C", dstDir,
	}, cmdRunner.RunCommands[0])
}

func createdTmpDir(t *testing.T, fs boshsys.FileSystem) string {
	dstDir := filepath.Join(os.TempDir(), "TestCompressor")
	err := fs.MkdirAll(dstDir, os.ModePerm)
	assert.NoError(t, err)
	return dstDir
}

func fixtureSrcDir(t *testing.T) string {
	pwd, err := os.Getwd()
	assert.NoError(t, err)
	return filepath.Join(pwd, "..", "..", "..", "..", "fixtures", "test_filtered_copy_to_temp")
}

func fixtureSrcTgz(t *testing.T) string {
	pwd, err := os.Getwd()
	assert.NoError(t, err)
	return filepath.Join(pwd, "..", "..", "..", "..", "fixtures", "compressor-decompress-file-to-dir.tgz")
}

func getCompressorDependencies() (boshsys.FileSystem, boshsys.CmdRunner) {
	logger := boshlog.NewLogger(boshlog.LEVEL_NONE)
	cmdRunner := boshsys.NewExecCmdRunner(logger)
	fs := boshsys.NewOsFileSystem(logger, cmdRunner)
	return fs, cmdRunner
}
<commit_msg>Removed invalid test since copier was extracted from tar compressor<commit_after>package commands

import (
	boshlog "bosh/logger"
	boshsys "bosh/system"
	fakesys "bosh/system/fakes"
	"github.com/stretchr/testify/assert"
	"os"
	"path/filepath"
	"testing"
)

func TestCompressFilesInDir(t *testing.T) {
	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	srcDir := fixtureSrcDir(t)
	tgzName, err := dc.CompressFilesInDir(srcDir)
	assert.NoError(t, err)

	defer os.Remove(tgzName)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	_, _, err = cmdRunner.RunCommand("tar", "--no-same-owner", "-xzpf", tgzName, "-C", dstDir)
	assert.NoError(t, err)

	// regular files
	content, err := fs.ReadFile(dstDir + "/app.stdout.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is app stdout")

	content, err = fs.ReadFile(dstDir + "/app.stderr.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is app stderr")

	// file in a directory
	content, err = fs.ReadFile(dstDir + "/other_logs/other_app.stdout.log")
	assert.NoError(t, err)
	assert.Contains(t, content, "this is other app stdout")
}

func TestDecompressFileToDir(t *testing.T) {
	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	err := dc.DecompressFileToDir(fixtureSrcTgz(t), dstDir)
	assert.NoError(t, err)

	// regular files
	content, err := fs.ReadFile(dstDir + "/not-nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "not-nested-file")

	// nested directory with a file
	content, err = fs.ReadFile(dstDir + "/dir/nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "nested-file")

	// nested directory with a file inside another directory
	content, err = fs.ReadFile(dstDir + "/dir/nested-dir/double-nested-file")
	assert.NoError(t, err)
	assert.Contains(t, content, "double-nested-file")

	// directory without a file (empty)
	content, err = fs.ReadFile(dstDir + "/empty-dir")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "is a directory")

	// nested directory without a file (empty) inside another directory
	content, err = fs.ReadFile(dstDir + "/dir/empty-nested-dir")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "is a directory")
}

func TestDecompressFileToDirReturnsError(t *testing.T) {
	nonExistentDstDir := filepath.Join(os.TempDir(), "TestDecompressFileToDirReturnsError")

	fs, cmdRunner := getCompressorDependencies()
	dc := NewTarballCompressor(cmdRunner, fs)

	// propagates errors raised when untarring
	err := dc.DecompressFileToDir(fixtureSrcTgz(t), nonExistentDstDir)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), nonExistentDstDir) // path in error
}

func TestDecompressFileToDirUsesNoSameOwnerOption(t *testing.T) {
	fs, _ := getCompressorDependencies()
	cmdRunner := fakesys.NewFakeCmdRunner()
	dc := NewTarballCompressor(cmdRunner, fs)

	dstDir := createdTmpDir(t, fs)
	defer os.RemoveAll(dstDir)

	tarballPath := fixtureSrcTgz(t)
	err := dc.DecompressFileToDir(tarballPath, dstDir)
	assert.NoError(t, err)

	assert.Equal(t, 1, len(cmdRunner.RunCommands))
	assert.Equal(t, []string{
		"tar", "--no-same-owner",
		"-xzvf", tarballPath,
		"-C", dstDir,
	}, cmdRunner.RunCommands[0])
}

func createdTmpDir(t *testing.T, fs boshsys.FileSystem) string {
	dstDir := filepath.Join(os.TempDir(), "TestCompressor")
	err := fs.MkdirAll(dstDir, os.ModePerm)
	assert.NoError(t, err)
	return dstDir
}

func fixtureSrcDir(t *testing.T) string {
	pwd, err := os.Getwd()
	assert.NoError(t, err)
	return filepath.Join(pwd, "..", "..", "..", "..", "fixtures", "test_filtered_copy_to_temp")
}

func fixtureSrcTgz(t *testing.T) string {
	pwd, err := os.Getwd()
	assert.NoError(t, err)
	return filepath.Join(pwd, "..", "..", "..", "..", "fixtures", "compressor-decompress-file-to-dir.tgz")
}

func getCompressorDependencies() (boshsys.FileSystem, boshsys.CmdRunner) {
	logger := boshlog.NewLogger(boshlog.LEVEL_NONE)
	cmdRunner := boshsys.NewExecCmdRunner(logger)
	fs := boshsys.NewOsFileSystem(logger, cmdRunner)
	return fs, cmdRunner
}
<|endoftext|>