<commit_before>/*
   Copyright 2014 Krishna Raman <kraman@gmail.com>

   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.
*/

/*
Go bindings for libnetfilter_queue

This library provides access to packets in the IPTables netfilter queue (NFQUEUE).
The libnetfilter_queue library is part of the http://netfilter.org/projects/libnetfilter_queue/ project.
*/
package netfilter

/*
#cgo pkg-config: libnetfilter_queue
#cgo CFLAGS: -Wall -I/usr/include
#cgo LDFLAGS: -L/usr/lib64/

#include "netfilter.h"
*/
import "C"

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"os"
	"sync"
	"time"
	"unsafe"
)

type NFPacket struct {
	Packet         gopacket.Packet
	verdictChannel chan Verdict
}

//Set the verdict for the packet
func (p *NFPacket) SetVerdict(v Verdict) {
	p.verdictChannel <- v
}

//Set the verdict for the packet
func (p *NFPacket) SetRequeueVerdict(newQueueId uint16) {
	v := uint(NF_QUEUE)
	q := (uint(newQueueId) << 16)
	v = v | q
	p.verdictChannel <- Verdict(v)
}

type NFQueue struct {
	h       *C.struct_nfq_handle
	qh      *C.struct_nfq_q_handle
	fd      C.int
	packets chan NFPacket
	idx     uint32
}

//Verdict for a packet
type Verdict C.uint

const (
	AF_INET = 2
	AF_INET6 = 10

	NF_DROP   Verdict = 0
	NF_ACCEPT Verdict = 1
	NF_STOLEN Verdict = 2
	NF_QUEUE  Verdict = 3
	NF_REPEAT Verdict = 4
	NF_STOP   Verdict = 5

	NF_DEFAULT_PACKET_SIZE uint32 = 0xffff
)

var theTable = make(map[uint32]*chan NFPacket, 0)
var theTabeLock sync.RWMutex

//Create and bind to queue specified by queueId
func NewNFQueue(queueId uint16, maxPacketsInQueue uint32, packetSize uint32) (*NFQueue, error) {
	var nfq = NFQueue{}
	var err error
	var ret C.int

	if nfq.h, err = C.nfq_open(); err != nil {
		return nil, fmt.Errorf("Error opening NFQueue handle: %v\n", err)
	}

	if ret, err = C.nfq_unbind_pf(nfq.h, AF_INET); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error unbinding existing NFQ handler from AF_INET protocol family: %v\n", err)
	}

	if ret, err = C.nfq_unbind_pf(nfq.h, AF_INET6); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error unbinding existing NFQ handler from AF_INET6 protocol family: %v\n", err)
	}

	if ret, err := C.nfq_bind_pf(nfq.h, AF_INET); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error binding to AF_INET protocol family: %v\n", err)
	}

	if ret, err := C.nfq_bind_pf(nfq.h, AF_INET6); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error binding to AF_INET6 protocol family: %v\n", err)
	}

	nfq.packets = make(chan NFPacket)
	nfq.idx = uint32(time.Now().UnixNano())
	theTabeLock.Lock()
	theTable[nfq.idx] = &nfq.packets
	theTabeLock.Unlock()
	if nfq.qh, err = C.CreateQueue(nfq.h, C.u_int16_t(queueId), C.u_int32_t(nfq.idx)); err != nil || nfq.qh == nil {
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Error binding to queue: %v\n", err)
	}

	if ret, err = C.nfq_set_queue_maxlen(nfq.qh, C.u_int32_t(maxPacketsInQueue)); err != nil || ret < 0 {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to set max packets in queue: %v\n", err)
	}

	if C.nfq_set_mode(nfq.qh, C.u_int8_t(2), C.uint(packetSize)) < 0 {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to set packets copy mode: %v\n", err)
	}

	if nfq.fd, err = C.nfq_fd(nfq.h); err != nil {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to get queue file-descriptor. %v", err)
	}

	go nfq.run()

	return &nfq, nil
}

//Unbind and close the queue
func (nfq *NFQueue) Close() {
	C.nfq_destroy_queue(nfq.qh)
	C.nfq_close(nfq.h)
	theTabeLock.Lock()
	delete(theTable, nfq.idx)
	theTabeLock.Unlock()
}

//Get the channel for packets
func (nfq *NFQueue) GetPackets() <-chan NFPacket {
	return nfq.packets
}

func (nfq *NFQueue) run() {
	C.Run(nfq.h, nfq.fd)
}

//export go_callback
func go_callback(queueId C.int, data *C.uchar, len C.int, idx uint32) Verdict {
	xdata := C.GoBytes(unsafe.Pointer(data), len)

	var packet gopacket.Packet
	if xdata[0] & 0xf0 == 0x40 {
		packet = gopacket.NewPacket(xdata, layers.LayerTypeIPv4, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
	} else {
		packet = gopacket.NewPacket(xdata, layers.LayerTypeIPv6, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
	}

	p := NFPacket{verdictChannel: make(chan Verdict), Packet: packet}
	theTabeLock.RLock()
	cb, ok := theTable[idx]
	theTabeLock.RUnlock()
	if !ok {
		fmt.Fprintf(os.Stderr, "Dropping, unexpectedly due to bad idx=%d\n", idx)
		return NF_DROP
	}
	select {
	case *cb <- p:
		v := <-p.verdictChannel
		return v
	default:
		fmt.Fprintf(os.Stderr, "Dropping, unexpectedly due to no recv, idx=%d\n", idx)
		return NF_DROP
	}
}
<commit_msg>contants instead of magic values, extra \n added<commit_after>/*
   Copyright 2014 Krishna Raman <kraman@gmail.com>

   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.
*/

/*
Go bindings for libnetfilter_queue

This library provides access to packets in the IPTables netfilter queue (NFQUEUE).
The libnetfilter_queue library is part of the http://netfilter.org/projects/libnetfilter_queue/ project.
*/
package netfilter

/*
#cgo pkg-config: libnetfilter_queue
#cgo CFLAGS: -Wall -I/usr/include
#cgo LDFLAGS: -L/usr/lib64/

#include "netfilter.h"
*/
import "C"

import (
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"os"
	"sync"
	"time"
	"unsafe"
)

type NFPacket struct {
	Packet         gopacket.Packet
	verdictChannel chan Verdict
}

//Set the verdict for the packet
func (p *NFPacket) SetVerdict(v Verdict) {
	p.verdictChannel <- v
}

//Set the verdict for the packet
func (p *NFPacket) SetRequeueVerdict(newQueueId uint16) {
	v := uint(NF_QUEUE)
	q := (uint(newQueueId) << 16)
	v = v | q
	p.verdictChannel <- Verdict(v)
}

type NFQueue struct {
	h       *C.struct_nfq_handle
	qh      *C.struct_nfq_q_handle
	fd      C.int
	packets chan NFPacket
	idx     uint32
}

//Verdict for a packet
type Verdict C.uint

const (
	AF_INET = 2
	AF_INET6 = 10

	NF_DROP   Verdict = 0
	NF_ACCEPT Verdict = 1
	NF_STOLEN Verdict = 2
	NF_QUEUE  Verdict = 3
	NF_REPEAT Verdict = 4
	NF_STOP   Verdict = 5

	NF_DEFAULT_PACKET_SIZE uint32 = 0xffff

	ipv4version = 0x40
)

var theTable = make(map[uint32]*chan NFPacket, 0)
var theTabeLock sync.RWMutex

//Create and bind to queue specified by queueId
func NewNFQueue(queueId uint16, maxPacketsInQueue uint32, packetSize uint32) (*NFQueue, error) {
	var nfq = NFQueue{}
	var err error
	var ret C.int

	if nfq.h, err = C.nfq_open(); err != nil {
		return nil, fmt.Errorf("Error opening NFQueue handle: %v\n", err)
	}

	if ret, err = C.nfq_unbind_pf(nfq.h, AF_INET); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error unbinding existing NFQ handler from AF_INET protocol family: %v\n", err)
	}

	if ret, err = C.nfq_unbind_pf(nfq.h, AF_INET6); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error unbinding existing NFQ handler from AF_INET6 protocol family: %v\n", err)
	}

	if ret, err := C.nfq_bind_pf(nfq.h, AF_INET); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error binding to AF_INET protocol family: %v\n", err)
	}

	if ret, err := C.nfq_bind_pf(nfq.h, AF_INET6); err != nil || ret < 0 {
		return nil, fmt.Errorf("Error binding to AF_INET6 protocol family: %v\n", err)
	}

	nfq.packets = make(chan NFPacket)
	nfq.idx = uint32(time.Now().UnixNano())
	theTabeLock.Lock()
	theTable[nfq.idx] = &nfq.packets
	theTabeLock.Unlock()
	if nfq.qh, err = C.CreateQueue(nfq.h, C.u_int16_t(queueId), C.u_int32_t(nfq.idx)); err != nil || nfq.qh == nil {
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Error binding to queue: %v\n", err)
	}

	if ret, err = C.nfq_set_queue_maxlen(nfq.qh, C.u_int32_t(maxPacketsInQueue)); err != nil || ret < 0 {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to set max packets in queue: %v\n", err)
	}

	if C.nfq_set_mode(nfq.qh, C.u_int8_t(2), C.uint(packetSize)) < 0 {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to set packets copy mode: %v\n", err)
	}

	if nfq.fd, err = C.nfq_fd(nfq.h); err != nil {
		C.nfq_destroy_queue(nfq.qh)
		C.nfq_close(nfq.h)
		return nil, fmt.Errorf("Unable to get queue file-descriptor. %v\n", err)
	}

	go nfq.run()

	return &nfq, nil
}

//Unbind and close the queue
func (nfq *NFQueue) Close() {
	C.nfq_destroy_queue(nfq.qh)
	C.nfq_close(nfq.h)
	theTabeLock.Lock()
	delete(theTable, nfq.idx)
	theTabeLock.Unlock()
}

//Get the channel for packets
func (nfq *NFQueue) GetPackets() <-chan NFPacket {
	return nfq.packets
}

func (nfq *NFQueue) run() {
	C.Run(nfq.h, nfq.fd)
}

