<commit_before>/*
Copyright (c) 2014-2016 VMware, Inc. 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 flags

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/kr/pretty"
	"github.com/vmware/govmomi/task"
	"github.com/vmware/govmomi/vim25/progress"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
	"github.com/vmware/govmomi/vim25/xml"
)

type OutputWriter interface {
	Write(io.Writer) error
}

type OutputFlag struct {
	common

	JSON bool
	XML  bool
	TTY  bool
	Dump bool
	Out  io.Writer

	formatError  bool
	formatIndent bool
}

var outputFlagKey = flagKey("output")

func NewOutputFlag(ctx context.Context) (*OutputFlag, context.Context) {
	if v := ctx.Value(outputFlagKey); v != nil {
		return v.(*OutputFlag), ctx
	}

	v := &OutputFlag{Out: os.Stdout}
	ctx = context.WithValue(ctx, outputFlagKey, v)
	return v, ctx
}

func (flag *OutputFlag) Register(ctx context.Context, f *flag.FlagSet) {
	flag.RegisterOnce(func() {
		f.BoolVar(&flag.JSON, "json", false, "Enable JSON output")
		f.BoolVar(&flag.XML, "xml", false, "Enable XML output")
		f.BoolVar(&flag.Dump, "dump", false, "Enable Go output")
		// Avoid adding more flags for now..
		flag.formatIndent = os.Getenv("GOVC_INDENT") != "false"      // Default to indented output
		flag.formatError = os.Getenv("GOVC_FORMAT_ERROR") != "false" // Default to formatted errors
	})
}

func (flag *OutputFlag) Process(ctx context.Context) error {
	return flag.ProcessOnce(func() error {
		if !flag.All() {
			// Assume we have a tty if not outputting JSON
			flag.TTY = true
		}

		return nil
	})
}

// Log outputs the specified string, prefixed with the current time.
// A newline is not automatically added. If the specified string
// starts with a '\r', the current line is cleared first.
func (flag *OutputFlag) Log(s string) (int, error) {
	if len(s) > 0 && s[0] == '\r' {
		flag.Write([]byte{'\r', 033, '[', 'K'})
		s = s[1:]
	}

	return flag.WriteString(time.Now().Format("[02-01-06 15:04:05] ") + s)
}

func (flag *OutputFlag) Write(b []byte) (int, error) {
	if !flag.TTY {
		return 0, nil
	}

	n, err := os.Stdout.Write(b)
	os.Stdout.Sync()
	return n, err
}

func (flag *OutputFlag) WriteString(s string) (int, error) {
	return flag.Write([]byte(s))
}

func (flag *OutputFlag) All() bool {
	return flag.JSON || flag.XML || flag.Dump
}

func dumpValue(val interface{}) interface{} {
	type dumper interface {
		Dump() interface{}
	}

	if d, ok := val.(dumper); ok {
		return d.Dump()
	}

	rval := reflect.ValueOf(val)
	if rval.Type().Kind() != reflect.Ptr {
		return val
	}

	rval = rval.Elem()
	if rval.Type().Kind() == reflect.Struct {
		f := rval.Field(0)
		if f.Type().Kind() == reflect.Slice {
			// common case for the various 'type infoResult'
			if f.Len() == 1 {
				return f.Index(0).Interface()
			}
			return f.Interface()
		}

		if rval.NumField() == 1 && rval.Type().Field(0).Anonymous {
			// common case where govc type wraps govmomi type to implement OutputWriter
			return f.Interface()
		}
	}

	return val
}

func (flag *OutputFlag) WriteResult(result OutputWriter) error {
	var err error

	switch {
	case flag.Dump:
		format := "%#v\n"
		if flag.formatIndent {
			format = "%# v\n"
		}
		_, err = pretty.Fprintf(flag.Out, format, dumpValue(result))
	case flag.JSON:
		e := json.NewEncoder(flag.Out)
		if flag.formatIndent {
			e.SetIndent("", "  ")
		}
		err = e.Encode(result)
	case flag.XML:
		e := xml.NewEncoder(flag.Out)
		if flag.formatIndent {
			e.Indent("", "  ")
		}
		err = e.Encode(dumpValue(result))
		if err == nil {
			fmt.Fprintln(flag.Out)
		}
	default:
		err = result.Write(flag.Out)
	}

	return err
}

func (flag *OutputFlag) WriteError(err error) bool {
	if flag.formatError {
		flag.Out = os.Stderr
		return flag.WriteResult(&errorOutput{err}) == nil
	}
	return false
}

type errorOutput struct {
	error
}

func (e errorOutput) Write(w io.Writer) error {
	reason := e.error.Error()
	var messages []string
	var faults []types.LocalizableMessage

	switch err := e.error.(type) {
	case task.Error:
		faults = err.LocalizedMethodFault.Fault.GetMethodFault().FaultMessage
		if err.Description != nil {
			reason = fmt.Sprintf("%s (%s)", reason, err.Description.Message)
		}
	default:
		if soap.IsSoapFault(err) {
			detail := soap.ToSoapFault(err).Detail.Fault
			if f, ok := detail.(types.BaseMethodFault); ok {
				faults = f.GetMethodFault().FaultMessage
			}
		}
	}

	for _, m := range faults {
		if m.Message != "" && !strings.HasPrefix(m.Message, "[context]") {
			messages = append(messages, fmt.Sprintf("%s (%s)", m.Message, m.Key))
		}
	}

	messages = append(messages, reason)

	for _, message := range messages {
                if _, err := fmt.Fprintf(w, "%s: %s\n", os.Args[0], message); err != nil {
                        return err
                }
	}

	return nil
}

func (e errorOutput) Dump() interface{} {
	if f, ok := e.error.(task.Error); ok {
		return f.LocalizedMethodFault
	}
	if soap.IsSoapFault(e.error) {
		return soap.ToSoapFault(e.error)
	}
	if soap.IsVimFault(e.error) {
		return soap.ToVimFault(e.error)
	}
	return e
}

func (e errorOutput) canEncode() bool {
	switch e.error.(type) {
	case task.Error:
		return true
	}
	return soap.IsSoapFault(e.error) || soap.IsVimFault(e.error)
}

// cannotEncode causes cli.Run to output err.Error() as it would without an error format specified
var cannotEncode = errors.New("cannot encode error")

func (e errorOutput) MarshalJSON() ([]byte, error) {
	_, ok := e.error.(json.Marshaler)
	if ok || e.canEncode() {
		return json.Marshal(e.error)
	}
	return nil, cannotEncode
}

func (e errorOutput) MarshalXML(encoder *xml.Encoder, start xml.StartElement) error {
	_, ok := e.error.(xml.Marshaler)
	if ok || e.canEncode() {
		return encoder.Encode(e.error)
	}
	return cannotEncode
}

type progressLogger struct {
	flag   *OutputFlag
	prefix string

	wg sync.WaitGroup

	sink chan chan progress.Report
	done chan struct{}
}

func newProgressLogger(flag *OutputFlag, prefix string) *progressLogger {
	p := &progressLogger{
		flag:   flag,
		prefix: prefix,

		sink: make(chan chan progress.Report),
		done: make(chan struct{}),
	}

	p.wg.Add(1)

	go p.loopA()

	return p
}

// loopA runs before Sink() has been called.
func (p *progressLogger) loopA() {
	var err error

	defer p.wg.Done()

	tick := time.NewTicker(100 * time.Millisecond)
	defer tick.Stop()

	called := false

	for stop := false; !stop; {
		select {
		case ch := <-p.sink:
			err = p.loopB(tick, ch)
			stop = true
			called = true
		case <-p.done:
			stop = true
		case <-tick.C:
			line := fmt.Sprintf("\r%s", p.prefix)
			p.flag.Log(line)
		}
	}

	if err != nil && err != io.EOF {
		p.flag.Log(fmt.Sprintf("\r%sError: %s\n", p.prefix, err))
	} else if called {
		p.flag.Log(fmt.Sprintf("\r%sOK\n", p.prefix))
	}
}

// loopA runs after Sink() has been called.
func (p *progressLogger) loopB(tick *time.Ticker, ch <-chan progress.Report) error {
	var r progress.Report
	var ok bool
	var err error

	for ok = true; ok; {
		select {
		case r, ok = <-ch:
			if !ok {
				break
			}
			err = r.Error()
		case <-tick.C:
			line := fmt.Sprintf("\r%s", p.prefix)
			if r != nil {
				line += fmt.Sprintf("(%.0f%%", r.Percentage())
				detail := r.Detail()
				if detail != "" {
					line += fmt.Sprintf(", %s", detail)
				}
				line += ")"
			}
			p.flag.Log(line)
		}
	}

	return err
}

func (p *progressLogger) Sink() chan<- progress.Report {
	ch := make(chan progress.Report)
	p.sink <- ch
	return ch
}

func (p *progressLogger) Wait() {
	close(p.done)
	p.wg.Wait()
}

func (flag *OutputFlag) ProgressLogger(prefix string) *progressLogger {
	return newProgressLogger(flag, prefix)
}
<commit_msg>gofmt<commit_after>/*
Copyright (c) 2014-2016 VMware, Inc. 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 flags

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/kr/pretty"
	"github.com/vmware/govmomi/task"
	"github.com/vmware/govmomi/vim25/progress"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
	"github.com/vmware/govmomi/vim25/xml"
)

type OutputWriter interface {
	Write(io.Writer) error
}

type OutputFlag struct {
	common

	JSON bool
	XML  bool
	TTY  bool
	Dump bool
	Out  io.Writer

	formatError  bool
	formatIndent bool
}

var outputFlagKey = flagKey("output")

func NewOutputFlag(ctx context.Context) (*OutputFlag, context.Context) {
	if v := ctx.Value(outputFlagKey); v != nil {
		return v.(*OutputFlag), ctx
	}

	v := &OutputFlag{Out: os.Stdout}
	ctx = context.WithValue(ctx, outputFlagKey, v)
	return v, ctx
}

func (flag *OutputFlag) Register(ctx context.Context, f *flag.FlagSet) {
	flag.RegisterOnce(func() {
		f.BoolVar(&flag.JSON, "json", false, "Enable JSON output")
		f.BoolVar(&flag.XML, "xml", false, "Enable XML output")
		f.BoolVar(&flag.Dump, "dump", false, "Enable Go output")
		// Avoid adding more flags for now..
		flag.formatIndent = os.Getenv("GOVC_INDENT") != "false"      // Default to indented output
		flag.formatError = os.Getenv("GOVC_FORMAT_ERROR") != "false" // Default to formatted errors
	})
}

func (flag *OutputFlag) Process(ctx context.Context) error {
	return flag.ProcessOnce(func() error {
		if !flag.All() {
			// Assume we have a tty if not outputting JSON
			flag.TTY = true
		}

		return nil
	})
}

// Log outputs the specified string, prefixed with the current time.
// A newline is not automatically added. If the specified string
// starts with a '\r', the current line is cleared first.
func (flag *OutputFlag) Log(s string) (int, error) {
	if len(s) > 0 && s[0] == '\r' {
		flag.Write([]byte{'\r', 033, '[', 'K'})
		s = s[1:]
	}

	return flag.WriteString(time.Now().Format("[02-01-06 15:04:05] ") + s)
}

func (flag *OutputFlag) Write(b []byte) (int, error) {
	if !flag.TTY {
		return 0, nil
	}

	n, err := os.Stdout.Write(b)
	os.Stdout.Sync()
	return n, err
}

func (flag *OutputFlag) WriteString(s string) (int, error) {
	return flag.Write([]byte(s))
}

func (flag *OutputFlag) All() bool {
	return flag.JSON || flag.XML || flag.Dump
}

func dumpValue(val interface{}) interface{} {
	type dumper interface {
		Dump() interface{}
	}

	if d, ok := val.(dumper); ok {
		return d.Dump()
	}

	rval := reflect.ValueOf(val)
	if rval.Type().Kind() != reflect.Ptr {
		return val
	}

	rval = rval.Elem()
	if rval.Type().Kind() == reflect.Struct {
		f := rval.Field(0)
		if f.Type().Kind() == reflect.Slice {
			// common case for the various 'type infoResult'
			if f.Len() == 1 {
				return f.Index(0).Interface()
			}
			return f.Interface()
		}

		if rval.NumField() == 1 && rval.Type().Field(0).Anonymous {
			// common case where govc type wraps govmomi type to implement OutputWriter
			return f.Interface()
		}
	}

	return val
}

func (flag *OutputFlag) WriteResult(result OutputWriter) error {
	var err error

	switch {
	case flag.Dump:
		format := "%#v\n"
		if flag.formatIndent {
			format = "%# v\n"
		}
		_, err = pretty.Fprintf(flag.Out, format, dumpValue(result))
	case flag.JSON:
		e := json.NewEncoder(flag.Out)
		if flag.formatIndent {
			e.SetIndent("", "  ")
		}
		err = e.Encode(result)
	case flag.XML:
		e := xml.NewEncoder(flag.Out)
		if flag.formatIndent {
			e.Indent("", "  ")
		}
		err = e.Encode(dumpValue(result))
		if err == nil {
			fmt.Fprintln(flag.Out)
		}
	default:
		err = result.Write(flag.Out)
	}

	return err
}

func (flag *OutputFlag) WriteError(err error) bool {
	if flag.formatError {
		flag.Out = os.Stderr
		return flag.WriteResult(&errorOutput{err}) == nil
	}
	return false
}

type errorOutput struct {
	error
}

func (e errorOutput) Write(w io.Writer) error {
	reason := e.error.Error()
	var messages []string
	var faults []types.LocalizableMessage

	switch err := e.error.(type) {
	case task.Error:
		faults = err.LocalizedMethodFault.Fault.GetMethodFault().FaultMessage
		if err.Description != nil {
			reason = fmt.Sprintf("%s (%s)", reason, err.Description.Message)
		}
	default:
		if soap.IsSoapFault(err) {
			detail := soap.ToSoapFault(err).Detail.Fault
			if f, ok := detail.(types.BaseMethodFault); ok {
				faults = f.GetMethodFault().FaultMessage
			}
		}
	}

	for _, m := range faults {
		if m.Message != "" && !strings.HasPrefix(m.Message, "[context]") {
			messages = append(messages, fmt.Sprintf("%s (%s)", m.Message, m.Key))
		}
	}

	messages = append(messages, reason)

	for _, message := range messages {
		if _, err := fmt.Fprintf(w, "%s: %s\n", os.Args[0], message); err != nil {
			return err
		}
	}

	return nil
}

func (e errorOutput) Dump() interface{} {
	if f, ok := e.error.(task.Error); ok {
		return f.LocalizedMethodFault
	}
	if soap.IsSoapFault(e.error) {
		return soap.ToSoapFault(e.error)
	}
	if soap.IsVimFault(e.error) {
		return soap.ToVimFault(e.error)
	}
	return e
}

func (e errorOutput) canEncode() bool {
	switch e.error.(type) {
	case task.Error:
		return true
	}
	return soap.IsSoapFault(e.error) || soap.IsVimFault(e.error)
}

// cannotEncode causes cli.Run to output err.Error() as it would without an error format specified
var cannotEncode = errors.New("cannot encode error")

func (e errorOutput) MarshalJSON() ([]byte, error) {
	_, ok := e.error.(json.Marshaler)
	if ok || e.canEncode() {
		return json.Marshal(e.error)
	}
	return nil, cannotEncode
}

func (e errorOutput) MarshalXML(encoder *xml.Encoder, start xml.StartElement) error {
	_, ok := e.error.(xml.Marshaler)
	if ok || e.canEncode() {
		return encoder.Encode(e.error)
	}
	return cannotEncode
}

type progressLogger struct {
	flag   *OutputFlag
	prefix string

	wg sync.WaitGroup

	sink chan chan progress.Report
	done chan struct{}
}

func newProgressLogger(flag *OutputFlag, prefix string) *progressLogger {
	p := &progressLogger{
		flag:   flag,
		prefix: prefix,

		sink: make(chan chan progress.Report),
		done: make(chan struct{}),
	}

	p.wg.Add(1)

	go p.loopA()

	return p
}

// loopA runs before Sink() has been called.
func (p *progressLogger) loopA() {
	var err error

	defer p.wg.Done()

	tick := time.NewTicker(100 * time.Millisecond)
	defer tick.Stop()

	called := false

	for stop := false; !stop; {
		select {
		case ch := <-p.sink:
			err = p.loopB(tick, ch)
			stop = true
			called = true
		case <-p.done:
			stop = true
		case <-tick.C:
			line := fmt.Sprintf("\r%s", p.prefix)
			p.flag.Log(line)
		}
	}

	if err != nil && err != io.EOF {
		p.flag.Log(fmt.Sprintf("\r%sError: %s\n", p.prefix, err))
	} else if called {
		p.flag.Log(fmt.Sprintf("\r%sOK\n", p.prefix))
	}
}

// loopA runs after Sink() has been called.
func (p *progressLogger) loopB(tick *time.Ticker, ch <-chan progress.Report) error {
	var r progress.Report
	var ok bool
	var err error

	for ok = true; ok; {
		select {
		case r, ok = <-ch:
			if !ok {
				break
			}
			err = r.Error()
		case <-tick.C:
			line := fmt.Sprintf("\r%s", p.prefix)
			if r != nil {
				line += fmt.Sprintf("(%.0f%%", r.Percentage())
				detail := r.Detail()
				if detail != "" {
					line += fmt.Sprintf(", %s", detail)
				}
				line += ")"
			}
			p.flag.Log(line)
		}
	}

	return err
}

func (p *progressLogger) Sink() chan<- progress.Report {
	ch := make(chan progress.Report)
	p.sink <- ch
	return ch
}

func (p *progressLogger) Wait() {
	close(p.done)
	p.wg.Wait()
}

func (flag *OutputFlag) ProgressLogger(prefix string) *progressLogger {
	return newProgressLogger(flag, prefix)
}
<|endoftext|>


<commit_before>package runner

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden/client"
	"code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const MNT_DETACH = 0x2

var DataDir string
var TarBin = os.Getenv("GARDEN_TAR_PATH")

type RunningGarden struct {
	client.Client

	runner  GardenRunner
	process ifrit.Process

	debugIP   string
	debugPort int

	Pid int

	Tmpdir string

	DepotDir  string
	DataDir   string
	GraphPath string

	logger lager.Logger
}

type GardenRunner struct {
	*ginkgomon.Runner
	Cmd            *exec.Cmd
	TmpDir         string
	GraphPath      string
	ConsoleSockets string
	DepotDir       string
	DebugIp        string
	DebugPort      int
	Network, Addr  string
}

func init() {
	DataDir = os.Getenv("GARDEN_TEST_GRAPHPATH")
	if DataDir == "" {
		// This must be set outside of the Ginkgo node directory (tmpDir) because
		// otherwise the Concourse worker may run into one of the AUFS kernel
		// module bugs that cause the VM to become unresponsive.
		DataDir = "/tmp/aufs_mount"
	}
}

func NewGardenRunner(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin, network, address string, user *syscall.Credential, argv ...string) GardenRunner {
	r := GardenRunner{}

	r.Network = network
	r.Addr = address
	r.TmpDir = filepath.Join(
		os.TempDir(),
		fmt.Sprintf("test-garden-%d", ginkgo.GinkgoParallelNode()),
	)

	r.GraphPath = filepath.Join(DataDir, fmt.Sprintf("node-%d", ginkgo.GinkgoParallelNode()))
	r.DepotDir = filepath.Join(r.TmpDir, "containers")
	r.ConsoleSockets = filepath.Join(r.TmpDir, "console-sockets")

	MustMountTmpfs(r.GraphPath)

	r.Cmd = cmd(r.TmpDir, r.DepotDir, r.GraphPath, r.ConsoleSockets, r.Network, r.Addr, bin, initBin, nstarBin, dadooBin, grootfsBin, tarBin, rootfs, user, argv...)
	r.Cmd.Env = append(os.Environ(), fmt.Sprintf("TMPDIR=%s", r.TmpDir))

	for i, arg := range r.Cmd.Args {
		if arg == "--debug-bind-ip" {
			r.DebugIp = r.Cmd.Args[i+1]
		}
		if arg == "--debug-bind-port" {
			r.DebugPort, _ = strconv.Atoi(r.Cmd.Args[i+1])
		}
	}

	r.Runner = ginkgomon.New(ginkgomon.Config{
		Name:              "guardian",
		Command:           r.Cmd,
		AnsiColorCode:     "31m",
		StartCheck:        "guardian.started",
		StartCheckTimeout: 30 * time.Second,
	})

	return r
}

func Start(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin string, user *syscall.Credential, argv ...string) *RunningGarden {
	runner := NewGardenRunner(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin, "unix", fmt.Sprintf("/tmp/garden_%d.sock", GinkgoParallelNode()), user, argv...)

	r := &RunningGarden{
		runner:   runner,
		DepotDir: runner.DepotDir,

		DataDir:   DataDir,
		GraphPath: runner.GraphPath,
		Tmpdir:    runner.TmpDir,
		logger:    lagertest.NewTestLogger("garden-runner"),

		debugIP:   runner.DebugIp,
		debugPort: runner.DebugPort,

		Client: client.New(connection.New(runner.Network, runner.Addr)),
	}

	r.process = ifrit.Invoke(r.runner)
	r.Pid = runner.Cmd.Process.Pid

	return r
}

func (r *RunningGarden) Kill() error {
	r.process.Signal(syscall.SIGKILL)
	select {
	case err := <-r.process.Wait():
		return err
	case <-time.After(time.Second * 10):
		r.process.Signal(syscall.SIGKILL)
		return errors.New("timed out waiting for garden to shutdown after 10 seconds")
	}
}

func (r *RunningGarden) DestroyAndStop() error {
	if err := r.DestroyContainers(); err != nil {
		return err
	}

	if err := r.Stop(); err != nil {
		return err
	}

	return nil
}

func (r *RunningGarden) Stop() error {
	r.process.Signal(syscall.SIGTERM)

	var err error
	for i := 0; i < 5; i++ {
		select {
		case err := <-r.process.Wait():
			return err
		case <-time.After(time.Second * 5):
			r.process.Signal(syscall.SIGTERM)
			err = errors.New("timed out waiting for garden to shutdown after 5 seconds")
		}
	}

	r.process.Signal(syscall.SIGKILL)
	return err
}

func cmd(tmpdir, depotDir, graphPath, consoleSocketsPath, network, addr, bin, initBin, nstarBin, dadooBin, grootfsBin, tarBin, rootfs string, user *syscall.Credential, argv ...string) *exec.Cmd {
	Expect(os.MkdirAll(tmpdir, 0755)).To(Succeed())
	Expect(os.MkdirAll(depotDir, 0755)).To(Succeed())

	appendDefaultFlag := func(ar []string, key, value string) []string {
		for _, a := range argv {
			if a == key {
				return ar
			}
		}

		if value != "" {
			return append(ar, key, value)
		} else {
			return append(ar, key)
		}
	}

	gardenArgs := make([]string, len(argv))
	copy(gardenArgs, argv)

	switch network {
	case "tcp":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-ip", strings.Split(addr, ":")[0])
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-port", strings.Split(addr, ":")[1])
	case "unix":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-socket", addr)
	}

	if rootfs != "" {
		gardenArgs = appendDefaultFlag(gardenArgs, "--default-rootfs", rootfs)
	}

	gardenArgs = appendDefaultFlag(gardenArgs, "--depot", depotDir)
	gardenArgs = appendDefaultFlag(gardenArgs, "--graph", graphPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--console-sockets-path", consoleSocketsPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tag", fmt.Sprintf("%d", GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--network-pool", fmt.Sprintf("10.254.%d.0/22", 4*GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--init-bin", initBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--dadoo-bin", dadooBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--nstar-bin", nstarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tar-bin", tarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--port-pool-start", fmt.Sprintf("%d", GinkgoParallelNode()*7000))

	cmd := exec.Command(bin, append([]string{"server"}, gardenArgs...)...)
	if user != nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
		cmd.SysProcAttr.Credential = user

		uidGid := fmt.Sprintf("%d:%d", user.Uid, user.Gid)
		Expect(exec.Command("chown", "-R", uidGid, tmpdir).Run()).To(Succeed())
	}

	return cmd
}

func (r *RunningGarden) Cleanup() {
	// unmount aufs since the docker graph driver leaves this around,
	// otherwise the following commands might fail
	retry := retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil)

	err := retry.Run(func() error {
		if err := os.RemoveAll(path.Join(r.GraphPath, "aufs")); err == nil {
			return nil // if we can remove it, it's already unmounted
		}

		if err := syscall.Unmount(path.Join(r.GraphPath, "aufs"), MNT_DETACH); err != nil {
			r.logger.Error("failed-unmount-attempt", err)
			return err
		}

		return nil
	})

	if err != nil {
		r.logger.Error("failed-to-unmount", err)
	}

	MustUnmountTmpfs(r.GraphPath)

	// In the kernel version 3.19.0-51-generic the code bellow results in
	// hanging the running VM. We are not deleting the node-X directories. They
	// are empty and the next test will re-use them. We will stick with that
	// workaround until we can test on a newer kernel that will hopefully not
	// have this bug.
	//
	// if err := os.RemoveAll(r.GraphPath); err != nil {
	// 	r.logger.Error("remove-graph", err)
	// }

	r.logger.Info("cleanup-tempdirs")
	if err := os.RemoveAll(r.Tmpdir); err != nil {
		r.logger.Error("cleanup-tempdirs-failed", err, lager.Data{"tmpdir": r.Tmpdir})
	} else {
		r.logger.Info("tempdirs-removed")
	}
}

func (r *RunningGarden) DestroyContainers() error {
	containers, err := r.Containers(nil)
	if err != nil {
		return err
	}

	for _, container := range containers {
		r.Destroy(container.Handle())
	}

	return nil
}

type debugVars struct {
	NumGoRoutines int `json:"numGoRoutines"`
}

func (r *RunningGarden) NumGoroutines() (int, error) {
	debugURL := fmt.Sprintf("http://%s:%d/debug/vars", r.debugIP, r.debugPort)
	res, err := http.Get(debugURL)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	decoder := json.NewDecoder(res.Body)
	var debugVarsData debugVars
	err = decoder.Decode(&debugVarsData)
	if err != nil {
		return 0, err
	}

	return debugVarsData.NumGoRoutines, nil
}

func (r *RunningGarden) Buffer() *gbytes.Buffer {
	return r.runner.Buffer()
}

func (r *RunningGarden) ExitCode() int {
	return r.runner.ExitCode()
}
<commit_msg>Add logging around flake<commit_after>package runner

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden/client"
	"code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const MNT_DETACH = 0x2

var DataDir string
var TarBin = os.Getenv("GARDEN_TAR_PATH")

type RunningGarden struct {
	client.Client

	runner  GardenRunner
	process ifrit.Process

	debugIP   string
	debugPort int

	Pid int

	Tmpdir string

	DepotDir  string
	DataDir   string
	GraphPath string

	logger lager.Logger
}

type GardenRunner struct {
	*ginkgomon.Runner
	Cmd            *exec.Cmd
	TmpDir         string
	GraphPath      string
	ConsoleSockets string
	DepotDir       string
	DebugIp        string
	DebugPort      int
	Network, Addr  string
}

func init() {
	DataDir = os.Getenv("GARDEN_TEST_GRAPHPATH")
	if DataDir == "" {
		// This must be set outside of the Ginkgo node directory (tmpDir) because
		// otherwise the Concourse worker may run into one of the AUFS kernel
		// module bugs that cause the VM to become unresponsive.
		DataDir = "/tmp/aufs_mount"
	}
}

func NewGardenRunner(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin, network, address string, user *syscall.Credential, argv ...string) GardenRunner {
	r := GardenRunner{}

	r.Network = network
	r.Addr = address
	r.TmpDir = filepath.Join(
		os.TempDir(),
		fmt.Sprintf("test-garden-%d", ginkgo.GinkgoParallelNode()),
	)

	r.GraphPath = filepath.Join(DataDir, fmt.Sprintf("node-%d", ginkgo.GinkgoParallelNode()))
	r.DepotDir = filepath.Join(r.TmpDir, "containers")
	r.ConsoleSockets = filepath.Join(r.TmpDir, "console-sockets")

	MustMountTmpfs(r.GraphPath)

	r.Cmd = cmd(r.TmpDir, r.DepotDir, r.GraphPath, r.ConsoleSockets, r.Network, r.Addr, bin, initBin, nstarBin, dadooBin, grootfsBin, tarBin, rootfs, user, argv...)
	r.Cmd.Env = append(os.Environ(), fmt.Sprintf("TMPDIR=%s", r.TmpDir))

	for i, arg := range r.Cmd.Args {
		if arg == "--debug-bind-ip" {
			r.DebugIp = r.Cmd.Args[i+1]
		}
		if arg == "--debug-bind-port" {
			r.DebugPort, _ = strconv.Atoi(r.Cmd.Args[i+1])
		}
	}

	r.Runner = ginkgomon.New(ginkgomon.Config{
		Name:              "guardian",
		Command:           r.Cmd,
		AnsiColorCode:     "31m",
		StartCheck:        "guardian.started",
		StartCheckTimeout: 30 * time.Second,
	})

	return r
}

func Start(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin string, user *syscall.Credential, argv ...string) *RunningGarden {
	runner := NewGardenRunner(bin, initBin, nstarBin, dadooBin, grootfsBin, rootfs, tarBin, "unix", fmt.Sprintf("/tmp/garden_%d.sock", GinkgoParallelNode()), user, argv...)

	r := &RunningGarden{
		runner:   runner,
		DepotDir: runner.DepotDir,

		DataDir:   DataDir,
		GraphPath: runner.GraphPath,
		Tmpdir:    runner.TmpDir,
		logger:    lagertest.NewTestLogger("garden-runner"),

		debugIP:   runner.DebugIp,
		debugPort: runner.DebugPort,

		Client: client.New(connection.New(runner.Network, runner.Addr)),
	}

	r.process = ifrit.Invoke(r.runner)
	r.Pid = runner.Cmd.Process.Pid

	return r
}

func (r *RunningGarden) Kill() error {
	r.process.Signal(syscall.SIGKILL)
	select {
	case err := <-r.process.Wait():
		return err
	case <-time.After(time.Second * 10):
		r.process.Signal(syscall.SIGKILL)
		return errors.New("timed out waiting for garden to shutdown after 10 seconds")
	}
}

func (r *RunningGarden) DestroyAndStop() error {
	if err := r.DestroyContainers(); err != nil {
		return err
	}

	if err := r.Stop(); err != nil {
		return err
	}

	return nil
}

func (r *RunningGarden) Stop() error {
	r.process.Signal(syscall.SIGTERM)

	var err error
	for i := 0; i < 5; i++ {
		select {
		case err := <-r.process.Wait():
			return err
		case <-time.After(time.Second * 5):
			r.process.Signal(syscall.SIGTERM)
			err = errors.New("timed out waiting for garden to shutdown after 5 seconds")
		}
	}

	r.process.Signal(syscall.SIGKILL)
	return err
}

func cmd(tmpdir, depotDir, graphPath, consoleSocketsPath, network, addr, bin, initBin, nstarBin, dadooBin, grootfsBin, tarBin, rootfs string, user *syscall.Credential, argv ...string) *exec.Cmd {
	Expect(os.MkdirAll(tmpdir, 0755)).To(Succeed())
	Expect(os.MkdirAll(depotDir, 0755)).To(Succeed())

	appendDefaultFlag := func(ar []string, key, value string) []string {
		for _, a := range argv {
			if a == key {
				return ar
			}
		}

		if value != "" {
			return append(ar, key, value)
		} else {
			return append(ar, key)
		}
	}

	gardenArgs := make([]string, len(argv))
	copy(gardenArgs, argv)

	switch network {
	case "tcp":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-ip", strings.Split(addr, ":")[0])
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-port", strings.Split(addr, ":")[1])
	case "unix":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-socket", addr)
	}

	if rootfs != "" {
		gardenArgs = appendDefaultFlag(gardenArgs, "--default-rootfs", rootfs)
	}

	gardenArgs = appendDefaultFlag(gardenArgs, "--depot", depotDir)
	gardenArgs = appendDefaultFlag(gardenArgs, "--graph", graphPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--console-sockets-path", consoleSocketsPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tag", fmt.Sprintf("%d", GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--network-pool", fmt.Sprintf("10.254.%d.0/22", 4*GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--init-bin", initBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--dadoo-bin", dadooBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--nstar-bin", nstarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tar-bin", tarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--port-pool-start", fmt.Sprintf("%d", GinkgoParallelNode()*7000))

	cmd := exec.Command(bin, append([]string{"server"}, gardenArgs...)...)
	if user != nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{}
		cmd.SysProcAttr.Credential = user

		uidGid := fmt.Sprintf("%d:%d", user.Uid, user.Gid)
		out, err := exec.Command("chown", "-R", uidGid, tmpdir).CombinedOutput()
		Expect(err).NotTo(HaveOccurred(),
			fmt.Sprintf("Execing 'chown -R %s %s' yielded the following output : %s", uidGid, tmpdir, string(out)),
		)
	}

	return cmd
}

func (r *RunningGarden) Cleanup() {
	// unmount aufs since the docker graph driver leaves this around,
	// otherwise the following commands might fail
	retry := retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil)

	err := retry.Run(func() error {
		if err := os.RemoveAll(path.Join(r.GraphPath, "aufs")); err == nil {
			return nil // if we can remove it, it's already unmounted
		}

		if err := syscall.Unmount(path.Join(r.GraphPath, "aufs"), MNT_DETACH); err != nil {
			r.logger.Error("failed-unmount-attempt", err)
			return err
		}

		return nil
	})

	if err != nil {
		r.logger.Error("failed-to-unmount", err)
	}

	MustUnmountTmpfs(r.GraphPath)

	// In the kernel version 3.19.0-51-generic the code bellow results in
	// hanging the running VM. We are not deleting the node-X directories. They
	// are empty and the next test will re-use them. We will stick with that
	// workaround until we can test on a newer kernel that will hopefully not
	// have this bug.
	//
	// if err := os.RemoveAll(r.GraphPath); err != nil {
	// 	r.logger.Error("remove-graph", err)
	// }

	r.logger.Info("cleanup-tempdirs")
	if err := os.RemoveAll(r.Tmpdir); err != nil {
		r.logger.Error("cleanup-tempdirs-failed", err, lager.Data{"tmpdir": r.Tmpdir})
	} else {
		r.logger.Info("tempdirs-removed")
	}
}

func (r *RunningGarden) DestroyContainers() error {
	containers, err := r.Containers(nil)
	if err != nil {
		return err
	}

	for _, container := range containers {
		r.Destroy(container.Handle())
	}

	return nil
}

type debugVars struct {
	NumGoRoutines int `json:"numGoRoutines"`
}

func (r *RunningGarden) NumGoroutines() (int, error) {
	debugURL := fmt.Sprintf("http://%s:%d/debug/vars", r.debugIP, r.debugPort)
	res, err := http.Get(debugURL)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	decoder := json.NewDecoder(res.Body)
	var debugVarsData debugVars
	err = decoder.Decode(&debugVarsData)
	if err != nil {
		return 0, err
	}

	return debugVarsData.NumGoRoutines, nil
}

func (r *RunningGarden) Buffer() *gbytes.Buffer {
	return r.runner.Buffer()
}

func (r *RunningGarden) ExitCode() int {
	return r.runner.ExitCode()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
  "time"
	"encoding/binary"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket"
	"flag"
)

var NETFLOW_V5_HEADER_SIZE int = 24;
var NETFLOW_V5_RECORD_SIZE int = 48;
var PROTOCOL_TCP uint8 = 6
var PROTOCOL_UDP uint8 = 17
var NETFLOW_PORT int = 2055
var NANOSECOND int64 = 1000000000

func construct_ethernet() *layers.Ethernet {
	return &layers.Ethernet{}
}

func construct_ip(srcaddr string, dstaddr string) *layers.IPv4 {
	return &layers.IPv4{
		SrcIP: net.ParseIP(srcaddr),
		DstIP: net.ParseIP(dstaddr),
	}
}

func construct_udp() *layers.UDP {
	return &layers.UDP{}
}

type NETFLOW_v5_header struct {
	Version			uint16
	Count			uint16
	Sys_uptime		uint32
	Unix_secs		uint32
	Unix_nsecs		uint32
	Flow_sequence		uint32
	Engine_type		uint8
	Engine_id		uint8
	Sampling_interval	uint16
}

type NETFLOW_v5_record struct {
	Srcaddr		uint32
	Dstaddr		uint32
	Nexthop		uint32
	Input		uint16
	Output		uint16
	DPkts		uint32
	DOctets		uint32
	First		uint32
	Last		uint32
	Srcport		uint16
	Dstport		uint16
	Pad1		uint8
	Tcp_flags	uint8
	Prot		uint8
	Tos		uint8
	Src_as  	uint16
	Dst_as  	uint16
	Src_mask 	uint8
	Dst_mask	uint8
	Pad2		uint16
}

func v4_to_uint32(addr net.IP) uint32 {
	var ret uint32;
	ret |= uint32(addr[0])
	ret |= uint32(addr[1]) << 8
	ret |= uint32(addr[2]) << 16
	ret |= uint32(addr[3]) << 24
	return ret
}

func construct_v5_header(count uint16, sampling uint16) NETFLOW_v5_header {
	header := NETFLOW_v5_header{
		Version:		5,		//Netflow v5
		Count:			count,		//Number of records in this packet
		Sys_uptime:		0,		//Ignore for now
		Unix_secs:		0,		//Ignore for now
		Unix_nsecs:		0,		//Ignore for now
		Flow_sequence:		0,		//Ignore for now. Eventually want to track sequence numbers
		Engine_type:		0,		//Ignore for now
		Engine_id:		0,		//Ignore for now
		Sampling_interval:	sampling,	//TODO
	}
	return header
}

func insert_v5_header(header NETFLOW_v5_header, buf []byte, offset int) int {
        binary.BigEndian.PutUint16(buf[offset:], header.Version)
        binary.BigEndian.PutUint16(buf[offset + 2:], header.Count)
        binary.BigEndian.PutUint32(buf[offset + 4:], header.Sys_uptime)
        binary.BigEndian.PutUint32(buf[offset + 8:], header.Unix_secs)
        binary.BigEndian.PutUint32(buf[offset + 12:], header.Unix_nsecs)
        binary.BigEndian.PutUint32(buf[offset + 16:], header.Flow_sequence)
        buf[offset + 20] = header.Engine_type
        buf[offset + 21] = header.Engine_id
        binary.BigEndian.PutUint16(buf[offset + 22:], header.Sampling_interval)

	return NETFLOW_V5_HEADER_SIZE
}

func construct_v5_record(srcaddr string, dstaddr string,
	pkts uint32, l3_bytes uint32, srcport uint16, dstport uint16,
	protocol uint8, src_as uint16, dst_as uint16) NETFLOW_v5_record {

	srcip := v4_to_uint32(net.ParseIP(srcaddr))
	dstip := v4_to_uint32(net.ParseIP(dstaddr))

	record := NETFLOW_v5_record {
		Srcaddr:		srcip,
		Dstaddr:		dstip,
		Nexthop:		0,				//Ignore for now
		Input:			0,				//Do something with this later
		Output:			0,				//^^
		DPkts:			pkts,
		DOctets:		l3_bytes,
		First:			0,				//Ignore for now
		Last:			0,				//Ignore for now
		Srcport:		srcport,
		Dstport:		dstport,
		Pad1:			0,
		Tcp_flags:		0,				//Something with this later
		Prot:			PROTOCOL_TCP,
		Tos:			0,
		Src_as:			src_as,
		Dst_as:			dst_as,
		Src_mask:		0,
		Dst_mask:		0,
	}
	return record
}

func insert_v5_record(record NETFLOW_v5_record, buf []byte, offset int) int {
        binary.BigEndian.PutUint32(buf[offset:], record.Srcaddr)
        binary.BigEndian.PutUint32(buf[offset + 4:], record.Dstaddr)
        binary.BigEndian.PutUint32(buf[offset + 8:], record.Nexthop)
        binary.BigEndian.PutUint16(buf[offset + 12:], record.Input)
        binary.BigEndian.PutUint16(buf[offset + 14:], record.Output)
        binary.BigEndian.PutUint32(buf[offset + 16:], record.DPkts)
        binary.BigEndian.PutUint32(buf[offset + 20:], record.DOctets)
        binary.BigEndian.PutUint32(buf[offset + 24:], record.First)
        binary.BigEndian.PutUint32(buf[offset + 28:], record.Last)
        binary.BigEndian.PutUint16(buf[offset + 32:], record.Srcport)
        binary.BigEndian.PutUint16(buf[offset + 34:], record.Dstport)
        buf[offset + 36] = record.Pad1
        buf[offset + 37] = record.Tcp_flags
        buf[offset + 38] = record.Prot
        buf[offset + 39] = record.Tos
        binary.BigEndian.PutUint16(buf[offset + 40:], record.Src_as)
        binary.BigEndian.PutUint16(buf[offset + 42:], record.Dst_as)
        buf[offset + 44] = record.Src_mask
        buf[offset + 45] = record.Dst_mask
        binary.BigEndian.PutUint16(buf[offset + 46:], record.Pad2)
	return NETFLOW_V5_RECORD_SIZE;
}

func construct_payload(num_records uint16) gopacket.Payload {

	buf := gopacket.NewSerializeBuffer()
//        payload := buf.Bytes()
        //Allocate the space we will need for the header
        bytes,err := buf.PrependBytes(NETFLOW_V5_HEADER_SIZE + NETFLOW_V5_RECORD_SIZE*int(num_records))
	if err != nil {
		return nil
	}

	offset := 0

	header := construct_v5_header(num_records, 1000)
	offset += insert_v5_header(header, bytes, offset)

	var record NETFLOW_v5_record;
	for i := 0; i < int(num_records); i++ {
		record = construct_v5_record("1.1.1.1", "2.2.2.2", 5, 256, 80, 5050, 6, 237, 237)
		insert_v5_record(record, bytes, offset)

	}

	return gopacket.Payload(bytes)
}

func chk(err error){
	if err != nil {
		panic(err)
	}
}

func init_connection(addr net.IP) *net.UDPConn {
	conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: addr, Port: 0})
	chk(err)
	return conn
}

func send_packet(conn *net.UDPConn, addr net.IP, port int, pkt []byte) {
	_, err := conn.WriteToUDP(pkt, &net.UDPAddr{IP: addr, Port: port})
	chk(err)
}

func main() {

	dst_ip := flag.String("dst", "127.0.0.1", "Destination IP to send the spoofed netflow")
	dst_port := flag.Int("port", NETFLOW_PORT, "Destination Port to send the spoofed netflow")
	rate := flag.Int64("rate", 1, "Rate in Packets/s")
  runtime := flag.Int64("time", 10, "Time in seconds to send packets")
  flows_per_packet := flag.Uint("fpp", 1, "flows per packet, max of 30")
  flag.Parse()

	dst_addr := net.ParseIP(*dst_ip)

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{}

	l2 := construct_ethernet()
	l3 := construct_ip("1.2.3.4", "5.6.7.8")
	l4 := construct_udp()

	payload := construct_payload(uint16(*flows_per_packet))
	//LayerCake
	gopacket.SerializeLayers(buf, opts,
		l2,
		l3,
		l4,
		payload)
	packetData := buf.Bytes()

	//Send the packet to lo
	conn := init_connection(dst_addr)
	send_packet(conn, dst_addr, *dst_port, packetData)

  //Simple way for now. Token based approach later
  throttle := time.Tick(time.Duration(*rate)*time.Second)
  for i := 0; int64(i) < (*rate)*(*runtime); i++ {
    <-throttle
    go send_packet(conn, dst_addr, *dst_port, packetData)
  }
	fmt.Println("fin")
}
<commit_msg>Add actually logging/log levels<commit_after>package main

import (
	"net"
  "io"
  "time"
  "log"
  "os"
	"encoding/binary"
	"code.google.com/p/gopacket/layers"
	"code.google.com/p/gopacket"
	"flag"
)
var (
  NETFLOW_V5_HEADER_SIZE int = 24;
  NETFLOW_V5_RECORD_SIZE int = 48;
  PROTOCOL_TCP uint8 = 6
  PROTOCOL_UDP uint8 = 17
  NETFLOW_PORT int = 2055
  NANOSECOND int64 = 1000000000
  Trace *log.Logger
  Info *log.Logger
  Warning *log.Logger
  Error *log.Logger
)

func construct_ethernet() *layers.Ethernet {
	return &layers.Ethernet{}
}

func construct_ip(srcaddr string, dstaddr string) *layers.IPv4 {
	return &layers.IPv4{
		SrcIP: net.ParseIP(srcaddr),
		DstIP: net.ParseIP(dstaddr),
	}
}

func construct_udp() *layers.UDP {
	return &layers.UDP{}
}

type NETFLOW_v5_header struct {
	Version			uint16
	Count			uint16
	Sys_uptime		uint32
	Unix_secs		uint32
	Unix_nsecs		uint32
	Flow_sequence		uint32
	Engine_type		uint8
	Engine_id		uint8
	Sampling_interval	uint16
}

type NETFLOW_v5_record struct {
	Srcaddr		uint32
	Dstaddr		uint32
	Nexthop		uint32
	Input		uint16
	Output		uint16
	DPkts		uint32
	DOctets		uint32
	First		uint32
	Last		uint32
	Srcport		uint16
	Dstport		uint16
	Pad1		uint8
	Tcp_flags	uint8
	Prot		uint8
	Tos		uint8
	Src_as  	uint16
	Dst_as  	uint16
	Src_mask 	uint8
	Dst_mask	uint8
	Pad2		uint16
}

func v4_to_uint32(addr net.IP) uint32 {
	var ret uint32;
	ret |= uint32(addr[0])
	ret |= uint32(addr[1]) << 8
	ret |= uint32(addr[2]) << 16
	ret |= uint32(addr[3]) << 24
	return ret
}

func construct_v5_header(count uint16, sampling uint16) NETFLOW_v5_header {
	header := NETFLOW_v5_header{
		Version:		5,		//Netflow v5
		Count:			count,		//Number of records in this packet
		Sys_uptime:		0,		//Ignore for now
		Unix_secs:		0,		//Ignore for now
		Unix_nsecs:		0,		//Ignore for now
		Flow_sequence:		0,		//Ignore for now. Eventually want to track sequence numbers
		Engine_type:		0,		//Ignore for now
		Engine_id:		0,		//Ignore for now
		Sampling_interval:	sampling,	//TODO
	}
	return header
}

func insert_v5_header(header NETFLOW_v5_header, buf []byte, offset int) int {
        binary.BigEndian.PutUint16(buf[offset:], header.Version)
        binary.BigEndian.PutUint16(buf[offset + 2:], header.Count)
        binary.BigEndian.PutUint32(buf[offset + 4:], header.Sys_uptime)
        binary.BigEndian.PutUint32(buf[offset + 8:], header.Unix_secs)
        binary.BigEndian.PutUint32(buf[offset + 12:], header.Unix_nsecs)
        binary.BigEndian.PutUint32(buf[offset + 16:], header.Flow_sequence)
        buf[offset + 20] = header.Engine_type
        buf[offset + 21] = header.Engine_id
        binary.BigEndian.PutUint16(buf[offset + 22:], header.Sampling_interval)

	return NETFLOW_V5_HEADER_SIZE
}

func construct_v5_record(srcaddr string, dstaddr string,
	pkts uint32, l3_bytes uint32, srcport uint16, dstport uint16,
	protocol uint8, src_as uint16, dst_as uint16) NETFLOW_v5_record {

	srcip := v4_to_uint32(net.ParseIP(srcaddr))
	dstip := v4_to_uint32(net.ParseIP(dstaddr))

	record := NETFLOW_v5_record {
		Srcaddr:		srcip,
		Dstaddr:		dstip,
		Nexthop:		0,				//Ignore for now
		Input:			0,				//Do something with this later
		Output:			0,				//^^
		DPkts:			pkts,
		DOctets:		l3_bytes,
		First:			0,				//Ignore for now
		Last:			0,				//Ignore for now
		Srcport:		srcport,
		Dstport:		dstport,
		Pad1:			0,
		Tcp_flags:		0,				//Something with this later
		Prot:			PROTOCOL_TCP,
		Tos:			0,
		Src_as:			src_as,
		Dst_as:			dst_as,
		Src_mask:		0,
		Dst_mask:		0,
	}
	return record
}

func insert_v5_record(record NETFLOW_v5_record, buf []byte, offset int) int {
        binary.BigEndian.PutUint32(buf[offset:], record.Srcaddr)
        binary.BigEndian.PutUint32(buf[offset + 4:], record.Dstaddr)
        binary.BigEndian.PutUint32(buf[offset + 8:], record.Nexthop)
        binary.BigEndian.PutUint16(buf[offset + 12:], record.Input)
        binary.BigEndian.PutUint16(buf[offset + 14:], record.Output)
        binary.BigEndian.PutUint32(buf[offset + 16:], record.DPkts)
        binary.BigEndian.PutUint32(buf[offset + 20:], record.DOctets)
        binary.BigEndian.PutUint32(buf[offset + 24:], record.First)
        binary.BigEndian.PutUint32(buf[offset + 28:], record.Last)
        binary.BigEndian.PutUint16(buf[offset + 32:], record.Srcport)
        binary.BigEndian.PutUint16(buf[offset + 34:], record.Dstport)
        buf[offset + 36] = record.Pad1
        buf[offset + 37] = record.Tcp_flags
        buf[offset + 38] = record.Prot
        buf[offset + 39] = record.Tos
        binary.BigEndian.PutUint16(buf[offset + 40:], record.Src_as)
        binary.BigEndian.PutUint16(buf[offset + 42:], record.Dst_as)
        buf[offset + 44] = record.Src_mask
        buf[offset + 45] = record.Dst_mask
        binary.BigEndian.PutUint16(buf[offset + 46:], record.Pad2)
	return NETFLOW_V5_RECORD_SIZE;
}

func construct_payload(num_records uint16) gopacket.Payload {

	buf := gopacket.NewSerializeBuffer()
//        payload := buf.Bytes()
        //Allocate the space we will need for the header
        bytes,err := buf.PrependBytes(NETFLOW_V5_HEADER_SIZE + NETFLOW_V5_RECORD_SIZE*int(num_records))
	if err != nil {
		return nil
	}

	offset := 0

	header := construct_v5_header(num_records, 1000)
	offset += insert_v5_header(header, bytes, offset)

	var record NETFLOW_v5_record;
	for i := 0; i < int(num_records); i++ {
		record = construct_v5_record("1.1.1.1", "2.2.2.2", 5, 256, 80, 5050, 6, 237, 237)
		insert_v5_record(record, bytes, offset)

	}

	return gopacket.Payload(bytes)
}

func chk(err error){
	if err != nil {
		panic(err)
	}
}

func init_connection(addr net.IP) *net.UDPConn {
	conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: addr, Port: 0})
	chk(err)
	return conn
}

func send_packet(conn *net.UDPConn, addr net.IP, port int, pkt []byte) {
	_, err := conn.WriteToUDP(pkt, &net.UDPAddr{IP: addr, Port: port})
	chk(err)
}

func Init (
  traceHandle io.Writer,
  infoHandle io.Writer,
  warningHandle io.Writer,
  errorHandle io.Writer) {

  Trace = log.New(traceHandle, "TRACE: ",
    log.Ldate|log.Ltime|log.Lshortfile)

  Info = log.New(infoHandle,
    "INFO: ",
    log.Ldate|log.Ltime|log.Lshortfile)

  Warning = log.New(warningHandle,
    "WARNING: ",
    log.Ldate|log.Ltime|log.Lshortfile)

  Error = log.New(errorHandle,
    "ERROR: ",
    log.Ldate|log.Ltime|log.Lshortfile)
}

func main() {
  Init(os.Stdout, os.Stdout, os.Stdout, os.Stderr)

	dst_ip := flag.String("dst", "127.0.0.1", "Destination IP to send the spoofed netflow")
	dst_port := flag.Int("port", NETFLOW_PORT, "Destination Port to send the spoofed netflow")
	rate := flag.Int64("rate", 1, "Rate in Packets/s")
  runtime := flag.Int64("time", 10, "Time in seconds to send packets")
  flows_per_packet := flag.Uint("fpp", 1, "flows per packet, max of 30")
  flag.Parse()

	dst_addr := net.ParseIP(*dst_ip)

	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{}

	l2 := construct_ethernet()
	l3 := construct_ip("1.2.3.4", "5.6.7.8")
	l4 := construct_udp()

	payload := construct_payload(uint16(*flows_per_packet))
	//LayerCake
	gopacket.SerializeLayers(buf, opts,
		l2,
		l3,
		l4,
		payload)
	packetData := buf.Bytes()

	//Send the packet to lo
	conn := init_connection(dst_addr)
	send_packet(conn, dst_addr, *dst_port, packetData)

  //Simple way for now. Token based approach later
  throttle := time.Tick(time.Duration(*rate)*time.Second)
  for i := 0; int64(i) < (*rate)*(*runtime); i++ {
    <-throttle
    go send_packet(conn, dst_addr, *dst_port, packetData)
  }
	Info.Println("fin")
}
<|endoftext|>


<commit_before>package sudoku

//GridModification is a series of CellModifications to apply to a Grid.
type GridModifcation []*CellModification

//CellModification represents a modification to be made to a given Cell in a
//grid.
type CellModification struct {
	//The cell representing the cell to modify. The cell's analog (at the same
	//row, col address) will be modified in the new grid.
	Cell Cell
	//The number to put in the cell. Negative numbers signify no changes.
	Number int
	//The excludes to proactively set. Invalid numbers will be ignored.
	//Indexes not listed will be left the same.
	ExcludesChanges map[int]bool
	//The marks to proactively set. Invalid numbers will be ignored.
	//Indexes not listed will be left the same.
	MarksChanges map[int]bool
}

//TODO: audit all uses of step/compoundstep.Apply()

//TOOD: make readOnlyCellImpl. Test if neighbors should be derived or not.
//Everything should be actual contiguous memory, no pointers (except for
//grid). Likely should make cellImpl embed a readOnlyCellImpl and only
//override items it needs to.

//TODO: make readOnlyGridImpl. Two possible approaches: a version that is
//incredibly easy to copy and then do minor tweaks. Or a version that stores a
//dictionary of cell configs, and any time you grab a Cell we look it up in
//the dict or in the ancestors' dicts.

//newCellModification returns a CellModification for the given cell that is a
//no-op.
func newCellModification(cell Cell) *CellModification {
	return &CellModification{
		Cell:            cell,
		Number:          -1,
		ExcludesChanges: make(map[int]bool),
		MarksChanges:    make(map[int]bool),
	}
}

//equivalent returns true if the other grid modification is equivalent to this one.
func (m GridModifcation) equivalent(other GridModifcation) bool {
	if len(m) != len(other) {
		return false
	}
	for i, modification := range m {
		otherModification := other[i]
		if modification.Cell.ref().String() != otherModification.Cell.ref().String() {
			return false
		}
		if modification.Number != otherModification.Number {
			return false
		}

		if len(modification.ExcludesChanges) != len(otherModification.ExcludesChanges) {
			return false
		}

		for key, val := range modification.ExcludesChanges {
			otherVal, ok := otherModification.ExcludesChanges[key]
			if !ok {
				return false
			}
			if val != otherVal {
				return false
			}
		}

		if len(modification.MarksChanges) != len(otherModification.MarksChanges) {
			return false
		}

		for key, val := range modification.MarksChanges {
			otherVal, ok := otherModification.MarksChanges[key]
			if !ok {
				return false
			}
			if val != otherVal {
				return false
			}
		}
	}
	return true
}

func (self *gridImpl) CopyWithModifications(modifications GridModifcation) Grid {

	//TODO: test this implementation deeply! Lots of crazy stuff that could go
	//wrong.

	result := new(gridImpl)

	//Copy in everything
	*result = *self

	cellNumberModified := false

	for _, modification := range modifications {
		cell := result.cellImpl(modification.Cell.Row(), modification.Cell.Col())

		if modification.Number >= 0 && modification.Number <= DIM {
			//cell.setNumber will handle setting all of the impossibles
			if cell.setNumber(modification.Number) {
				cellNumberModified = true
			}
		}

		for key, val := range modification.ExcludesChanges {
			//Key is 1-indexed
			key--
			cell.excluded[key] = val
		}

		for key, val := range modification.MarksChanges {
			//Key is 1-indexed
			key--
			cell.marks[key] = val
		}
	}

	if cellNumberModified {

		//At least one cell's number was modified, which means we need to fix
		//up the queue, numFilledCells, Invalid, Solved.

		filledCellsCount := 0

		for _, cell := range result.cells {
			if cell.number == 0 {
				continue
			}
			filledCellsCount++
		}

		result.filledCellsCount = filledCellsCount

		//Check if we're invalid.

		invalid := false

		for _, cell := range result.cells {
			//Make sure we have at least one possibility per cell
			foundPossibility := false
			for i := 0; i < DIM; i++ {
				if cell.impossibles[i] == 0 {
					foundPossibility = true
					break
				}
			}
			if !foundPossibility {
				invalid = true
				break
			}
		}

		if !invalid {
			//Let's do a deep check
			invalid = gridGroupsInvalid(result)
		}

		result.invalid = invalid

		if filledCellsCount == DIM*DIM && !result.invalid {
			//All cells are filled and it's not invalid, so it's solved!
			result.solved = true
		} else {
			//No way it's solved
			result.solved = false
		}

		result.theQueue.fix()
	}

	return result

}

func (self *mutableGridImpl) CopyWithModifications(modifications GridModifcation) Grid {
	//TODO: when we have an honest-to-god readonly grid impl, optimize this.
	result := self.MutableCopy()

	for _, modification := range modifications {
		cell := modification.Cell.MutableInGrid(result)

		if modification.Number >= 0 && modification.Number <= DIM {
			cell.SetNumber(modification.Number)
		}

		for key, val := range modification.ExcludesChanges {
			//setExcluded will skip invalid entries
			cell.SetExcluded(key, val)
		}

		for key, val := range modification.MarksChanges {
			//SetMark will skip invalid numbers
			cell.SetMark(key, val)
		}
	}

	return result
}
<commit_msg>When doing CopyWithModifications on a gridImpl, fix up all of the gridrefs to the real grid.<commit_after>package sudoku

//GridModification is a series of CellModifications to apply to a Grid.
type GridModifcation []*CellModification

//CellModification represents a modification to be made to a given Cell in a
//grid.
type CellModification struct {
	//The cell representing the cell to modify. The cell's analog (at the same
	//row, col address) will be modified in the new grid.
	Cell Cell
	//The number to put in the cell. Negative numbers signify no changes.
	Number int
	//The excludes to proactively set. Invalid numbers will be ignored.
	//Indexes not listed will be left the same.
	ExcludesChanges map[int]bool
	//The marks to proactively set. Invalid numbers will be ignored.
	//Indexes not listed will be left the same.
	MarksChanges map[int]bool
}

//TODO: audit all uses of step/compoundstep.Apply()

//TOOD: make readOnlyCellImpl. Test if neighbors should be derived or not.
//Everything should be actual contiguous memory, no pointers (except for
//grid). Likely should make cellImpl embed a readOnlyCellImpl and only
//override items it needs to.

//TODO: make readOnlyGridImpl. Two possible approaches: a version that is
//incredibly easy to copy and then do minor tweaks. Or a version that stores a
//dictionary of cell configs, and any time you grab a Cell we look it up in
//the dict or in the ancestors' dicts.

//newCellModification returns a CellModification for the given cell that is a
//no-op.
func newCellModification(cell Cell) *CellModification {
	return &CellModification{
		Cell:            cell,
		Number:          -1,
		ExcludesChanges: make(map[int]bool),
		MarksChanges:    make(map[int]bool),
	}
}

//equivalent returns true if the other grid modification is equivalent to this one.
func (m GridModifcation) equivalent(other GridModifcation) bool {
	if len(m) != len(other) {
		return false
	}
	for i, modification := range m {
		otherModification := other[i]
		if modification.Cell.ref().String() != otherModification.Cell.ref().String() {
			return false
		}
		if modification.Number != otherModification.Number {
			return false
		}

		if len(modification.ExcludesChanges) != len(otherModification.ExcludesChanges) {
			return false
		}

		for key, val := range modification.ExcludesChanges {
			otherVal, ok := otherModification.ExcludesChanges[key]
			if !ok {
				return false
			}
			if val != otherVal {
				return false
			}
		}

		if len(modification.MarksChanges) != len(otherModification.MarksChanges) {
			return false
		}

		for key, val := range modification.MarksChanges {
			otherVal, ok := otherModification.MarksChanges[key]
			if !ok {
				return false
			}
			if val != otherVal {
				return false
			}
		}
	}
	return true
}

func (self *gridImpl) CopyWithModifications(modifications GridModifcation) Grid {

	//TODO: test this implementation deeply! Lots of crazy stuff that could go
	//wrong.

	result := new(gridImpl)

	//Copy in everything
	*result = *self

	for i := 0; i < DIM*DIM; i++ {
		cell := &result.cells[i]
		cell.gridRef = result
	}

	cellNumberModified := false

	for _, modification := range modifications {
		cell := result.cellImpl(modification.Cell.Row(), modification.Cell.Col())

		if modification.Number >= 0 && modification.Number <= DIM {
			//cell.setNumber will handle setting all of the impossibles
			if cell.setNumber(modification.Number) {
				cellNumberModified = true
			}
		}

		for key, val := range modification.ExcludesChanges {
			//Key is 1-indexed
			key--
			cell.excluded[key] = val
		}

		for key, val := range modification.MarksChanges {
			//Key is 1-indexed
			key--
			cell.marks[key] = val
		}
	}

	if cellNumberModified {

		//At least one cell's number was modified, which means we need to fix
		//up the queue, numFilledCells, Invalid, Solved.

		filledCellsCount := 0

		for _, cell := range result.cells {
			if cell.number == 0 {
				continue
			}
			filledCellsCount++
		}

		result.filledCellsCount = filledCellsCount

		//Check if we're invalid.

		invalid := false

		for _, cell := range result.cells {
			//Make sure we have at least one possibility per cell
			foundPossibility := false
			for i := 0; i < DIM; i++ {
				if cell.impossibles[i] == 0 {
					foundPossibility = true
					break
				}
			}
			if !foundPossibility {
				invalid = true
				break
			}
		}

		if !invalid {
			//Let's do a deep check
			invalid = gridGroupsInvalid(result)
		}

		result.invalid = invalid

		if filledCellsCount == DIM*DIM && !result.invalid {
			//All cells are filled and it's not invalid, so it's solved!
			result.solved = true
		} else {
			//No way it's solved
			result.solved = false
		}

		result.theQueue.fix()
	}

	return result

}

func (self *mutableGridImpl) CopyWithModifications(modifications GridModifcation) Grid {
	//TODO: when we have an honest-to-god readonly grid impl, optimize this.
	result := self.MutableCopy()

	for _, modification := range modifications {
		cell := modification.Cell.MutableInGrid(result)

		if modification.Number >= 0 && modification.Number <= DIM {
			cell.SetNumber(modification.Number)
		}

		for key, val := range modification.ExcludesChanges {
			//setExcluded will skip invalid entries
			cell.SetExcluded(key, val)
		}

		for key, val := range modification.MarksChanges {
			//SetMark will skip invalid numbers
			cell.SetMark(key, val)
		}
	}

	return result
}
<|endoftext|>


<commit_before>package release

import (
	"fmt"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/cluster"
	fluxmetrics "github.com/weaveworks/flux/metrics"
	"github.com/weaveworks/flux/policy"
	"github.com/weaveworks/flux/registry"
	"github.com/weaveworks/flux/update"
)

type ServiceUpdate struct {
	ServiceID     flux.ServiceID
	Service       cluster.Service
	ManifestPath  string
	ManifestBytes []byte
	Updates       []update.ContainerUpdate
}

func Release(rc *ReleaseContext, spec update.ReleaseSpec, cause update.Cause, logger log.Logger) (commitRef string, results update.Result, err error) {
	started := time.Now()
	defer func(start time.Time) {
		releaseDuration.With(
			fluxmetrics.LabelSuccess, fmt.Sprint(err == nil),
			fluxmetrics.LabelReleaseType, update.ReleaseSpec(spec).ReleaseType(),
			fluxmetrics.LabelReleaseKind, string(spec.Kind),
		).Observe(time.Since(started).Seconds())
	}(started)

	logger = log.NewContext(logger).With("type", "release")
	// We time each stage of this process, and expose as metrics.
	var timer *metrics.Timer

	// FIXME pull from the repository? Or rely on something else to do that.
	// ALSO: clean up in the result of failure, afterwards

	// From here in, we collect the results of the calculations.
	results = update.Result{}

	// Figure out the services involved.
	timer = NewStageTimer("select_services")
	var updates []*ServiceUpdate
	updates, err = selectServices(rc, &spec, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	// If the request was for a specific service and the service was
	// not in the cluster, then add a result with a skipped status
	for _, v := range spec.ServiceSpecs {
		if v == update.ServiceSpecAll {
			continue
		}
		id, err := v.AsID()
		if err != nil {
			continue
		}
		if _, ok := results[id]; !ok {
			results[id] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  NotInRepo,
			}
		}
	}

	// Look up images, and calculate updates
	timer = NewStageTimer("lookup_images")
	// Figure out how the services are to be updated.
	updates, err = calculateImageUpdates(rc, updates, &spec, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	// At this point we may have filtered the updates we can do down
	// to nothing. Check and exit early if so.
	logger.Log("updates", len(updates))
	if len(updates) == 0 {
		logger.Log("exit", "no images to update for services given")
		return "", results, nil
	}

	// If it's a dry run, we're done.
	if spec.Kind == update.ReleaseKindPlan {
		logger.Log("exit", "dry-run")
		return "", results, nil
	}

	timer = NewStageTimer("push_changes")
	err = rc.PushChanges(updates, &spec, cause, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	revision, err := rc.Repo.HeadRevision()
	return revision, results, err
}

// Take the spec given in the job, and figure out which services are
// in question based on the running services and those defined in the
// repo. Fill in the release results along the way.
func selectServices(rc *ReleaseContext, spec *update.ReleaseSpec, results update.Result) ([]*ServiceUpdate, error) {
	// Build list of filters
	filtList, err := filters(spec, rc)
	if err != nil {
		return nil, err
	}
	// Find and filter services
	return rc.SelectServices(
		results,
		filtList...,
	)
}

// filters converts a ReleaseSpec (and Lock config) into ServiceFilters
func filters(spec *update.ReleaseSpec, rc *ReleaseContext) ([]ServiceFilter, error) {
	// Image filter
	var filtList []ServiceFilter
	if spec.ImageSpec != update.ImageSpecLatest {
		id, err := flux.ParseImageID(spec.ImageSpec.String())
		if err != nil {
			return nil, err
		}
		filtList = append(filtList, &SpecificImageFilter{id})
	}

	// Service filter
	ids := []flux.ServiceID{}
specs:
	for _, s := range spec.ServiceSpecs {
		switch s {
		case update.ServiceSpecAll:
			// "<all>" Overrides any other filters
			ids = []flux.ServiceID{}
			break specs
		case update.ServiceSpecAutomated:
			// "<automated>" Overrides any other filters
			automated, err := rc.ServicesWithPolicy(policy.Automated)
			if err != nil {
				return nil, err
			}
			ids = automated.ToSlice()
			break specs
		}
		id, err := flux.ParseServiceID(string(s))
		if err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	if len(ids) > 0 {
		filtList = append(filtList, &IncludeFilter{ids})
	}

	// Exclude filter
	if len(spec.Excludes) > 0 {
		filtList = append(filtList, &ExcludeFilter{spec.Excludes})
	}

	// Locked filter
	lockedSet, err := rc.ServicesWithPolicy(policy.Locked)
	if err != nil {
		return nil, err
	}
	filtList = append(filtList, &LockedFilter{lockedSet.ToSlice()})
	return filtList, nil
}

// Find all the image updates that should be performed, and do
// replacements. (For a dry-run, we don't strictly need to do the
// replacements, since we won't be committing any changes back;
// however we do want to see if we *can* do the replacements, because
// if not, it indicates there's likely some problem with the running
// system vs the definitions given in the repo.)
func calculateImageUpdates(rc *ReleaseContext, candidates []*ServiceUpdate, spec *update.ReleaseSpec, results update.Result) ([]*ServiceUpdate, error) {
	// Compile an `ImageMap` of all relevant images
	var images ImageMap
	var err error

	switch spec.ImageSpec {
	case update.ImageSpecLatest:
		images, err = CollectUpdateImages(rc.Registry, candidates)
	default:
		var image flux.ImageID
		image, err = spec.ImageSpec.AsID()
		if err == nil {
			images, err = ExactImages(rc.Registry, []flux.ImageID{image})
		}
	}

	if err != nil {
		return nil, err
	}

	// Look through all the services' containers to see which have an
	// image that could be updated.
	var updates []*ServiceUpdate
	for _, u := range candidates {
		containers, err := u.Service.ContainersOrError()
		if err != nil {
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusFailed,
				Error:  err.Error(),
			}
			continue
		}

		// If at least one container used an image in question, we say
		// we're skipping it rather than ignoring it. This is mainly
		// for the purpose of filtering the output.
		ignoredOrSkipped := update.ReleaseStatusIgnored
		var containerUpdates []update.ContainerUpdate

		for _, container := range containers {
			currentImageID, err := flux.ParseImageID(container.Image)
			if err != nil {
				// We may hope never to find a malformed image ID, but
				// anything is possible.
				return nil, err
			}

			latestImage := images.LatestImage(currentImageID.Repository())
			if latestImage == nil {
				ignoredOrSkipped = update.ReleaseStatusUnknown
				continue
			}

			if currentImageID == latestImage.ID {
				ignoredOrSkipped = update.ReleaseStatusSkipped
				continue
			}

			u.ManifestBytes, err = rc.Manifests.UpdateDefinition(u.ManifestBytes, latestImage.ID)
			if err != nil {
				return nil, err
			}

			containerUpdates = append(containerUpdates, update.ContainerUpdate{
				Container: container.Name,
				Current:   currentImageID,
				Target:    latestImage.ID,
			})
		}

		switch {
		case len(containerUpdates) > 0:
			u.Updates = containerUpdates
			updates = append(updates, u)
			results[u.ServiceID] = update.ServiceResult{
				Status:       update.ReleaseStatusSuccess,
				PerContainer: containerUpdates,
			}
		case ignoredOrSkipped == update.ReleaseStatusSkipped:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  ImageUpToDate,
			}
		case ignoredOrSkipped == update.ReleaseStatusIgnored:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusIgnored,
				Error:  "does not use image(s)",
			}
		case ignoredOrSkipped == update.ReleaseStatusUnknown:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  ImageNotFound,
			}
		}
	}

	return updates, nil
}

func commitMessageFromReleaseSpec(spec *update.ReleaseSpec) string {
	image := strings.Trim(spec.ImageSpec.String(), "<>")
	var services []string
	for _, s := range spec.ServiceSpecs {
		services = append(services, strings.Trim(s.String(), "<>"))
	}
	return fmt.Sprintf("Release %s to %s", image, strings.Join(services, ", "))
}

// CollectUpdateImages is a convenient shim to
// `CollectAvailableImages`.
func CollectUpdateImages(registry registry.Registry, updateable []*ServiceUpdate) (ImageMap, error) {
	var servicesToCheck []cluster.Service
	for _, update := range updateable {
		servicesToCheck = append(servicesToCheck, update.Service)
	}
	return CollectAvailableImages(registry, servicesToCheck)
}
<commit_msg>Mark container ignored on release if repos don't match<commit_after>package release

import (
	"fmt"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/cluster"
	fluxmetrics "github.com/weaveworks/flux/metrics"
	"github.com/weaveworks/flux/policy"
	"github.com/weaveworks/flux/registry"
	"github.com/weaveworks/flux/update"
)

type ServiceUpdate struct {
	ServiceID     flux.ServiceID
	Service       cluster.Service
	ManifestPath  string
	ManifestBytes []byte
	Updates       []update.ContainerUpdate
}

func Release(rc *ReleaseContext, spec update.ReleaseSpec, cause update.Cause, logger log.Logger) (commitRef string, results update.Result, err error) {
	started := time.Now()
	defer func(start time.Time) {
		releaseDuration.With(
			fluxmetrics.LabelSuccess, fmt.Sprint(err == nil),
			fluxmetrics.LabelReleaseType, update.ReleaseSpec(spec).ReleaseType(),
			fluxmetrics.LabelReleaseKind, string(spec.Kind),
		).Observe(time.Since(started).Seconds())
	}(started)

	logger = log.NewContext(logger).With("type", "release")
	// We time each stage of this process, and expose as metrics.
	var timer *metrics.Timer

	// FIXME pull from the repository? Or rely on something else to do that.
	// ALSO: clean up in the result of failure, afterwards

	// From here in, we collect the results of the calculations.
	results = update.Result{}

	// Figure out the services involved.
	timer = NewStageTimer("select_services")
	var updates []*ServiceUpdate
	updates, err = selectServices(rc, &spec, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	// If the request was for a specific service and the service was
	// not in the cluster, then add a result with a skipped status
	for _, v := range spec.ServiceSpecs {
		if v == update.ServiceSpecAll {
			continue
		}
		id, err := v.AsID()
		if err != nil {
			continue
		}
		if _, ok := results[id]; !ok {
			results[id] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  NotInRepo,
			}
		}
	}

	// Look up images, and calculate updates
	timer = NewStageTimer("lookup_images")
	// Figure out how the services are to be updated.
	updates, err = calculateImageUpdates(rc, updates, &spec, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	// At this point we may have filtered the updates we can do down
	// to nothing. Check and exit early if so.
	logger.Log("updates", len(updates))
	if len(updates) == 0 {
		logger.Log("exit", "no images to update for services given")
		return "", results, nil
	}

	// If it's a dry run, we're done.
	if spec.Kind == update.ReleaseKindPlan {
		logger.Log("exit", "dry-run")
		return "", results, nil
	}

	timer = NewStageTimer("push_changes")
	err = rc.PushChanges(updates, &spec, cause, results)
	timer.ObserveDuration()
	if err != nil {
		return "", nil, err
	}

	revision, err := rc.Repo.HeadRevision()
	return revision, results, err
}

// Take the spec given in the job, and figure out which services are
// in question based on the running services and those defined in the
// repo. Fill in the release results along the way.
func selectServices(rc *ReleaseContext, spec *update.ReleaseSpec, results update.Result) ([]*ServiceUpdate, error) {
	// Build list of filters
	filtList, err := filters(spec, rc)
	if err != nil {
		return nil, err
	}
	// Find and filter services
	return rc.SelectServices(
		results,
		filtList...,
	)
}

// filters converts a ReleaseSpec (and Lock config) into ServiceFilters
func filters(spec *update.ReleaseSpec, rc *ReleaseContext) ([]ServiceFilter, error) {
	// Image filter
	var filtList []ServiceFilter
	if spec.ImageSpec != update.ImageSpecLatest {
		id, err := flux.ParseImageID(spec.ImageSpec.String())
		if err != nil {
			return nil, err
		}
		filtList = append(filtList, &SpecificImageFilter{id})
	}

	// Service filter
	ids := []flux.ServiceID{}
specs:
	for _, s := range spec.ServiceSpecs {
		switch s {
		case update.ServiceSpecAll:
			// "<all>" Overrides any other filters
			ids = []flux.ServiceID{}
			break specs
		case update.ServiceSpecAutomated:
			// "<automated>" Overrides any other filters
			automated, err := rc.ServicesWithPolicy(policy.Automated)
			if err != nil {
				return nil, err
			}
			ids = automated.ToSlice()
			break specs
		}
		id, err := flux.ParseServiceID(string(s))
		if err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	if len(ids) > 0 {
		filtList = append(filtList, &IncludeFilter{ids})
	}

	// Exclude filter
	if len(spec.Excludes) > 0 {
		filtList = append(filtList, &ExcludeFilter{spec.Excludes})
	}

	// Locked filter
	lockedSet, err := rc.ServicesWithPolicy(policy.Locked)
	if err != nil {
		return nil, err
	}
	filtList = append(filtList, &LockedFilter{lockedSet.ToSlice()})
	return filtList, nil
}

// Find all the image updates that should be performed, and do
// replacements. (For a dry-run, we don't strictly need to do the
// replacements, since we won't be committing any changes back;
// however we do want to see if we *can* do the replacements, because
// if not, it indicates there's likely some problem with the running
// system vs the definitions given in the repo.)
func calculateImageUpdates(rc *ReleaseContext, candidates []*ServiceUpdate, spec *update.ReleaseSpec, results update.Result) ([]*ServiceUpdate, error) {
	// Compile an `ImageMap` of all relevant images
	var images ImageMap
	var repo string
	var err error

	switch spec.ImageSpec {
	case update.ImageSpecLatest:
		images, err = CollectUpdateImages(rc.Registry, candidates)
	default:
		var image flux.ImageID
		image, err = spec.ImageSpec.AsID()
		if err == nil {
			repo = image.Repository()
			images, err = ExactImages(rc.Registry, []flux.ImageID{image})
		}
	}

	if err != nil {
		return nil, err
	}

	// Look through all the services' containers to see which have an
	// image that could be updated.
	var updates []*ServiceUpdate
	for _, u := range candidates {
		containers, err := u.Service.ContainersOrError()
		if err != nil {
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusFailed,
				Error:  err.Error(),
			}
			continue
		}

		// If at least one container used an image in question, we say
		// we're skipping it rather than ignoring it. This is mainly
		// for the purpose of filtering the output.
		ignoredOrSkipped := update.ReleaseStatusIgnored
		var containerUpdates []update.ContainerUpdate

		for _, container := range containers {
			currentImageID, err := flux.ParseImageID(container.Image)
			if err != nil {
				// We may hope never to find a malformed image ID, but
				// anything is possible.
				return nil, err
			}

			latestImage := images.LatestImage(currentImageID.Repository())
			if latestImage == nil {
				if currentImageID.Repository() != repo {
					ignoredOrSkipped = update.ReleaseStatusIgnored
				} else {
					ignoredOrSkipped = update.ReleaseStatusUnknown
				}
				continue
			}

			if currentImageID == latestImage.ID {
				ignoredOrSkipped = update.ReleaseStatusSkipped
				continue
			}

			u.ManifestBytes, err = rc.Manifests.UpdateDefinition(u.ManifestBytes, latestImage.ID)
			if err != nil {
				return nil, err
			}

			containerUpdates = append(containerUpdates, update.ContainerUpdate{
				Container: container.Name,
				Current:   currentImageID,
				Target:    latestImage.ID,
			})
		}

		switch {
		case len(containerUpdates) > 0:
			u.Updates = containerUpdates
			updates = append(updates, u)
			results[u.ServiceID] = update.ServiceResult{
				Status:       update.ReleaseStatusSuccess,
				PerContainer: containerUpdates,
			}
		case ignoredOrSkipped == update.ReleaseStatusSkipped:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  ImageUpToDate,
			}
		case ignoredOrSkipped == update.ReleaseStatusIgnored:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusIgnored,
				Error:  "does not use image(s)",
			}
		case ignoredOrSkipped == update.ReleaseStatusUnknown:
			results[u.ServiceID] = update.ServiceResult{
				Status: update.ReleaseStatusSkipped,
				Error:  ImageNotFound,
			}
		}
	}

	return updates, nil
}

func commitMessageFromReleaseSpec(spec *update.ReleaseSpec) string {
	image := strings.Trim(spec.ImageSpec.String(), "<>")
	var services []string
	for _, s := range spec.ServiceSpecs {
		services = append(services, strings.Trim(s.String(), "<>"))
	}
	return fmt.Sprintf("Release %s to %s", image, strings.Join(services, ", "))
}

// CollectUpdateImages is a convenient shim to
// `CollectAvailableImages`.
func CollectUpdateImages(registry registry.Registry, updateable []*ServiceUpdate) (ImageMap, error) {
	var servicesToCheck []cluster.Service
	for _, update := range updateable {
		servicesToCheck = append(servicesToCheck, update.Service)
	}
	return CollectAvailableImages(registry, servicesToCheck)
}
<|endoftext|>