//export go_callback
func go_callback(queueId C.int, data *C.uchar, len C.int, idx uint32) Verdict {
	xdata := C.GoBytes(unsafe.Pointer(data), len)

	var packet gopacket.Packet
	if xdata[0] & 0xf0 == ipv4version {
		packet = gopacket.NewPacket(xdata, layers.LayerTypeIPv4, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
	} else {
		packet = gopacket.NewPacket(xdata, layers.LayerTypeIPv6, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
	}

	p := NFPacket{verdictChannel: make(chan Verdict), Packet: packet}
	theTabeLock.RLock()
	cb, ok := theTable[idx]
	theTabeLock.RUnlock()
	if !ok {
		fmt.Fprintf(os.Stderr, "Dropping, unexpectedly due to bad idx=%d\n", idx)
		return NF_DROP
	}
	select {
	case *cb <- p:
		v := <-p.verdictChannel
		return v
	default:
		fmt.Fprintf(os.Stderr, "Dropping, unexpectedly due to no recv, idx=%d\n", idx)
		return NF_DROP
	}
}
<|endoftext|>


<commit_before>/*
D. J. Bernstein's netstrings for Go.
*/
package netstring

import (
	"errors"
	"io"
)

var Incomplete = errors.New("The netstring is incomplete")
var Garbled = errors.New("The netstring was not correctly formatted and could not be read")

type Netstring struct {
	buffer   []byte
	complete bool
}

// Construct a netstring wrapping a byte slice, for output.
// n.IsComplete() will return true.
func From(buf []byte) *Netstring {
	return &Netstring{
		buffer:   buf,
		complete: true,
	}
}

// Construct an empty netstring, for input.
// n.IsComplete() will return false.
func ForReading() *Netstring {
	return &Netstring{buffer: nil, complete: false}
}

// Returns true if the number of bytes advertized in the netstring's length have been read into its buffer.
// Operations that require the netstring's contents will not be available until it is complete.
func (n *Netstring) IsComplete() bool { return n.complete }

// Returns the advertized length of the netstring.
// If n.IsComplete() then this is the length of the data in the buffer, too.
// If this value is negative, in means that no length has been read yet.
func (n *Netstring) Length() int {
	if n.buffer == nil {
		return -1
	}
	return cap(n.buffer)
}

// Returns the bytes in the netstring if it is complete, otherwise returns Incomplete as an error.
func (n *Netstring) Bytes() ([]byte, error) {
	if n.complete {
		return n.buffer, nil
	}
	return nil, Incomplete
}

// Read a netstring from input.
// Returns any errors from input except io.EOF.
// Returns Garbled if the input was not a valid netstring.
// Returns Incomplete if the input was shorter than a full netstring.
// To resume reading where you left off, call Readfrom(input) again.
// Calling Readfrom(input) on a complete netstring does nothing.
func (n *Netstring) ReadFrom(input io.Reader) error {
	var err error
	if n.buffer == nil {
		var length int
		length, err = n.readLength(input)
		if err != nil {
			return err
		}
		err = n.readColon(input)
		if err != nil {
			return err
		}
		n.buffer = make([]byte, 0, length) // capacity stores the length
	}
	if len(n.buffer) < cap(n.buffer) {
		// slice n.buffer to the part between len and cap
		dest := n.buffer[len(n.buffer):cap(n.buffer)]
		var count int
		count, err = input.Read(dest)

		// slice n.buffer to add on count bytes
		if count > 0 {
			n.buffer = n.buffer[:len(n.buffer)+count]
		}

		switch {
		case err == io.EOF: // we still expect to read a comma, so EOF here is always incomplete
			return Incomplete
		case err != nil:
			return err
		case len(n.buffer) < cap(n.buffer):
			return Incomplete
		}
	}
	if !n.complete {
		err = n.readComma(input)
		if err != nil {
			return err
		}
		n.complete = true
	}
	return nil
}

func (n *Netstring) readLength(input io.Reader) (int, error) {
	//TODO
	return 0, nil
}

func (n *Netstring) readColon(input io.Reader) error {
	//TODO
	return nil
}

func (n *Netstring) readComma(input io.Reader) error {
	//TODO
	return nil
}
<commit_msg>make ReadFrom work<commit_after>/*
D. J. Bernstein's netstrings for Go.
*/
package netstring

import (
	"bufio"
	"errors"
	"io"
	"strconv"
)

var Incomplete = errors.New("The netstring is incomplete")
var Garbled = errors.New("The netstring was not correctly formatted and could not be read")

type Netstring struct {
	buffer   []byte
	complete bool
}

// Construct a netstring wrapping a byte slice, for output.
// n.IsComplete() will return true.
func From(buf []byte) *Netstring {
	return &Netstring{
		buffer:   buf,
		complete: true,
	}
}

// Construct an empty netstring, for input.
// n.IsComplete() will return false.
func ForReading() *Netstring {
	return &Netstring{buffer: nil, complete: false}
}

// Returns true if the number of bytes advertized in the netstring's length have been read into its buffer.
// Operations that require the netstring's contents will not be available until it is complete.
func (n *Netstring) IsComplete() bool { return n.complete }

// Returns the advertized length of the netstring.
// If n.IsComplete() then this is the length of the data in the buffer, too.
// If this value is negative, in means that no length has been read yet.
func (n *Netstring) Length() int {
	if n.buffer == nil {
		return -1
	}
	return cap(n.buffer)
}

// Returns the bytes in the netstring if it is complete, otherwise returns Incomplete as an error.
func (n *Netstring) Bytes() ([]byte, error) {
	if n.complete {
		return n.buffer, nil
	}
	return nil, Incomplete
}

// Read a netstring from input.
// Returns any errors from input except io.EOF.
// Returns Garbled if the input was not a valid netstring.
// Returns Incomplete if the input was shorter than a full netstring.
// To resume reading where you left off, call Readfrom(input) again.
// Calling Readfrom(input) on a complete netstring does nothing.
func (n *Netstring) ReadFrom(input io.Reader) error {
	var err error
	if n.buffer == nil {
		var length int
		length, err = n.readLength(input)
		if err != nil {
			return err
		}
		n.buffer = make([]byte, 0, length) // capacity stores the length
	}
	if len(n.buffer) < cap(n.buffer) {
		// slice n.buffer to the part between len and cap
		dest := n.buffer[len(n.buffer):cap(n.buffer)]
		var count int
		count, err = input.Read(dest)

		// slice n.buffer to add on count bytes
		if count > 0 {
			n.buffer = n.buffer[:len(n.buffer)+count]
		}

		switch {
		case err == io.EOF: // we still expect to read a comma, so EOF here is always incomplete
			return Incomplete
		case err != nil:
			return err
		case len(n.buffer) < cap(n.buffer):
			return Incomplete
		}
	}
	if !n.complete {
		err = n.readComma(input)
		if err != nil {
			return err
		}
		n.complete = true
	}
	return nil
}

func (n *Netstring) readLength(input io.Reader) (int, error) {
	digits := make([]byte, 0, 10)
	for {
		digit := make([]byte, 1)
		_, err := input.Read(digit)
		if err != nil {
			return -1, err
		}
		switch rune(digit[0]) {
		case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
			digits = append(digits, digit[0])
			if len(digits) > 10 {
				return -1, Garbled
			}
		case ':':
			length, err := strconv.Atoi(string(digits))
			if err != nil {
				return -1, Garbled
			}
			return length, nil
		default:
			return -1, Garbled
		}
	}
}

func (n *Netstring) readComma(input io.Reader) error {
	c, _, err := bufio.NewReaderSize(input, 1).ReadRune()
	switch {
	case err == io.EOF:
		return Incomplete
	case err != nil:
		return err
	case c == ',':
		return nil
	default:
		return Garbled
	}
}
<|endoftext|>


<commit_before>package build

import (
	"fmt"
	"io/ioutil"
	"neon/util"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"sort"
	"strings"

	"strconv"

	anko_core "github.com/mattn/anko/builtins"
	"github.com/mattn/anko/parser"
	"github.com/mattn/anko/vm"
)

// Build context
type Context struct {
	VM          *vm.Env
	Properties  []string
	Environment map[string]string
	Index       *Index
	Stack       *Stack
	Thread      *int
}

// NewContext make a new build context
func NewContext(build *Build) (*Context, error) {
	v := vm.NewEnv()
	anko_core.LoadAllBuiltins(v)
	LoadBuiltins(v)
	properties := build.GetProperties()
	environment := build.GetEnvironment()
	context := &Context{
		VM:          v,
		Properties:  properties.Fields(),
		Environment: environment,
		Index:       NewIndex(),
		Stack:       NewStack(),
	}
	for _, script := range build.Scripts {
		source, err := ioutil.ReadFile(script)
		if err != nil {
			return nil, fmt.Errorf("reading script '%s': %v", script, err)
		}
		_, err = v.Execute(string(source))
		if err != nil {
			return nil, fmt.Errorf("evaluating script '%s': %v", script, FormatScriptError(err))
		}
	}
	err := context.setInitialProperties(build, properties)
	if err != nil {
		return nil, fmt.Errorf("evaluating properties: %v", err)
	}
	return context, nil
}

// NewThreadContext builds a context in a thread
func (context *Context) NewThreadContext(thread int, data interface{}) *Context {
	copy := context.Copy()
	copy.Thread = &thread
	copy.SetProperty("_thread", thread)
	copy.SetProperty("_data", data)
	return copy
}

func (context *Context) Copy() *Context {
	properties := make([]string, len(context.Properties))
	for i := 0; i < len(context.Properties); i++ {
		properties[i] = context.Properties[i]
	}
	environment := make(map[string]string)
	for name, value := range context.Environment {
		environment[name] = value
	}
	copy := Context{
		VM:          context.VM.Copy(),
		Properties:  properties,
		Environment: environment,
		Index:       context.Index.Copy(),
		Stack:       context.Stack.Copy(),
	}
	return &copy
}

// Set initial build properties
func (context *Context) setInitialProperties(build *Build, object util.Object) error {
	context.SetProperty("_OS", runtime.GOOS)
	context.SetProperty("_ARCH", runtime.GOARCH)
	context.SetProperty("_NCPU", runtime.NumCPU())
	context.SetProperty("_BASE", build.Dir)
	context.SetProperty("_HERE", build.Here)
	todo := object.Fields()
	var crash error
	for len(todo) > 0 {
		var done []string
		for _, name := range todo {
			value := object[name]
			eval, err := context.EvaluateObject(value)
			if err == nil {
				context.SetProperty(name, eval)
				done = append(done, name)
			} else {
				crash = err
			}
		}
		if len(done) == 0 {
			return crash
		}
		var next []string
		for _, name := range todo {
			found := false
			for _, n := range done {
				if name == n {
					found = true
				}
			}
			if !found {
				next = append(next, name)
			}
		}
		todo = next
	}
	return nil
}

// Set property with given to given value
func (context *Context) SetProperty(name string, value interface{}) {
	context.VM.Define(name, value)
}

// Get property value with given name
func (context *Context) GetProperty(name string) (interface{}, error) {
	value, err := context.VM.Get(name)
	if err != nil {
		return nil, err
	}
	return util.ValueToInterface(value), nil
}

// Evaluate given expression in context and return its value
func (context *Context) EvaluateExpression(source string) (interface{}, error) {
	value, err := context.VM.Execute(source)
	if err != nil {
		return nil, FormatScriptError(err)
	}
	return util.ValueToInterface(value), nil
}

// Evaluate a given object, that is replace '#{foo}' in strings with the value
// of property foo
func (context *Context) EvaluateObject(object interface{}) (interface{}, error) {
	switch value := object.(type) {
	case string:
		evaluated, err := context.EvaluateString(value)
		if err != nil {
			return nil, err
		}
		return evaluated, nil
	case bool:
		return value, nil
	case int:
		return value, nil
	case int32:
		return value, nil
	case int64:
		return value, nil
	case float64:
		return value, nil
	default:
		if value == nil {
			return nil, nil
		}
		switch reflect.TypeOf(object).Kind() {
		case reflect.Slice:
			slice := reflect.ValueOf(object)
			elements := make([]interface{}, slice.Len())
			for index := 0; index < slice.Len(); index++ {
				val, err := context.EvaluateObject(slice.Index(index).Interface())
				if err != nil {
					return nil, err
				}
				elements[index] = val
			}
			return elements, nil
		case reflect.Map:
			dict := reflect.ValueOf(object)
			elements := make(map[interface{}]interface{})
			for _, key := range dict.MapKeys() {
				keyEval, err := context.EvaluateObject(key.Interface())
				if err != nil {
					return nil, err
				}
				valueEval, err := context.EvaluateObject(dict.MapIndex(key).Interface())
				if err != nil {
					return nil, err
				}
				elements[keyEval] = valueEval
			}
			return elements, nil
		default:
			return nil, fmt.Errorf("no serializer for type '%T'", object)
		}
	}
}

// Evaluate a string by replacing '#{foo}' with value of property foo
func (context *Context) EvaluateString(text string) (string, error) {
	r := regexp.MustCompile(`#{.*?}`)
	var errors []error
	replaced := r.ReplaceAllStringFunc(text, func(expression string) string {
		name := expression[2 : len(expression)-1]
		var value interface{}
		value, err := context.EvaluateExpression(name)
		if err != nil {
			errors = append(errors, err)
			return ""
		} else {
			var str string
			str, err = PropertyToString(value, false)
			if err != nil {
				errors = append(errors, err)
				return ""
			} else {
				return str
			}
		}
	})
	if len(errors) > 0 {
		return replaced, errors[0]
	} else {
		return replaced, nil
	}
}

// Evaluate environment in context and return it as a slice of strings
func (context *Context) EvaluateEnvironment(build *Build) ([]string, error) {
	environment := make(map[string]string)
	for _, line := range os.Environ() {
		index := strings.Index(line, "=")
		name := line[:index]
		value := line[index+1:]
		environment[name] = value
	}
	environment["_BASE"] = build.Dir
	environment["_HERE"] = build.Here
	var variables []string
	for name := range context.Environment {
		variables = append(variables, name)
	}
	sort.Strings(variables)
	for _, name := range variables {
		value := context.Environment[name]
		r := regexp.MustCompile(`[$#]{.*?}`)
		replaced := r.ReplaceAllStringFunc(value, func(expression string) string {
			name := expression[2 : len(expression)-1]
			if expression[0:1] == "$" {
				value, ok := environment[name]
				if !ok {
					return expression
				} else {
					return value
				}
			} else {
				value, err := context.EvaluateExpression(name)
				if err != nil {
					return expression
				} else {
					str, _ := PropertyToString(value, false)
					return str
				}
			}
		})
		environment[name] = replaced
	}
	var lines []string
	for name, value := range environment {
		line := name + "=" + value
		lines = append(lines, line)
	}
	return lines, nil
}

// Run steps in context
func (context *Context) Run(steps []Step) error {
	context.Index.Expand()
	for index, step := range steps {
		context.Index.Set(index)
		err := step.Run(context)
		if err != nil {
			return err
		}
	}
	context.Index.Shrink()
	return nil
}

// Find files in the context:
// - dir: the search root directory
// - includes: the list of globs to include
// - excludes: the list of globs to exclude
// - folder: tells if we should include folders
// Return the list of files as a slice of strings
func (context *Context) FindFiles(dir string, includes, excludes []string, folder bool) ([]string, error) {
	dir, err := context.EvaluateString(dir)
	if err != nil {
		return nil, fmt.Errorf("evaluating source directory: %v", err)
	}
	var included []string
	for _, include := range includes {
		pattern, err := context.EvaluateString(include)
		if err != nil {
			return nil, fmt.Errorf("evaluating pattern: %v", err)
		}
		included = append(included, pattern)
	}
	var excluded []string
	for _, exclude := range excludes {
		pattern, err := context.EvaluateString(exclude)
		if err != nil {
			return nil, fmt.Errorf("evaluating pattern: %v", err)
		}
		excluded = append(excluded, pattern)
	}
	return util.FindFiles(dir, included, excluded, folder)
}

// Message print a message on the console
func (context *Context) Message(text string, args ...interface{}) {
	if context.Thread != nil {
		prefix := strconv.Itoa(*context.Thread) + "| "
		text = prefix + strings.Replace(text, "\n", "\n"+prefix, -1)
	}
	Message(text, args...)
}

// FormatScriptError adds line and column numbers on parser or vm errors.
func FormatScriptError(err error) error {
	if e, ok := err.(*parser.Error); ok {
		return fmt.Errorf("%s (at line %d, column %d)", err, e.Pos.Line, e.Pos.Column)
	} else if e, ok := err.(*vm.Error); ok {
		return fmt.Errorf("%s (at line %d, column %d)", err, e.Pos.Line, e.Pos.Column)
	} else {
		return err
	}
}
<commit_msg>Release 0.10.0: Threads and task replace change<commit_after>package build

import (
	"fmt"
	"io/ioutil"
	"neon/util"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"sort"
	"strings"

	anko_core "github.com/mattn/anko/builtins"
	"github.com/mattn/anko/parser"
	"github.com/mattn/anko/vm"
)

// Build context
type Context struct {
	VM          *vm.Env
	Properties  []string
	Environment map[string]string
	Index       *Index
	Stack       *Stack
}

// NewContext make a new build context
func NewContext(build *Build) (*Context, error) {
	v := vm.NewEnv()
	anko_core.LoadAllBuiltins(v)
	LoadBuiltins(v)
	properties := build.GetProperties()
	environment := build.GetEnvironment()
	context := &Context{
		VM:          v,
		Properties:  properties.Fields(),
		Environment: environment,
		Index:       NewIndex(),
		Stack:       NewStack(),
	}
	for _, script := range build.Scripts {
		source, err := ioutil.ReadFile(script)
		if err != nil {
			return nil, fmt.Errorf("reading script '%s': %v", script, err)
		}
		_, err = v.Execute(string(source))
		if err != nil {
			return nil, fmt.Errorf("evaluating script '%s': %v", script, FormatScriptError(err))
		}
	}
	err := context.setInitialProperties(build, properties)
	if err != nil {
		return nil, fmt.Errorf("evaluating properties: %v", err)
	}
	return context, nil
}

// NewThreadContext builds a context in a thread
func (context *Context) NewThreadContext(thread int, data interface{}) *Context {
	copy := context.Copy()
	copy.SetProperty("_thread", thread)
	copy.SetProperty("_data", data)
	return copy
}

func (context *Context) Copy() *Context {
	properties := make([]string, len(context.Properties))
	for i := 0; i < len(context.Properties); i++ {
		properties[i] = context.Properties[i]
	}
	environment := make(map[string]string)
	for name, value := range context.Environment {
		environment[name] = value
	}
	copy := Context{
		VM:          context.VM.Copy(),
		Properties:  properties,
		Environment: environment,
		Index:       context.Index.Copy(),
		Stack:       context.Stack.Copy(),
	}
	return &copy
}

// Set initial build properties
func (context *Context) setInitialProperties(build *Build, object util.Object) error {
	context.SetProperty("_OS", runtime.GOOS)
	context.SetProperty("_ARCH", runtime.GOARCH)
	context.SetProperty("_NCPU", runtime.NumCPU())
	context.SetProperty("_BASE", build.Dir)
	context.SetProperty("_HERE", build.Here)
	todo := object.Fields()
	var crash error
	for len(todo) > 0 {
		var done []string
		for _, name := range todo {
			value := object[name]
			eval, err := context.EvaluateObject(value)
			if err == nil {
				context.SetProperty(name, eval)
				done = append(done, name)
			} else {
				crash = err
			}
		}
		if len(done) == 0 {
			return crash
		}
		var next []string
		for _, name := range todo {
			found := false
			for _, n := range done {
				if name == n {
					found = true
				}
			}
			if !found {
				next = append(next, name)
			}
		}
		todo = next
	}
	return nil
}

// Set property with given to given value
func (context *Context) SetProperty(name string, value interface{}) {
	context.VM.Define(name, value)
}

// Get property value with given name
func (context *Context) GetProperty(name string) (interface{}, error) {
	value, err := context.VM.Get(name)
	if err != nil {
		return nil, err
	}
	return util.ValueToInterface(value), nil
}

// Evaluate given expression in context and return its value
func (context *Context) EvaluateExpression(source string) (interface{}, error) {
	value, err := context.VM.Execute(source)
	if err != nil {
		return nil, FormatScriptError(err)
	}
	return util.ValueToInterface(value), nil
}

// Evaluate a given object, that is replace '#{foo}' in strings with the value
// of property foo
func (context *Context) EvaluateObject(object interface{}) (interface{}, error) {
	switch value := object.(type) {
	case string:
		evaluated, err := context.EvaluateString(value)
		if err != nil {
			return nil, err
		}
		return evaluated, nil
	case bool:
		return value, nil
	case int:
		return value, nil
	case int32:
		return value, nil
	case int64:
		return value, nil
	case float64:
		return value, nil
	default:
		if value == nil {
			return nil, nil
		}
		switch reflect.TypeOf(object).Kind() {
		case reflect.Slice:
			slice := reflect.ValueOf(object)
			elements := make([]interface{}, slice.Len())
			for index := 0; index < slice.Len(); index++ {
				val, err := context.EvaluateObject(slice.Index(index).Interface())
				if err != nil {
					return nil, err
				}
				elements[index] = val
			}
			return elements, nil
		case reflect.Map:
			dict := reflect.ValueOf(object)
			elements := make(map[interface{}]interface{})
			for _, key := range dict.MapKeys() {
				keyEval, err := context.EvaluateObject(key.Interface())
				if err != nil {
					return nil, err
				}
				valueEval, err := context.EvaluateObject(dict.MapIndex(key).Interface())
				if err != nil {
					return nil, err
				}
				elements[keyEval] = valueEval
			}
			return elements, nil
		default:
			return nil, fmt.Errorf("no serializer for type '%T'", object)
		}
	}
}

// Evaluate a string by replacing '#{foo}' with value of property foo
func (context *Context) EvaluateString(text string) (string, error) {
	r := regexp.MustCompile(`#{.*?}`)
	var errors []error
	replaced := r.ReplaceAllStringFunc(text, func(expression string) string {
		name := expression[2: len(expression)-1]
		var value interface{}
		value, err := context.EvaluateExpression(name)
		if err != nil {
			errors = append(errors, err)
			return ""
		} else {
			var str string
			str, err = PropertyToString(value, false)
			if err != nil {
				errors = append(errors, err)
				return ""
			} else {
				return str
			}
		}
	})
	if len(errors) > 0 {
		return replaced, errors[0]
	} else {
		return replaced, nil
	}
}

// Evaluate environment in context and return it as a slice of strings
func (context *Context) EvaluateEnvironment(build *Build) ([]string, error) {
	environment := make(map[string]string)
	for _, line := range os.Environ() {
		index := strings.Index(line, "=")
		name := line[:index]
		value := line[index+1:]
		environment[name] = value
	}
	environment["_BASE"] = build.Dir
	environment["_HERE"] = build.Here
	var variables []string
	for name := range context.Environment {
		variables = append(variables, name)
	}
	sort.Strings(variables)
	for _, name := range variables {
		value := context.Environment[name]
		r := regexp.MustCompile(`[$#]{.*?}`)
		replaced := r.ReplaceAllStringFunc(value, func(expression string) string {
			name := expression[2: len(expression)-1]
			if expression[0:1] == "$" {
				value, ok := environment[name]
				if !ok {
					return expression
				} else {
					return value
				}
			} else {
				value, err := context.EvaluateExpression(name)
				if err != nil {
					return expression
				} else {
					str, _ := PropertyToString(value, false)
					return str
				}
			}
		})
		environment[name] = replaced
	}
	var lines []string
	for name, value := range environment {
		line := name + "=" + value
		lines = append(lines, line)
	}
	return lines, nil
}

// Run steps in context
func (context *Context) Run(steps []Step) error {
	context.Index.Expand()
	for index, step := range steps {
		context.Index.Set(index)
		err := step.Run(context)
		if err != nil {
			return err
		}
	}
	context.Index.Shrink()
	return nil
}

// Find files in the context:
// - dir: the search root directory
// - includes: the list of globs to include
// - excludes: the list of globs to exclude
// - folder: tells if we should include folders
// Return the list of files as a slice of strings
func (context *Context) FindFiles(dir string, includes, excludes []string, folder bool) ([]string, error) {
	dir, err := context.EvaluateString(dir)
	if err != nil {
		return nil, fmt.Errorf("evaluating source directory: %v", err)
	}
	var included []string
	for _, include := range includes {
		pattern, err := context.EvaluateString(include)
		if err != nil {
			return nil, fmt.Errorf("evaluating pattern: %v", err)
		}
		included = append(included, pattern)
	}
	var excluded []string
	for _, exclude := range excludes {
		pattern, err := context.EvaluateString(exclude)
		if err != nil {
			return nil, fmt.Errorf("evaluating pattern: %v", err)
		}
		excluded = append(excluded, pattern)
	}
	return util.FindFiles(dir, included, excluded, folder)
}

// Message print a message on the console
func (context *Context) Message(text string, args ...interface{}) {
	Message(text, args...)
}

// FormatScriptError adds line and column numbers on parser or vm errors.
func FormatScriptError(err error) error {
	if e, ok := err.(*parser.Error); ok {
		return fmt.Errorf("%s (at line %d, column %d)", err, e.Pos.Line, e.Pos.Column)
	} else if e, ok := err.(*vm.Error); ok {
		return fmt.Errorf("%s (at line %d, column %d)", err, e.Pos.Line, e.Pos.Column)
	} else {
		return err
	}
}
<|endoftext|>


<commit_before>package admin

import (
	"errors"
	"fmt"

	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/args"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/with"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/billing"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	billingMethods "github.com/BytemarkHosting/bytemark-client/lib/requests/billing"
	brainMethods "github.com/BytemarkHosting/bytemark-client/lib/requests/brain"
	"github.com/BytemarkHosting/bytemark-client/util/log"
	"github.com/urfave/cli"
)

func init() {
	readUpdateFlags := func(c *app.Context) (usageStrategy *string, overcommitRatio *int, label *string) {
		if c.Context.IsSet("usage-strategy") {
			v := c.String("usage-strategy")
			usageStrategy = &v
		}

		if c.Context.IsSet("overcommit-ratio") {
			v := c.Int("overcommit-ratio")
			overcommitRatio = &v
		}

		if c.Context.IsSet("label") {
			v := c.String("label")
			label = &v
		}

		return
	}

	Commands = append(Commands, cli.Command{
		Name:   "update",
		Action: cli.ShowSubcommandHelp,
		Subcommands: []cli.Command{
			{
				Name:      "billing-definition",
				Usage:     "update a bmbilling definition",
				UsageText: "bytemark --admin update billing-definition [flags] [name] [value]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "name",
						Usage: "the name of the definition to set",
					},
					cli.StringFlag{
						Name:  "value",
						Usage: "the value of the definition to set",
					},
					cli.StringFlag{
						Name:  "group",
						Usage: "the group a user must be in to update the definition",
					},
				},
				Action: app.Action(args.Optional("name", "value"), with.RequiredFlags("name", "value"), with.Auth, func(ctx *app.Context) error {
					def := billing.Definition{
						Name:           ctx.String("name"),
						Value:          ctx.String("value"),
						UpdateGroupReq: ctx.String("group"),
					}
					if _, err := billingMethods.GetDefinition(ctx.Client(), def.Name); err != nil {
						if _, ok := err.(lib.NotFoundError); ok {
							ctx.LogErr("Couldn't find a definition called %s - aborting.", def.Name)
							return nil
						}
						return err
					}
					err := billingMethods.UpdateDefinition(ctx.Client(), def)
					if err == nil {
						ctx.LogErr("Updated %s to %s", def.Name, def.Value)
					}
					return err

				}),
			}, {
				Name:      "head",
				Usage:     "update the settings of a head",
				UsageText: "bytemark --admin update head <head> [--usage-strategy] [--overcommit-ratio] [--label]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "head",
						Usage: "the ID or label of the head to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the head",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the head",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the head",
					},
				},
				Action: app.Action(args.Optional("head", "usage-strategy", "overcommit-ratio", "label"), with.RequiredFlags("head"), with.Auth, func(c *app.Context) error {
					usageStrategy, overcommitRatio, label := readUpdateFlags(c)

					options := lib.UpdateHead{
						UsageStrategy:   usageStrategy,
						OvercommitRatio: overcommitRatio,
						Label:           label,
					}

					if err := c.Client().UpdateHead(c.String("head"), options); err != nil {
						return err
					}

					log.Outputf("Head %s updated\n", c.String("head"))

					return nil
				}),
			},
			{
				Name:      "tail",
				Usage:     "update the settings of a tail",
				UsageText: "bytemark --admin update tail <tail> [--usage-strategy] [--overcommit-ratio] [--label]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "tail",
						Usage: "the ID or label of the tail to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the tail",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the tail",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the tail",
					},
				},
				Action: app.Action(args.Optional("tail", "usage-strategy", "overcommit-ratio", "label"), with.RequiredFlags("tail"), with.Auth, func(c *app.Context) error {
					usageStrategy, overcommitRatio, label := readUpdateFlags(c)

					options := lib.UpdateTail{
						UsageStrategy:   usageStrategy,
						OvercommitRatio: overcommitRatio,
						Label:           label,
					}

					if err := c.Client().UpdateTail(c.String("tail"), options); err != nil {
						return err
					}

					log.Outputf("Tail %s updated\n", c.String("tail"))

					return nil
				}),
			},
			{
				Name:      "storage pool",
				Usage:     "update the settings of a storage pool",
				UsageText: "bytemark --admin update storage pool [--usage-strategy new-strategy] [--overcommit-ratio new-ratio] [--label new-label] [--migration-concurrency new-limit] <storage pool>",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "storage-pool",
						Usage: "the ID or label of the storage pool to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the storage pool",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the storage pool",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the storage pool",
					},
					cli.IntFlag{
						Name:  "migration-concurrency",
						Usage: "the number of concurrent migrations the storage pool can handle",
					},
				},
				Action: app.Action(args.Optional("storage-pool", "usage-strategy", "overcommit-ratio"), with.RequiredFlags("storage-pool"), with.Auth, func(c *app.Context) error {

					options := brain.StoragePool{
						UsageStrategy:        c.String("usage-strategy"),
						OvercommitRatio:      c.Int("overcommit-ratio"),
						Label:                c.String("label"),
						MigrationConcurrency: c.Int("migration-concurrency"),
					}

					if err := c.Client().UpdateStoragePool(c.String("storage-pool"), options); err != nil {
						return err
					}

					log.Outputf("Storage pool %s updated\n", c.String("storage-pool"))

					return nil
				}),
			},
			{
				Name:    "server",
				Aliases: []string{"vm"},
				Action:  cli.ShowSubcommandHelp,
				Subcommands: []cli.Command{
					{
						Name:        "migration",
						Usage:       "update the settings of an in-progress migration",
						UsageText:   "bytemark --admin update server migration <name> [--migrate-speed] [--migrate-downtime]",
						Description: `This command migrates a server to a new head. If a new head isn't supplied, a new one is picked automatically.`,
						Flags: []cli.Flag{
							cli.GenericFlag{
								Name:  "server",
								Usage: "the server to migrate",
								Value: new(app.VirtualMachineNameFlag),
							},
							cli.Int64Flag{
								Name:  "migrate-speed",
								Usage: "the max speed to migrate the server at",
							},
							cli.IntFlag{
								Name:  "migrate-downtime",
								Usage: "the max allowed downtime",
							},
						},
						Action: app.Action(args.Optional("server", "migrate-speed", "migrate-downtime"), with.RequiredFlags("server"), with.Auth, func(c *app.Context) error {
							vm := c.VirtualMachineName("server")

							var speed *int64
							var downtime *int

							if c.Context.IsSet("migrate-speed") {
								s := c.Int64("migrate-speed")
								speed = &s
							}
							if c.Context.IsSet("migrate-downtime") {
								d := c.Int("migrate-downtime")
								downtime = &d
							}

							if speed == nil && downtime == nil {
								return errors.New("Nothing to update")
							}

							if err := c.Client().UpdateVMMigration(vm, speed, downtime); err != nil {
								return err
							}

							log.Outputf("Migration for server %s updated\n", vm.String())

							return nil
						}),
					},
				},
			},
			{
				Name:        "migration",
				Usage:       "update a migration",
				UsageText:   "bytemark --admin update migration --id <id> --priority <priority> --cancel-disc <disc> --cancel-pool <pool> --cancel-tail <tail> | --cancel-all",
				Description: `This command allows you to update an ongoing migration job by altering its priority, cancelling migrating discs, pools, tails, or canceling everything for the current job`,
				Flags: []cli.Flag{
					cli.IntFlag{
						Name:  "id",
						Usage: "the id of the migration job",
					},
					cli.IntFlag{
						Name:  "priority",
						Usage: "the priority of the current job",
					},
					cli.StringSliceFlag{
						Name:  "cancel-disc",
						Usage: "the disc(s) to cancel migration of",
					},
					cli.StringSliceFlag{
						Name:  "cancel-pool",
						Usage: "the pool(s) to cancel migration of",
					},
					cli.StringSliceFlag{
						Name:  "cancel-tail",
						Usage: "the tail(s) to cancel migration of",
					},
					cli.BoolFlag{
						Name:  "cancel-all",
						Usage: "cancel the all migrations of the job",
					},
				},
				Action: app.Action(with.RequiredFlags("id"), with.Auth, func(c *app.Context) error {
					// read all flags?
					id := c.Context.Int("id")
					discs := c.Context.StringSlice("cancel-disc")
					pools := c.Context.StringSlice("cancel-pool")
					tails := c.Context.StringSlice("cancel-tail")

					totalLength := append(discs, pools...)
					totalLength = append(totalLength, tails...)

					// check the presence of cancel all first
					if c.Context.IsSet("cancel-all") {
						if len(totalLength) > 0 {

							return fmt.Errorf("You have set additional flags as well as --cancel-all. Nothing else can be specified when --cancel-all has been set.")
						}

						err := brainMethods.CancelMigrationJob(c.Client(), id)
						if err != nil {
							return err
						}
						c.LogErr("All migrations for job %d have been cancelled.", id)

						// make the call to the cancel all migration.
					}

					modifications := brain.MigrationJobModification{
						Cancel: brain.MigrationJobLocations{
							Discs: stringsToJsonNumbers(discs),
							Pools: stringsToJsonNumbers(pools),
							Tails: stringsToJsonNumbers(tails),
						},
						Options: brain.MigrationJobOptions{
							Priority: c.Context.Int("priority"),
						},
					}

					err := brainMethods.EditMigrationJob(c.Client(), id, modifications)
					if err != nil {
						return err
					}

					// if we have set a new priority, show a message.
					c.LogErr("Priority updated for Job %d", id)

					// not really sure if we should be showing a confirmation message for all individual things that have been cancelled.
					return err
				}),
			},
		},
	})
}
<commit_msg>remove unnecessary comments. display a confirmation message for all things that have been cancelled.<commit_after>package admin

import (
	"errors"
	"fmt"

	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/args"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/with"
	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/billing"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	billingMethods "github.com/BytemarkHosting/bytemark-client/lib/requests/billing"
	brainMethods "github.com/BytemarkHosting/bytemark-client/lib/requests/brain"
	"github.com/BytemarkHosting/bytemark-client/util/log"
	"github.com/urfave/cli"
)

func init() {
	readUpdateFlags := func(c *app.Context) (usageStrategy *string, overcommitRatio *int, label *string) {
		if c.Context.IsSet("usage-strategy") {
			v := c.String("usage-strategy")
			usageStrategy = &v
		}

		if c.Context.IsSet("overcommit-ratio") {
			v := c.Int("overcommit-ratio")
			overcommitRatio = &v
		}

		if c.Context.IsSet("label") {
			v := c.String("label")
			label = &v
		}

		return
	}

	Commands = append(Commands, cli.Command{
		Name:   "update",
		Action: cli.ShowSubcommandHelp,
		Subcommands: []cli.Command{
			{
				Name:      "billing-definition",
				Usage:     "update a bmbilling definition",
				UsageText: "bytemark --admin update billing-definition [flags] [name] [value]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "name",
						Usage: "the name of the definition to set",
					},
					cli.StringFlag{
						Name:  "value",
						Usage: "the value of the definition to set",
					},
					cli.StringFlag{
						Name:  "group",
						Usage: "the group a user must be in to update the definition",
					},
				},
				Action: app.Action(args.Optional("name", "value"), with.RequiredFlags("name", "value"), with.Auth, func(ctx *app.Context) error {
					def := billing.Definition{
						Name:           ctx.String("name"),
						Value:          ctx.String("value"),
						UpdateGroupReq: ctx.String("group"),
					}
					if _, err := billingMethods.GetDefinition(ctx.Client(), def.Name); err != nil {
						if _, ok := err.(lib.NotFoundError); ok {
							ctx.LogErr("Couldn't find a definition called %s - aborting.", def.Name)
							return nil
						}
						return err
					}
					err := billingMethods.UpdateDefinition(ctx.Client(), def)
					if err == nil {
						ctx.LogErr("Updated %s to %s", def.Name, def.Value)
					}
					return err

				}),
			}, {
				Name:      "head",
				Usage:     "update the settings of a head",
				UsageText: "bytemark --admin update head <head> [--usage-strategy] [--overcommit-ratio] [--label]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "head",
						Usage: "the ID or label of the head to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the head",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the head",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the head",
					},
				},
				Action: app.Action(args.Optional("head", "usage-strategy", "overcommit-ratio", "label"), with.RequiredFlags("head"), with.Auth, func(c *app.Context) error {
					usageStrategy, overcommitRatio, label := readUpdateFlags(c)

					options := lib.UpdateHead{
						UsageStrategy:   usageStrategy,
						OvercommitRatio: overcommitRatio,
						Label:           label,
					}

					if err := c.Client().UpdateHead(c.String("head"), options); err != nil {
						return err
					}

					log.Outputf("Head %s updated\n", c.String("head"))

					return nil
				}),
			},
			{
				Name:      "tail",
				Usage:     "update the settings of a tail",
				UsageText: "bytemark --admin update tail <tail> [--usage-strategy] [--overcommit-ratio] [--label]",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "tail",
						Usage: "the ID or label of the tail to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the tail",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the tail",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the tail",
					},
				},
				Action: app.Action(args.Optional("tail", "usage-strategy", "overcommit-ratio", "label"), with.RequiredFlags("tail"), with.Auth, func(c *app.Context) error {
					usageStrategy, overcommitRatio, label := readUpdateFlags(c)

					options := lib.UpdateTail{
						UsageStrategy:   usageStrategy,
						OvercommitRatio: overcommitRatio,
						Label:           label,
					}

					if err := c.Client().UpdateTail(c.String("tail"), options); err != nil {
						return err
					}

					log.Outputf("Tail %s updated\n", c.String("tail"))

					return nil
				}),
			},
			{
				Name:      "storage pool",
				Usage:     "update the settings of a storage pool",
				UsageText: "bytemark --admin update storage pool [--usage-strategy new-strategy] [--overcommit-ratio new-ratio] [--label new-label] [--migration-concurrency new-limit] <storage pool>",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "storage-pool",
						Usage: "the ID or label of the storage pool to be updated",
					},
					cli.StringFlag{
						Name:  "usage-strategy",
						Usage: "the usage strategy of the storage pool",
					},
					cli.IntFlag{
						Name:  "overcommit-ratio",
						Usage: "the overcommit ratio of the storage pool",
					},
					cli.StringFlag{
						Name:  "label",
						Usage: "the label of the storage pool",
					},
					cli.IntFlag{
						Name:  "migration-concurrency",
						Usage: "the number of concurrent migrations the storage pool can handle",
					},
				},
				Action: app.Action(args.Optional("storage-pool", "usage-strategy", "overcommit-ratio"), with.RequiredFlags("storage-pool"), with.Auth, func(c *app.Context) error {

					options := brain.StoragePool{
						UsageStrategy:        c.String("usage-strategy"),
						OvercommitRatio:      c.Int("overcommit-ratio"),
						Label:                c.String("label"),
						MigrationConcurrency: c.Int("migration-concurrency"),
					}

					if err := c.Client().UpdateStoragePool(c.String("storage-pool"), options); err != nil {
						return err
					}

					log.Outputf("Storage pool %s updated\n", c.String("storage-pool"))

					return nil
				}),
			},
			{
				Name:    "server",
				Aliases: []string{"vm"},
				Action:  cli.ShowSubcommandHelp,
				Subcommands: []cli.Command{
					{
						Name:        "migration",
						Usage:       "update the settings of an in-progress migration",
						UsageText:   "bytemark --admin update server migration <name> [--migrate-speed] [--migrate-downtime]",
						Description: `This command migrates a server to a new head. If a new head isn't supplied, a new one is picked automatically.`,
						Flags: []cli.Flag{
							cli.GenericFlag{
								Name:  "server",
								Usage: "the server to migrate",
								Value: new(app.VirtualMachineNameFlag),
							},
							cli.Int64Flag{
								Name:  "migrate-speed",
								Usage: "the max speed to migrate the server at",
							},
							cli.IntFlag{
								Name:  "migrate-downtime",
								Usage: "the max allowed downtime",
							},
						},
						Action: app.Action(args.Optional("server", "migrate-speed", "migrate-downtime"), with.RequiredFlags("server"), with.Auth, func(c *app.Context) error {
							vm := c.VirtualMachineName("server")

							var speed *int64
							var downtime *int

							if c.Context.IsSet("migrate-speed") {
								s := c.Int64("migrate-speed")
								speed = &s
							}
							if c.Context.IsSet("migrate-downtime") {
								d := c.Int("migrate-downtime")
								downtime = &d
							}

							if speed == nil && downtime == nil {
								return errors.New("Nothing to update")
							}

							if err := c.Client().UpdateVMMigration(vm, speed, downtime); err != nil {
								return err
							}

							log.Outputf("Migration for server %s updated\n", vm.String())

							return nil
						}),
					},
				},
			},
			{
				Name:        "migration",
				Usage:       "update a migration",
				UsageText:   "bytemark --admin update migration --id <id> --priority <priority> --cancel-disc <disc> --cancel-pool <pool> --cancel-tail <tail> | --cancel-all",
				Description: `This command allows you to update an ongoing migration job by altering its priority, cancelling migrating discs, pools, tails, or canceling everything for the current job`,
				Flags: []cli.Flag{
					cli.IntFlag{
						Name:  "id",
						Usage: "the id of the migration job",
					},
					cli.IntFlag{
						Name:  "priority",
						Usage: "the priority of the current job",
					},
					cli.StringSliceFlag{
						Name:  "cancel-disc",
						Usage: "the disc(s) to cancel migration of",
					},
					cli.StringSliceFlag{
						Name:  "cancel-pool",
						Usage: "the pool(s) to cancel migration of",
					},
					cli.StringSliceFlag{
						Name:  "cancel-tail",
						Usage: "the tail(s) to cancel migration of",
					},
					cli.BoolFlag{
						Name:  "cancel-all",
						Usage: "cancel the all migrations of the job",
					},
				},
				Action: app.Action(with.RequiredFlags("id"), with.Auth, func(c *app.Context) error {
					id := c.Context.Int("id")
					discs := c.Context.StringSlice("cancel-disc")
					pools := c.Context.StringSlice("cancel-pool")
					tails := c.Context.StringSlice("cancel-tail")

					allCancelled := append(discs, pools...)
					allCancelled = append(allCancelled, tails...)

					if c.Context.IsSet("cancel-all") {
						if len(allCancelled) > 0 {
							return fmt.Errorf("You have set additional flags as well as --cancel-all. Nothing else can be specified when --cancel-all has been set.")
						}

						err := brainMethods.CancelMigrationJob(c.Client(), id)
						if err != nil {
							return err
						}
						c.LogErr("All migrations for job %d have been cancelled.", id)
						return err
					}

					modifications := brain.MigrationJobModification{
						Cancel: brain.MigrationJobLocations{
							Discs: stringsToJsonNumbers(discs),
							Pools: stringsToJsonNumbers(pools),
							Tails: stringsToJsonNumbers(tails),
						},
						Options: brain.MigrationJobOptions{
							Priority: c.Context.Int("priority"),
						},
					}

					err := brainMethods.EditMigrationJob(c.Client(), id, modifications)
					if err != nil {
						return err
					}

					if c.Context.IsSet("priority") {
						c.LogErr("Priority updated for job %d", id)
					}

					for _, cancelled := range allCancelled {
						c.LogErr("Migration cancelled for %s on job %d", cancelled, id)
					}

					return err
				}),
			},
		},
	})
}
<|endoftext|>


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

package os_test

import (
	"fmt"
	"internal/testenv"
	"os"
	osexec "os/exec"
	"path/filepath"
	"runtime"
	"testing"
)

const executable_EnvVar = "OSTEST_OUTPUT_EXECPATH"

func TestExecutable(t *testing.T) {
	testenv.MustHaveExec(t)
	ep, err := os.Executable()
	if err != nil {
		t.Fatalf("Executable failed: %v", err)
	}
	// we want fn to be of the form "dir/prog"
	dir := filepath.Dir(filepath.Dir(ep))
	fn, err := filepath.Rel(dir, ep)
	if err != nil {
		t.Fatalf("filepath.Rel: %v", err)
	}

	cmd := &osexec.Cmd{}
	// make child start with a relative program path
	cmd.Dir = dir
	cmd.Path = fn
	// forge argv[0] for child, so that we can verify we could correctly
	// get real path of the executable without influenced by argv[0].
	cmd.Args = []string{"-", "-test.run=XXXX"}
	if runtime.GOOS == "openbsd" || runtime.GOOS == "aix" {
		// OpenBSD and AIX rely on argv[0]
		cmd.Args[0] = fn
	}
	cmd.Env = append(os.Environ(), fmt.Sprintf("%s=1", executable_EnvVar))
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("exec(self) failed: %v", err)
	}
	outs := string(out)
	if !filepath.IsAbs(outs) {
		t.Fatalf("Child returned %q, want an absolute path", out)
	}
	if !sameFile(outs, ep) {
		t.Fatalf("Child returned %q, not the same file as %q", out, ep)
	}
}

func sameFile(fn1, fn2 string) bool {
	fi1, err := os.Stat(fn1)
	if err != nil {
		return false
	}
	fi2, err := os.Stat(fn2)
	if err != nil {
		return false
	}
	return os.SameFile(fi1, fi2)
}

func init() {
	if e := os.Getenv(executable_EnvVar); e != "" {
		// first chdir to another path
		dir := "/"
		if runtime.GOOS == "windows" {
			cwd, err := os.Getwd()
			if err != nil {
				panic(err)
			}
			dir = filepath.VolumeName(cwd)
		}
		os.Chdir(dir)
		if ep, err := os.Executable(); err != nil {
			fmt.Fprint(os.Stderr, "ERROR: ", err)
		} else {
			fmt.Fprint(os.Stderr, ep)
		}
		os.Exit(0)
	}
}

func TestExecutableDeleted(t *testing.T) {
	testenv.MustHaveExec(t)
	switch runtime.GOOS {
	case "windows":
		t.Skip("windows does not support deleting running binary")
	case "openbsd", "freebsd", "aix":
		t.Skipf("%v does not support reading deleted binary name", runtime.GOOS)
	}

	dir := t.TempDir()

	src := filepath.Join(dir, "testdel.go")
	exe := filepath.Join(dir, "testdel.exe")

	err := os.WriteFile(src, []byte(testExecutableDeletion), 0666)
	if err != nil {
		t.Fatal(err)
	}

	out, err := osexec.Command(testenv.GoToolPath(t), "build", "-o", exe, src).CombinedOutput()
	t.Logf("build output:\n%s", out)
	if err != nil {
		t.Fatal(err)
	}

	out, err = osexec.Command(exe).CombinedOutput()
	t.Logf("exec output:\n%s", out)
	if err != nil {
		t.Fatal(err)
	}
}

const testExecutableDeletion = `package main

import (
	"fmt"
	"os"
)

func main() {
	before, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read executable name before deletion: %v\n", err)
		os.Exit(1)
	}

	err = os.Remove(before)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to remove executable: %v\n", err)
		os.Exit(1)
	}

	after, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read executable name after deletion: %v\n", err)
		os.Exit(1)
	}

	if before != after {
		fmt.Fprintf(os.Stderr, "before and after do not match: %v != %v\n", before, after)
		os.Exit(1)
	}
}
`
<commit_msg>os: skip TestExecutableDeleted on plan9<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package os_test

import (
	"fmt"
	"internal/testenv"
	"os"
	osexec "os/exec"
	"path/filepath"
	"runtime"
	"testing"
)

const executable_EnvVar = "OSTEST_OUTPUT_EXECPATH"

func TestExecutable(t *testing.T) {
	testenv.MustHaveExec(t)
	ep, err := os.Executable()
	if err != nil {
		t.Fatalf("Executable failed: %v", err)
	}
	// we want fn to be of the form "dir/prog"
	dir := filepath.Dir(filepath.Dir(ep))
	fn, err := filepath.Rel(dir, ep)
	if err != nil {
		t.Fatalf("filepath.Rel: %v", err)
	}

	cmd := &osexec.Cmd{}
	// make child start with a relative program path
	cmd.Dir = dir
	cmd.Path = fn
	// forge argv[0] for child, so that we can verify we could correctly
	// get real path of the executable without influenced by argv[0].
	cmd.Args = []string{"-", "-test.run=XXXX"}
	if runtime.GOOS == "openbsd" || runtime.GOOS == "aix" {
		// OpenBSD and AIX rely on argv[0]
		cmd.Args[0] = fn
	}
	cmd.Env = append(os.Environ(), fmt.Sprintf("%s=1", executable_EnvVar))
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("exec(self) failed: %v", err)
	}
	outs := string(out)
	if !filepath.IsAbs(outs) {
		t.Fatalf("Child returned %q, want an absolute path", out)
	}
	if !sameFile(outs, ep) {
		t.Fatalf("Child returned %q, not the same file as %q", out, ep)
	}
}

func sameFile(fn1, fn2 string) bool {
	fi1, err := os.Stat(fn1)
	if err != nil {
		return false
	}
	fi2, err := os.Stat(fn2)
	if err != nil {
		return false
	}
	return os.SameFile(fi1, fi2)
}

func init() {
	if e := os.Getenv(executable_EnvVar); e != "" {
		// first chdir to another path
		dir := "/"
		if runtime.GOOS == "windows" {
			cwd, err := os.Getwd()
			if err != nil {
				panic(err)
			}
			dir = filepath.VolumeName(cwd)
		}
		os.Chdir(dir)
		if ep, err := os.Executable(); err != nil {
			fmt.Fprint(os.Stderr, "ERROR: ", err)
		} else {
			fmt.Fprint(os.Stderr, ep)
		}
		os.Exit(0)
	}
}

func TestExecutableDeleted(t *testing.T) {
	testenv.MustHaveExec(t)
	switch runtime.GOOS {
	case "windows", "plan9":
		t.Skipf("%v does not support deleting running binary", runtime.GOOS)
	case "openbsd", "freebsd", "aix":
		t.Skipf("%v does not support reading deleted binary name", runtime.GOOS)
	}

	dir := t.TempDir()

	src := filepath.Join(dir, "testdel.go")
	exe := filepath.Join(dir, "testdel.exe")

	err := os.WriteFile(src, []byte(testExecutableDeletion), 0666)
	if err != nil {
		t.Fatal(err)
	}

	out, err := osexec.Command(testenv.GoToolPath(t), "build", "-o", exe, src).CombinedOutput()
	t.Logf("build output:\n%s", out)
	if err != nil {
		t.Fatal(err)
	}

	out, err = osexec.Command(exe).CombinedOutput()
	t.Logf("exec output:\n%s", out)
	if err != nil {
		t.Fatal(err)
	}
}

const testExecutableDeletion = `package main

import (
	"fmt"
	"os"
)

func main() {
	before, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read executable name before deletion: %v\n", err)
		os.Exit(1)
	}

	err = os.Remove(before)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to remove executable: %v\n", err)
		os.Exit(1)
	}

	after, err := os.Executable()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to read executable name after deletion: %v\n", err)
		os.Exit(1)
	}

	if before != after {
		fmt.Fprintf(os.Stderr, "before and after do not match: %v != %v\n", before, after)
		os.Exit(1)
	}
}
`
<|endoftext|>


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

// This file implements the Check function, which typechecks a package.

package types

import (
	"fmt"
	"go/ast"
	"go/scanner"
	"go/token"
	"sort"
)

// enable for debugging
const trace = false

type checker struct {
	ctxt  *Context
	fset  *token.FileSet
	files []*ast.File

	// lazily initialized
	pkgscope  *ast.Scope
	firsterr  error
	initexprs map[*ast.ValueSpec][]ast.Expr // "inherited" initialization expressions for constant declarations
	funclist  []function                    // list of functions/methods with correct signatures and non-empty bodies
	funcsig   *Signature                    // signature of currently typechecked function
	pos       []token.Pos                   // stack of expr positions; debugging support, used if trace is set
}

type function struct {
	obj  *ast.Object // for debugging/tracing only
	sig  *Signature
	body *ast.BlockStmt
}

// later adds a function with non-empty body to the list of functions
// that need to be processed after all package-level declarations
// are typechecked.
//
func (check *checker) later(obj *ast.Object, sig *Signature, body *ast.BlockStmt) {
	// functions implemented elsewhere (say in assembly) have no body
	if body != nil {
		check.funclist = append(check.funclist, function{obj, sig, body})
	}
}

// declare declares an object of the given kind and name (ident) in scope;
// decl is the corresponding declaration in the AST. An error is reported
// if the object was declared before.
//
// TODO(gri) This is very similar to the declare function in go/parser; it
// is only used to associate methods with their respective receiver base types.
// In a future version, it might be simpler and cleaner to do all the resolution
// in the type-checking phase. It would simplify the parser, AST, and also
// reduce some amount of code duplication.
//
func (check *checker) declare(scope *ast.Scope, kind ast.ObjKind, ident *ast.Ident, decl ast.Decl) {
	assert(ident.Obj == nil) // identifier already declared or resolved
	obj := ast.NewObj(kind, ident.Name)
	obj.Decl = decl
	ident.Obj = obj
	if ident.Name != "_" {
		if alt := scope.Insert(obj); alt != nil {
			prevDecl := ""
			if pos := alt.Pos(); pos.IsValid() {
				prevDecl = fmt.Sprintf("\n\tprevious declaration at %s", check.fset.Position(pos))
			}
			check.errorf(ident.Pos(), fmt.Sprintf("%s redeclared in this block%s", ident.Name, prevDecl))
		}
	}
}

func (check *checker) valueSpec(pos token.Pos, obj *ast.Object, lhs []*ast.Ident, typ ast.Expr, rhs []ast.Expr, iota int) {
	if len(lhs) == 0 {
		check.invalidAST(pos, "missing lhs in declaration")
		return
	}

	// determine type for all of lhs, if any
	// (but only set it for the object we typecheck!)
	var t Type
	if typ != nil {
		t = check.typ(typ, false)
	}

	// len(lhs) > 0
	if len(lhs) == len(rhs) {
		// check only lhs and rhs corresponding to obj
		var l, r ast.Expr
		for i, name := range lhs {
			if name.Obj == obj {
				l = lhs[i]
				r = rhs[i]
				break
			}
		}
		assert(l != nil)
		obj.Type = t
		check.assign1to1(l, r, nil, true, iota)
		return
	}

	// there must be a type or initialization expressions
	if t == nil && len(rhs) == 0 {
		check.invalidAST(pos, "missing type or initialization expression")
		t = Typ[Invalid]
	}

	// if we have a type, mark all of lhs
	if t != nil {
		for _, name := range lhs {
			name.Obj.Type = t
		}
	}

	// check initial values, if any
	if len(rhs) > 0 {
		// TODO(gri) should try to avoid this conversion
		lhx := make([]ast.Expr, len(lhs))
		for i, e := range lhs {
			lhx[i] = e
		}
		check.assignNtoM(lhx, rhs, true, iota)
	}
}

// object typechecks an object by assigning it a type; obj.Type must be nil.
// Callers must check obj.Type before calling object; this eliminates a call
// for each identifier that has been typechecked already, a common scenario.
//
func (check *checker) object(obj *ast.Object, cycleOk bool) {
	assert(obj.Type == nil)

	switch obj.Kind {
	case ast.Bad, ast.Pkg:
		// nothing to do

	case ast.Con, ast.Var:
		// The obj.Data field for constants and variables is initialized
		// to the respective (hypothetical, for variables) iota value by
		// the parser. The object's fields can be in one of the following
		// states:
		// Type != nil  =>  the constant value is Data
		// Type == nil  =>  the object is not typechecked yet, and Data can be:
		// Data is int  =>  Data is the value of iota for this declaration
		// Data == nil  =>  the object's expression is being evaluated
		if obj.Data == nil {
			check.errorf(obj.Pos(), "illegal cycle in initialization of %s", obj.Name)
			obj.Type = Typ[Invalid]
			return
		}
		spec := obj.Decl.(*ast.ValueSpec)
		iota := obj.Data.(int)
		obj.Data = nil
		// determine initialization expressions
		values := spec.Values
		if len(values) == 0 && obj.Kind == ast.Con {
			values = check.initexprs[spec]
		}
		check.valueSpec(spec.Pos(), obj, spec.Names, spec.Type, values, iota)

	case ast.Typ:
		typ := &NamedType{Obj: obj}
		obj.Type = typ // "mark" object so recursion terminates
		typ.Underlying = underlying(check.typ(obj.Decl.(*ast.TypeSpec).Type, cycleOk))
		// typecheck associated method signatures
		if obj.Data != nil {
			scope := obj.Data.(*ast.Scope)
			switch t := typ.Underlying.(type) {
			case *Struct:
				// struct fields must not conflict with methods
				for _, f := range t.Fields {
					if m := scope.Lookup(f.Name); m != nil {
						check.errorf(m.Pos(), "type %s has both field and method named %s", obj.Name, f.Name)
						// ok to continue
					}
				}
			case *Interface:
				// methods cannot be associated with an interface type
				for _, m := range scope.Objects {
					recv := m.Decl.(*ast.FuncDecl).Recv.List[0].Type
					check.errorf(recv.Pos(), "invalid receiver type %s (%s is an interface type)", obj.Name, obj.Name)
					// ok to continue
				}
			}
			// typecheck method signatures
			for _, obj := range scope.Objects {
				mdecl := obj.Decl.(*ast.FuncDecl)
				sig := check.typ(mdecl.Type, cycleOk).(*Signature)
				params, _ := check.collectParams(mdecl.Recv, false)
				sig.Recv = params[0] // the parser/assocMethod ensure there is exactly one parameter
				obj.Type = sig
				check.later(obj, sig, mdecl.Body)
			}
		}

	case ast.Fun:
		fdecl := obj.Decl.(*ast.FuncDecl)
		// methods are typechecked when their receivers are typechecked
		if fdecl.Recv == nil {
			sig := check.typ(fdecl.Type, cycleOk).(*Signature)
			if obj.Name == "init" && (len(sig.Params) != 0 || len(sig.Results) != 0) {
				check.errorf(fdecl.Pos(), "func init must have no arguments and no return values")
				// ok to continue
			}
			obj.Type = sig
			check.later(obj, sig, fdecl.Body)
		}

	default:
		panic("unreachable")
	}
}

// assocInitvals associates "inherited" initialization expressions
// with the corresponding *ast.ValueSpec in the check.initexprs map
// for constant declarations without explicit initialization expressions.
//
func (check *checker) assocInitvals(decl *ast.GenDecl) {
	var values []ast.Expr
	for _, s := range decl.Specs {
		if s, ok := s.(*ast.ValueSpec); ok {
			if len(s.Values) > 0 {
				values = s.Values
			} else {
				check.initexprs[s] = values
			}
		}
	}
	if len(values) == 0 {
		check.invalidAST(decl.Pos(), "no initialization values provided")
	}
}

// assocMethod associates a method declaration with the respective
// receiver base type. meth.Recv must exist.
//
func (check *checker) assocMethod(meth *ast.FuncDecl) {
	// The receiver type is one of the following (enforced by parser):
	// - *ast.Ident
	// - *ast.StarExpr{*ast.Ident}
	// - *ast.BadExpr (parser error)
	typ := meth.Recv.List[0].Type
	if ptr, ok := typ.(*ast.StarExpr); ok {
		typ = ptr.X
	}
	// determine receiver base type object
	var obj *ast.Object
	if ident, ok := typ.(*ast.Ident); ok && ident.Obj != nil {
		obj = ident.Obj
		if obj.Kind != ast.Typ {
			check.errorf(ident.Pos(), "%s is not a type", ident.Name)
			return // ignore this method
		}
		// TODO(gri) determine if obj was defined in this package
		/*
			if check.notLocal(obj) {
				check.errorf(ident.Pos(), "cannot define methods on non-local type %s", ident.Name)
				return // ignore this method
			}
		*/
	} else {
		// If it's not an identifier or the identifier wasn't declared/resolved,
		// the parser/resolver already reported an error. Nothing to do here.
		return // ignore this method
	}
	// declare method in receiver base type scope
	var scope *ast.Scope
	if obj.Data != nil {
		scope = obj.Data.(*ast.Scope)
	} else {
		scope = ast.NewScope(nil)
		obj.Data = scope
	}
	check.declare(scope, ast.Fun, meth.Name, meth)
}

func (check *checker) assocInitvalsOrMethod(decl ast.Decl) {
	switch d := decl.(type) {
	case *ast.GenDecl:
		if d.Tok == token.CONST {
			check.assocInitvals(d)
		}
	case *ast.FuncDecl:
		if d.Recv != nil {
			check.assocMethod(d)
		}
	}
}

func (check *checker) decl(decl ast.Decl) {
	switch d := decl.(type) {
	case *ast.BadDecl:
		// ignore
	case *ast.GenDecl:
		for _, spec := range d.Specs {
			switch s := spec.(type) {
			case *ast.ImportSpec:
				// nothing to do (handled by ast.NewPackage)
			case *ast.ValueSpec:
				for _, name := range s.Names {
					if obj := name.Obj; obj.Type == nil {
						check.object(obj, false)
					}
				}
			case *ast.TypeSpec:
				if obj := s.Name.Obj; obj.Type == nil {
					check.object(obj, false)
				}
			default:
				check.invalidAST(s.Pos(), "unknown ast.Spec node %T", s)
			}
		}
	case *ast.FuncDecl:
		// methods are checked when their respective base types are checked
		if d.Recv != nil {
			return
		}
		obj := d.Name.Obj
		// Initialization functions don't have an object associated with them
		// since they are not in any scope. Create a dummy object for them.
		if d.Name.Name == "init" {
			assert(obj == nil) // all other functions should have an object
			obj = ast.NewObj(ast.Fun, d.Name.Name)
			obj.Decl = d
			d.Name.Obj = obj
		}
		if obj.Type == nil {
			check.object(obj, false)
		}
	default:
		check.invalidAST(d.Pos(), "unknown ast.Decl node %T", d)
	}
}

// iterate calls f for each package-level declaration.
func (check *checker) iterate(f func(*checker, ast.Decl)) {
	for _, file := range check.files {
		for _, decl := range file.Decls {
			f(check, decl)
		}
	}
}

// sortedFiles returns the sorted list of package files given a package file map.
func sortedFiles(m map[string]*ast.File) []*ast.File {
	keys := make([]string, len(m))
	i := 0
	for k, _ := range m {
		keys[i] = k
		i++
	}
	sort.Strings(keys)

	files := make([]*ast.File, len(m))
	for i, k := range keys {
		files[i] = m[k]
	}

	return files
}

// A bailout panic is raised to indicate early termination.
type bailout struct{}

func check(ctxt *Context, fset *token.FileSet, files map[string]*ast.File) (pkg *ast.Package, err error) {
	// initialize checker
	check := checker{
		ctxt:      ctxt,
		fset:      fset,
		files:     sortedFiles(files),
		initexprs: make(map[*ast.ValueSpec][]ast.Expr),
	}

	// handle panics
	defer func() {
		switch p := recover().(type) {
		case nil, bailout:
			// normal return or early exit
			err = check.firsterr
		default:
			// unexpected panic: don't crash clients
			// panic(p) // enable for debugging
			// TODO(gri) add a test case for this scenario
			err = fmt.Errorf("types internal error: %v", p)
		}
	}()

	// resolve identifiers
	imp := ctxt.Import
	if imp == nil {
		imp = GcImport
	}
	pkg, err = ast.NewPackage(fset, files, imp, Universe)
	if err != nil {
		if list, _ := err.(scanner.ErrorList); len(list) > 0 {
			for _, err := range list {
				check.err(err)
			}
		} else {
			check.err(err)
		}
	}
	check.pkgscope = pkg.Scope

	// determine missing constant initialization expressions
	// and associate methods with types
	check.iterate((*checker).assocInitvalsOrMethod)

	// typecheck all declarations
	check.iterate((*checker).decl)

	// typecheck all function/method bodies
	// (funclist may grow when checking statements - do not use range clause!)
	for i := 0; i < len(check.funclist); i++ {
		f := check.funclist[i]
		if trace {
			s := "<function literal>"
			if f.obj != nil {
				s = f.obj.Name
			}
			fmt.Println("---", s)
		}
		check.funcsig = f.sig
		check.stmtList(f.body.List)
	}

	return
}
<commit_msg>go/types: "inherit" type in constant declarations w/o initialization expressions (bug fix)<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file implements the Check function, which typechecks a package.

package types

import (
	"fmt"
	"go/ast"
	"go/scanner"
	"go/token"
	"sort"
)

// enable for debugging
const trace = false

type checker struct {
	ctxt  *Context
	fset  *token.FileSet
	files []*ast.File

	// lazily initialized
	pkgscope *ast.Scope
	firsterr error
	initspec map[*ast.ValueSpec]*ast.ValueSpec // "inherited" type and initialization expressions for constant declarations
	funclist []function                        // list of functions/methods with correct signatures and non-empty bodies
	funcsig  *Signature                        // signature of currently typechecked function
	pos      []token.Pos                       // stack of expr positions; debugging support, used if trace is set
}

type function struct {
	obj  *ast.Object // for debugging/tracing only
	sig  *Signature
	body *ast.BlockStmt
}

// later adds a function with non-empty body to the list of functions
// that need to be processed after all package-level declarations
// are typechecked.
//
func (check *checker) later(obj *ast.Object, sig *Signature, body *ast.BlockStmt) {
	// functions implemented elsewhere (say in assembly) have no body
	if body != nil {
		check.funclist = append(check.funclist, function{obj, sig, body})
	}
}

// declare declares an object of the given kind and name (ident) in scope;
// decl is the corresponding declaration in the AST. An error is reported
// if the object was declared before.
//
// TODO(gri) This is very similar to the declare function in go/parser; it
// is only used to associate methods with their respective receiver base types.
// In a future version, it might be simpler and cleaner to do all the resolution
// in the type-checking phase. It would simplify the parser, AST, and also
// reduce some amount of code duplication.
//
func (check *checker) declare(scope *ast.Scope, kind ast.ObjKind, ident *ast.Ident, decl ast.Decl) {
	assert(ident.Obj == nil) // identifier already declared or resolved
	obj := ast.NewObj(kind, ident.Name)
	obj.Decl = decl
	ident.Obj = obj
	if ident.Name != "_" {
		if alt := scope.Insert(obj); alt != nil {
			prevDecl := ""
			if pos := alt.Pos(); pos.IsValid() {
				prevDecl = fmt.Sprintf("\n\tprevious declaration at %s", check.fset.Position(pos))
			}
			check.errorf(ident.Pos(), fmt.Sprintf("%s redeclared in this block%s", ident.Name, prevDecl))
		}
	}
}

func (check *checker) valueSpec(pos token.Pos, obj *ast.Object, lhs []*ast.Ident, typ ast.Expr, rhs []ast.Expr, iota int) {
	if len(lhs) == 0 {
		check.invalidAST(pos, "missing lhs in declaration")
		return
	}

	// determine type for all of lhs, if any
	// (but only set it for the object we typecheck!)
	var t Type
	if typ != nil {
		t = check.typ(typ, false)
	}

	// len(lhs) > 0
	if len(lhs) == len(rhs) {
		// check only lhs and rhs corresponding to obj
		var l, r ast.Expr
		for i, name := range lhs {
			if name.Obj == obj {
				l = lhs[i]
				r = rhs[i]
				break
			}
		}
		assert(l != nil)
		obj.Type = t
		check.assign1to1(l, r, nil, true, iota)
		return
	}

	// there must be a type or initialization expressions
	if t == nil && len(rhs) == 0 {
		check.invalidAST(pos, "missing type or initialization expression")
		t = Typ[Invalid]
	}

	// if we have a type, mark all of lhs
	if t != nil {
		for _, name := range lhs {
			name.Obj.Type = t
		}
	}

	// check initial values, if any
	if len(rhs) > 0 {
		// TODO(gri) should try to avoid this conversion
		lhx := make([]ast.Expr, len(lhs))
		for i, e := range lhs {
			lhx[i] = e
		}
		check.assignNtoM(lhx, rhs, true, iota)
	}
}

// object typechecks an object by assigning it a type; obj.Type must be nil.
// Callers must check obj.Type before calling object; this eliminates a call
// for each identifier that has been typechecked already, a common scenario.
//
func (check *checker) object(obj *ast.Object, cycleOk bool) {
	assert(obj.Type == nil)

	switch obj.Kind {
	case ast.Bad, ast.Pkg:
		// nothing to do

	case ast.Con, ast.Var:
		// The obj.Data field for constants and variables is initialized
		// to the respective (hypothetical, for variables) iota value by
		// the parser. The object's fields can be in one of the following
		// states:
		// Type != nil  =>  the constant value is Data
		// Type == nil  =>  the object is not typechecked yet, and Data can be:
		// Data is int  =>  Data is the value of iota for this declaration
		// Data == nil  =>  the object's expression is being evaluated
		if obj.Data == nil {
			check.errorf(obj.Pos(), "illegal cycle in initialization of %s", obj.Name)
			obj.Type = Typ[Invalid]
			return
		}
		spec := obj.Decl.(*ast.ValueSpec)
		iota := obj.Data.(int)
		obj.Data = nil
		// determine spec for type and initialization expressions
		init := spec
		if len(init.Values) == 0 && obj.Kind == ast.Con {
			init = check.initspec[spec]
		}
		check.valueSpec(spec.Pos(), obj, spec.Names, init.Type, init.Values, iota)

	case ast.Typ:
		typ := &NamedType{Obj: obj}
		obj.Type = typ // "mark" object so recursion terminates
		typ.Underlying = underlying(check.typ(obj.Decl.(*ast.TypeSpec).Type, cycleOk))
		// typecheck associated method signatures
		if obj.Data != nil {
			scope := obj.Data.(*ast.Scope)
			switch t := typ.Underlying.(type) {
			case *Struct:
				// struct fields must not conflict with methods
				for _, f := range t.Fields {
					if m := scope.Lookup(f.Name); m != nil {
						check.errorf(m.Pos(), "type %s has both field and method named %s", obj.Name, f.Name)
						// ok to continue
					}
				}
			case *Interface:
				// methods cannot be associated with an interface type
				for _, m := range scope.Objects {
					recv := m.Decl.(*ast.FuncDecl).Recv.List[0].Type
					check.errorf(recv.Pos(), "invalid receiver type %s (%s is an interface type)", obj.Name, obj.Name)
					// ok to continue
				}
			}
			// typecheck method signatures
			for _, obj := range scope.Objects {
				mdecl := obj.Decl.(*ast.FuncDecl)
				sig := check.typ(mdecl.Type, cycleOk).(*Signature)
				params, _ := check.collectParams(mdecl.Recv, false)
				sig.Recv = params[0] // the parser/assocMethod ensure there is exactly one parameter
				obj.Type = sig
				check.later(obj, sig, mdecl.Body)
			}
		}

	case ast.Fun:
		fdecl := obj.Decl.(*ast.FuncDecl)
		// methods are typechecked when their receivers are typechecked
		if fdecl.Recv == nil {
			sig := check.typ(fdecl.Type, cycleOk).(*Signature)
			if obj.Name == "init" && (len(sig.Params) != 0 || len(sig.Results) != 0) {
				check.errorf(fdecl.Pos(), "func init must have no arguments and no return values")
				// ok to continue
			}
			obj.Type = sig
			check.later(obj, sig, fdecl.Body)
		}

	default:
		panic("unreachable")
	}
}

// assocInitvals associates "inherited" initialization expressions
// with the corresponding *ast.ValueSpec in the check.initspec map
// for constant declarations without explicit initialization expressions.
//
func (check *checker) assocInitvals(decl *ast.GenDecl) {
	var last *ast.ValueSpec
	for _, s := range decl.Specs {
		if s, ok := s.(*ast.ValueSpec); ok {
			if len(s.Values) > 0 {
				last = s
			} else {
				check.initspec[s] = last
			}
		}
	}
	if last == nil {
		check.invalidAST(decl.Pos(), "no initialization values provided")
	}
}

// assocMethod associates a method declaration with the respective
// receiver base type. meth.Recv must exist.
//
func (check *checker) assocMethod(meth *ast.FuncDecl) {
	// The receiver type is one of the following (enforced by parser):
	// - *ast.Ident
	// - *ast.StarExpr{*ast.Ident}
	// - *ast.BadExpr (parser error)
	typ := meth.Recv.List[0].Type
	if ptr, ok := typ.(*ast.StarExpr); ok {
		typ = ptr.X
	}
	// determine receiver base type object
	var obj *ast.Object
	if ident, ok := typ.(*ast.Ident); ok && ident.Obj != nil {
		obj = ident.Obj
		if obj.Kind != ast.Typ {
			check.errorf(ident.Pos(), "%s is not a type", ident.Name)
			return // ignore this method
		}
		// TODO(gri) determine if obj was defined in this package
		/*
			if check.notLocal(obj) {
				check.errorf(ident.Pos(), "cannot define methods on non-local type %s", ident.Name)
				return // ignore this method
			}
		*/
	} else {
		// If it's not an identifier or the identifier wasn't declared/resolved,
		// the parser/resolver already reported an error. Nothing to do here.
		return // ignore this method
	}
	// declare method in receiver base type scope
	var scope *ast.Scope
	if obj.Data != nil {
		scope = obj.Data.(*ast.Scope)
	} else {
		scope = ast.NewScope(nil)
		obj.Data = scope
	}
	check.declare(scope, ast.Fun, meth.Name, meth)
}

func (check *checker) assocInitvalsOrMethod(decl ast.Decl) {
	switch d := decl.(type) {
	case *ast.GenDecl:
		if d.Tok == token.CONST {
			check.assocInitvals(d)
		}
	case *ast.FuncDecl:
		if d.Recv != nil {
			check.assocMethod(d)
		}
	}
}

func (check *checker) decl(decl ast.Decl) {
	switch d := decl.(type) {
	case *ast.BadDecl:
		// ignore
	case *ast.GenDecl:
		for _, spec := range d.Specs {
			switch s := spec.(type) {
			case *ast.ImportSpec:
				// nothing to do (handled by ast.NewPackage)
			case *ast.ValueSpec:
				for _, name := range s.Names {
					if obj := name.Obj; obj.Type == nil {
						check.object(obj, false)
					}
				}
			case *ast.TypeSpec:
				if obj := s.Name.Obj; obj.Type == nil {
					check.object(obj, false)
				}
			default:
				check.invalidAST(s.Pos(), "unknown ast.Spec node %T", s)
			}
		}
	case *ast.FuncDecl:
		// methods are checked when their respective base types are checked
		if d.Recv != nil {
			return
		}
		obj := d.Name.Obj
		// Initialization functions don't have an object associated with them
		// since they are not in any scope. Create a dummy object for them.
		if d.Name.Name == "init" {
			assert(obj == nil) // all other functions should have an object
			obj = ast.NewObj(ast.Fun, d.Name.Name)
			obj.Decl = d
			d.Name.Obj = obj
		}
		if obj.Type == nil {
			check.object(obj, false)
		}
	default:
		check.invalidAST(d.Pos(), "unknown ast.Decl node %T", d)
	}
}

// iterate calls f for each package-level declaration.
func (check *checker) iterate(f func(*checker, ast.Decl)) {
	for _, file := range check.files {
		for _, decl := range file.Decls {
			f(check, decl)
		}
	}
}

// sortedFiles returns the sorted list of package files given a package file map.
func sortedFiles(m map[string]*ast.File) []*ast.File {
	keys := make([]string, len(m))
	i := 0
	for k, _ := range m {
		keys[i] = k
		i++
	}
	sort.Strings(keys)

	files := make([]*ast.File, len(m))
	for i, k := range keys {
		files[i] = m[k]
	}

	return files
}

// A bailout panic is raised to indicate early termination.
type bailout struct{}

func check(ctxt *Context, fset *token.FileSet, files map[string]*ast.File) (pkg *ast.Package, err error) {
	// initialize checker
	check := checker{
		ctxt:     ctxt,
		fset:     fset,
		files:    sortedFiles(files),
		initspec: make(map[*ast.ValueSpec]*ast.ValueSpec),
	}

	// handle panics
	defer func() {
		switch p := recover().(type) {
		case nil, bailout:
			// normal return or early exit
			err = check.firsterr
		default:
			// unexpected panic: don't crash clients
			// panic(p) // enable for debugging
			// TODO(gri) add a test case for this scenario
			err = fmt.Errorf("types internal error: %v", p)
		}
	}()

	// resolve identifiers
	imp := ctxt.Import
	if imp == nil {
		imp = GcImport
	}
	pkg, err = ast.NewPackage(fset, files, imp, Universe)
	if err != nil {
		if list, _ := err.(scanner.ErrorList); len(list) > 0 {
			for _, err := range list {
				check.err(err)
			}
		} else {
			check.err(err)
		}
	}
	check.pkgscope = pkg.Scope

	// determine missing constant initialization expressions
	// and associate methods with types
	check.iterate((*checker).assocInitvalsOrMethod)

	// typecheck all declarations
	check.iterate((*checker).decl)

	// typecheck all function/method bodies
	// (funclist may grow when checking statements - do not use range clause!)
	for i := 0; i < len(check.funclist); i++ {
		f := check.funclist[i]
		if trace {
			s := "<function literal>"
			if f.obj != nil {
				s = f.obj.Name
			}
			fmt.Println("---", s)
		}
		check.funcsig = f.sig
		check.stmtList(f.body.List)
	}

	return
}
<|endoftext|>


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

// TODO(cw): ListenPacket test, Read() test, ipv6 test &
// Dial()/Listen() level tests

package net

import (
	"bytes"
	"flag"
	"os"
	"testing"
)

const ICMP_ECHO_REQUEST = 8
const ICMP_ECHO_REPLY = 0

// returns a suitable 'ping request' packet, with id & seq and a
// payload length of pktlen
func makePingRequest(id, seq, pktlen int, filler []byte) []byte {
	p := make([]byte, pktlen)
	copy(p[8:], bytes.Repeat(filler, (pktlen-8)/len(filler)+1))

	p[0] = ICMP_ECHO_REQUEST // type
	p[1] = 0                 // code
	p[2] = 0                 // cksum
	p[3] = 0                 // cksum
	p[4] = uint8(id >> 8)    // id
	p[5] = uint8(id & 0xff)  // id
	p[6] = uint8(seq >> 8)   // sequence
	p[7] = uint8(seq & 0xff) // sequence

	// calculate icmp checksum
	cklen := len(p)
	s := uint32(0)
	for i := 0; i < (cklen - 1); i += 2 {
		s += uint32(p[i+1])<<8 | uint32(p[i])
	}
	if cklen&1 == 1 {
		s += uint32(p[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)

	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	p[2] ^= uint8(^s & 0xff)
	p[3] ^= uint8(^s >> 8)

	return p
}

func parsePingReply(p []byte) (id, seq int) {
	id = int(p[4])<<8 | int(p[5])
	seq = int(p[6])<<8 | int(p[7])
	return
}

var srchost = flag.String("srchost", "", "Source of the ICMP ECHO request")
var dsthost = flag.String("dsthost", "localhost", "Destination for the ICMP ECHO request")

// test (raw) IP socket using ICMP
func TestICMP(t *testing.T) {
	if os.Getuid() != 0 {
		t.Logf("test disabled; must be root")
		return
	}

	var (
		laddr *IPAddr
		err   os.Error
	)
	if *srchost != "" {
		laddr, err = ResolveIPAddr(*srchost)
		if err != nil {
			t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *srchost, laddr, err)
		}
	}

	raddr, err := ResolveIPAddr(*dsthost)
	if err != nil {
		t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *dsthost, raddr, err)
	}

	c, err := ListenIP("ip4:icmp", laddr)
	if err != nil {
		t.Fatalf(`net.ListenIP("ip4:icmp", %v) = %v, %v`, *srchost, c, err)
	}

	sendid := os.Getpid() & 0xffff
	const sendseq = 61455
	const pingpktlen = 128
	sendpkt := makePingRequest(sendid, sendseq, pingpktlen, []byte("Go Go Gadget Ping!!!"))

	n, err := c.WriteToIP(sendpkt, raddr)
	if err != nil || n != pingpktlen {
		t.Fatalf(`net.WriteToIP(..., %v) = %v, %v`, raddr, n, err)
	}

	c.SetTimeout(100e6)
	resp := make([]byte, 1024)
	for {
		n, from, err := c.ReadFrom(resp)
		if err != nil {
			t.Fatalf(`ReadFrom(...) = %v, %v, %v`, n, from, err)
		}
		if resp[0] != ICMP_ECHO_REPLY {
			continue
		}
		rcvid, rcvseq := parsePingReply(resp)
		if rcvid != sendid || rcvseq != sendseq {
			t.Fatalf(`Ping reply saw id,seq=0x%x,0x%x (expected 0x%x, 0x%x)`, rcvid, rcvseq, sendid, sendseq)
		}
		return
	}
	t.Fatalf("saw no ping return")
}
<commit_msg>net: default to 127.0.0.1, not localhost, in TestICMP<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// TODO(cw): ListenPacket test, Read() test, ipv6 test &
// Dial()/Listen() level tests

package net

import (
	"bytes"
	"flag"
	"os"
	"testing"
)

const ICMP_ECHO_REQUEST = 8
const ICMP_ECHO_REPLY = 0

// returns a suitable 'ping request' packet, with id & seq and a
// payload length of pktlen
func makePingRequest(id, seq, pktlen int, filler []byte) []byte {
	p := make([]byte, pktlen)
	copy(p[8:], bytes.Repeat(filler, (pktlen-8)/len(filler)+1))

	p[0] = ICMP_ECHO_REQUEST // type
	p[1] = 0                 // code
	p[2] = 0                 // cksum
	p[3] = 0                 // cksum
	p[4] = uint8(id >> 8)    // id
	p[5] = uint8(id & 0xff)  // id
	p[6] = uint8(seq >> 8)   // sequence
	p[7] = uint8(seq & 0xff) // sequence

	// calculate icmp checksum
	cklen := len(p)
	s := uint32(0)
	for i := 0; i < (cklen - 1); i += 2 {
		s += uint32(p[i+1])<<8 | uint32(p[i])
	}
	if cklen&1 == 1 {
		s += uint32(p[cklen-1])
	}
	s = (s >> 16) + (s & 0xffff)
	s = s + (s >> 16)

	// place checksum back in header; using ^= avoids the
	// assumption the checksum bytes are zero
	p[2] ^= uint8(^s & 0xff)
	p[3] ^= uint8(^s >> 8)

	return p
}

func parsePingReply(p []byte) (id, seq int) {
	id = int(p[4])<<8 | int(p[5])
	seq = int(p[6])<<8 | int(p[7])
	return
}

var srchost = flag.String("srchost", "", "Source of the ICMP ECHO request")
// 127.0.0.1 because this is an IPv4-specific test.
var dsthost = flag.String("dsthost", "127.0.0.1", "Destination for the ICMP ECHO request")

// test (raw) IP socket using ICMP
func TestICMP(t *testing.T) {
	if os.Getuid() != 0 {
		t.Logf("test disabled; must be root")
		return
	}

	var (
		laddr *IPAddr
		err   os.Error
	)
	if *srchost != "" {
		laddr, err = ResolveIPAddr(*srchost)
		if err != nil {
			t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *srchost, laddr, err)
		}
	}

	raddr, err := ResolveIPAddr(*dsthost)
	if err != nil {
		t.Fatalf(`net.ResolveIPAddr("%v") = %v, %v`, *dsthost, raddr, err)
	}

	c, err := ListenIP("ip4:icmp", laddr)
	if err != nil {
		t.Fatalf(`net.ListenIP("ip4:icmp", %v) = %v, %v`, *srchost, c, err)
	}

	sendid := os.Getpid() & 0xffff
	const sendseq = 61455
	const pingpktlen = 128
	sendpkt := makePingRequest(sendid, sendseq, pingpktlen, []byte("Go Go Gadget Ping!!!"))

	n, err := c.WriteToIP(sendpkt, raddr)
	if err != nil || n != pingpktlen {
		t.Fatalf(`net.WriteToIP(..., %v) = %v, %v`, raddr, n, err)
	}

	c.SetTimeout(100e6)
	resp := make([]byte, 1024)
	for {
		n, from, err := c.ReadFrom(resp)
		if err != nil {
			t.Fatalf(`ReadFrom(...) = %v, %v, %v`, n, from, err)
		}
		if resp[0] != ICMP_ECHO_REPLY {
			continue
		}
		rcvid, rcvseq := parsePingReply(resp)
		if rcvid != sendid || rcvseq != sendseq {
			t.Fatalf(`Ping reply saw id,seq=0x%x,0x%x (expected 0x%x, 0x%x)`, rcvid, rcvseq, sendid, sendseq)
		}
		return
	}
	t.Fatalf("saw no ping return")
}
<|endoftext|>