<commit_before>package zmq

/*
#cgo pkg-config: libzmq
#include <zmq.h>
#include <stdlib.h>
*/
import "C"

import (
	"reflect"
	"unsafe"
)

// Socket represents a zero mq socket
type Socket struct {
	psocket unsafe.Pointer
}

// SocketType identifies the type of the socket
type SocketType C.int

// Bindings to available socket types
const (
	Req    = SocketType(C.ZMQ_REQ)
	Rep    = SocketType(C.ZMQ_REP)
	Router = SocketType(C.ZMQ_ROUTER)
	Dealer = SocketType(C.ZMQ_DEALER)
	Pull   = SocketType(C.ZMQ_PULL)
	Push   = SocketType(C.ZMQ_PUSH)
	Pub    = SocketType(C.ZMQ_PUB)
	Sub    = SocketType(C.ZMQ_SUB)
	Xsub   = SocketType(C.ZMQ_XSUB)
	Xpub   = SocketType(C.ZMQ_XPUB)
	Pair   = SocketType(C.ZMQ_PAIR)
)

// SendFlag identifies the flags passed to zeromq send command
type SendFlag C.int

// Bindings to available send flags
const (
	SndMore  = SendFlag(C.ZMQ_SNDMORE)
	DontWait = SendFlag(C.ZMQ_DONTWAIT)
)

// Close 0mq socket.
func (s *Socket) Close() error {
	rc, err := C.zmq_close(s.psocket)
	if rc == 0 {
		return nil
	}
	return err
}

// Bind the socket to the given address
func (s *Socket) Bind(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_bind(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Unbind the socket from the given address
func (s *Socket) Unbind(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_unbind(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Connect the socket to the given address
func (s *Socket) Connect(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_connect(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Disconnect the socket from the given address
func (s *Socket) Disconnect(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_disconnect(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Send data to the socket
func (s *Socket) Send(data []byte, flag SendFlag) error {
	var pdata unsafe.Pointer
	var msg C.zmq_msg_t
	if len(data) == 0 {
		pdata = unsafe.Pointer(&data)
	} else {
		pdata = unsafe.Pointer(&data[0])
	}

	sizeData := C.size_t(len(data))
	// The slice is reused as an unsafe pointer to avoid copy
	// There might be a problem if go gc collect data slice
	// before the message is effectively send.
	// TODO try to use a leaky bucket to mitigate gc collect
	// and improve performances
	C.zmq_msg_init_data(&msg, pdata, sizeData, nil, nil)
	for {
		rc, err := C.zmq_msg_send(&msg, s.psocket, C.int(flag))
		// Retry send on an interrupted system call
		if rc == -1 && C.zmq_errno() == C.int(C.EINTR) {
			continue
		}
		if rc == -1 {
			return err
		}
		break
	}
	C.zmq_msg_close(&msg)
	return nil
}

// SendMultipart sends a message with on or several frames to the socket
func (s *Socket) SendMultipart(data [][]byte, flag SendFlag) error {
	moreFlag := flag | SndMore
	for _, v := range data[:len(data)-1] {
		err := s.Send(v, moreFlag)
		if err != nil {
			return err
		}
	}
	err := s.Send(data[len(data)-1], flag)
	if err != nil {
		return err
	}
	return nil
}

// RecvMultipart receives a multi part message from the socket
func (s *Socket) RecvMultipart(flag SendFlag) (*MessageMultipart, error) {
	msg := &MessageMultipart{}
	msg.parts = make([]*MessagePart, 10)
	i := 0
	for {
		msgPart, err := s.Recv(flag)
		if err != nil {
			return nil, err
		}
		msg.parts[i] = msgPart
		i += 1
		if !msgPart.HasMore() {
			break
		}
	}
	// Make slice iterable
	msg.parts = msg.parts[:i]
	msg.aggregateData()
	return msg, nil
}

var messagePartPool = make(chan *MessagePart, 100)

// Recv receives a message part from the socket
// It is necessary to call CloseMsg on each MessagePart to avoid memory leak
// when the data is not needed anymore
func (s *Socket) Recv(flag SendFlag) (*MessagePart, error) {
	var msg C.zmq_msg_t
	rc, err := C.zmq_msg_init(&msg)
	if rc != 0 {
		return nil, err
	}
	for {
		rc, err = C.zmq_msg_recv(&msg, s.psocket, 0)
		// Retry receive on an interrupted system call
		if rc == -1 && C.zmq_errno() == C.int(C.EINTR) {
			continue
		}
		if rc == -1 {
			C.zmq_msg_close(&msg)
			return nil, err
		}
		break
	}
	data := buildSliceFromMsg(&msg)

	var msgPart *MessagePart
	select {
	case msgPart = <-messagePartPool:
	default:
		msgPart = &MessagePart{}
	}

	msgPart.Data = data
	msgPart.zmqMsg = (*zmqMsg)(&msg)

	return msgPart, nil
}

// SocketOptionInt identifies socket option which returns int value
type SocketOptionInt C.int
// SocketOptionUint64 identifies socket option which returns uint64 value
type SocketOptionUint64 C.int
// SocketOptionInt64 identifies socket option which returns int64 value
type SocketOptionInt64 C.int
// SocketOptionString identifies socket option which returns string value
type SocketOptionString C.int

// Bindings to socket options
const (
	Type                 = SocketOptionInt(C.ZMQ_TYPE)
	Rcvmore              = SocketOptionInt(C.ZMQ_RCVMORE)
	Sndhwm               = SocketOptionInt(C.ZMQ_SNDHWM)
	Rcvhwm               = SocketOptionInt(C.ZMQ_RCVHWM)
	Affinity             = SocketOptionUint64(C.ZMQ_AFFINITY)
	Identity             = SocketOptionString(C.ZMQ_IDENTITY)
	Rate                 = SocketOptionInt(C.ZMQ_RATE)
	RecoveryIvl          = SocketOptionInt(C.ZMQ_RECOVERY_IVL)
	Sndbuf               = SocketOptionInt(C.ZMQ_SNDBUF)
	Rcvbuf               = SocketOptionInt(C.ZMQ_RCVBUF)
	Linger               = SocketOptionInt(C.ZMQ_LINGER)
	ReconnectIvl         = SocketOptionInt(C.ZMQ_RECONNECT_IVL)
	ReconnectIvlMax      = SocketOptionInt(C.ZMQ_RECONNECT_IVL_MAX)
	Backlog              = SocketOptionInt(C.ZMQ_BACKLOG)
	Maxmsgsize           = SocketOptionInt64(C.ZMQ_MAXMSGSIZE)
	MulticastHops        = SocketOptionInt(C.ZMQ_MULTICAST_HOPS)
	Rcvtimeo             = SocketOptionInt(C.ZMQ_RCVTIMEO)
	Sndtimeo             = SocketOptionInt(C.ZMQ_SNDTIMEO)
	Ipv4only             = SocketOptionInt(C.ZMQ_IPV4ONLY)
	DelayAttachOnConnect = SocketOptionInt(C.ZMQ_DELAY_ATTACH_ON_CONNECT)
	Fd                   = SocketOptionInt(C.ZMQ_FD)
	Events               = SocketOptionInt(C.ZMQ_EVENTS)
	LastEndpoint         = SocketOptionString(C.ZMQ_LAST_ENDPOINT)
	TcpKeepalive         = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE)
	TcpKeepaliveIdle     = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_IDLE)
	TcpKeepaliveCnt      = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_CNT)
	TcpKeepaliveIntvl    = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_INTVL)

	Subscribe       = SocketOptionString(C.ZMQ_SUBSCRIBE)
	Unsubscribe     = SocketOptionString(C.ZMQ_UNSUBSCRIBE)
	RouterMandatory = SocketOptionInt(C.ZMQ_ROUTER_MANDATORY)
	XpubVerbose     = SocketOptionInt(C.ZMQ_XPUB_VERBOSE)
)

func (s *Socket) getOption(option C.int, v interface{}, size *C.size_t) error {
	value := reflect.ValueOf(v)
	pvalue := unsafe.Pointer(value.Pointer())
	rc, err := C.zmq_getsockopt(s.psocket, option, pvalue, size)
	if rc == -1 {
		return err
	}
	return nil
}

// GetOptionInt gets the value of a socket option as an int
func (s *Socket) GetOptionInt(option SocketOptionInt) (int, error) {
	var value int
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionUint64 gets the value of a socket option as an uint64
func (s *Socket) GetOptionUint64(option SocketOptionUint64) (uint64, error) {
	var value uint64
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionInt64 gets the value of a socket option as an int64
func (s *Socket) GetOptionInt64(option SocketOptionUint64) (int64, error) {
	var value int64
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionString gets the value of a socket option as a string
func (s *Socket) GetOptionString(option SocketOptionString) (string, error) {
	var value [1024]byte
	sizeString := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &sizeString)
	if sizeString > 0 {
		// Remove \x00 from zmq string
		return string(value[:sizeString-1]), err
	}
	return "", nil
}

func (s *Socket) setOption(option C.int, v interface{}, size C.size_t) error {
	value := reflect.ValueOf(v)
	pvalue := unsafe.Pointer(value.Pointer())
	rc, err := C.zmq_setsockopt(s.psocket, option, pvalue, size)
	if rc == -1 {
		return err
	}
	return nil
}

// SetOptionInt sets a int socket option to the given value
func (s *Socket) SetOptionInt(option SocketOptionInt, value int) error {
	val := C.int(value)
	size := C.size_t(unsafe.Sizeof(val))
	return s.setOption(C.int(option), &val, size)
}

// SetOptionInt64 sets a int 64 socket option to the given value
func (s *Socket) SetOptionInt64(option SocketOptionInt64, value int64) error {
	size := C.size_t(unsafe.Sizeof(value))
	return s.setOption(C.int(option), &value, size)
}

// SetOptionUint64 sets a uint 64 socket option to the given value
func (s *Socket) SetOptionUint64(option SocketOptionUint64, value uint64) error {
	size := C.size_t(unsafe.Sizeof(value))
	return s.setOption(C.int(option), &value, size)
}

// SetOptionString sets a string socket option to the given value. Can be nil
func (s *Socket) SetOptionString(option SocketOptionString, value *string) error {
	if value == nil {
		return s.setOption(C.int(option), nil, 0)
	}
	size := C.size_t(len(*value))
	cstr := C.CString(*value)
	err := s.setOption(C.int(option), cstr, size)
	C.free(unsafe.Pointer(cstr))
	return err
}

// SocketEvent identifies socket events available
type SocketEvent C.int

// Bindings to socket events
const (
	EventConnected        = SocketEvent(C.ZMQ_EVENT_CONNECTED)
	EventConnectDelayed  = SocketEvent(C.ZMQ_EVENT_CONNECT_DELAYED)
	EventConnectRetried = SocketEvent(C.ZMQ_EVENT_CONNECT_RETRIED)
	EventListening        = SocketEvent(C.ZMQ_EVENT_LISTENING)
	EventBindFailed      = SocketEvent(C.ZMQ_EVENT_BIND_FAILED)
	EventAccepted         = SocketEvent(C.ZMQ_EVENT_ACCEPTED)
	EventAcceptFailed    = SocketEvent(C.ZMQ_EVENT_ACCEPT_FAILED)
	EventClosed           = SocketEvent(C.ZMQ_EVENT_CLOSED)
	EventCloseFailed     = SocketEvent(C.ZMQ_EVENT_CLOSE_FAILED)
	EventDisconnected     = SocketEvent(C.ZMQ_EVENT_DISCONNECTED)
	EventAll              = SocketEvent(C.ZMQ_EVENT_ALL)
)

// Monitor binds event to the socket
func (s *Socket) Monitor(endpoint string, events SocketEvent) error {
	cstr := C.CString(endpoint)
	rc, err := C.zmq_socket_monitor(s.psocket, cstr, C.int(events))
	if rc == -1 {
		return err
	}
	return nil
}
<commit_msg>Fix msg part appending<commit_after>package zmq

/*
#cgo pkg-config: libzmq
#include <zmq.h>
#include <stdlib.h>
*/
import "C"

import (
	"reflect"
	"unsafe"
)

// Socket represents a zero mq socket
type Socket struct {
	psocket unsafe.Pointer
}

// SocketType identifies the type of the socket
type SocketType C.int

// Bindings to available socket types
const (
	Req    = SocketType(C.ZMQ_REQ)
	Rep    = SocketType(C.ZMQ_REP)
	Router = SocketType(C.ZMQ_ROUTER)
	Dealer = SocketType(C.ZMQ_DEALER)
	Pull   = SocketType(C.ZMQ_PULL)
	Push   = SocketType(C.ZMQ_PUSH)
	Pub    = SocketType(C.ZMQ_PUB)
	Sub    = SocketType(C.ZMQ_SUB)
	Xsub   = SocketType(C.ZMQ_XSUB)
	Xpub   = SocketType(C.ZMQ_XPUB)
	Pair   = SocketType(C.ZMQ_PAIR)
)

// SendFlag identifies the flags passed to zeromq send command
type SendFlag C.int

// Bindings to available send flags
const (
	SndMore  = SendFlag(C.ZMQ_SNDMORE)
	DontWait = SendFlag(C.ZMQ_DONTWAIT)
)

// Close 0mq socket.
func (s *Socket) Close() error {
	rc, err := C.zmq_close(s.psocket)
	if rc == 0 {
		return nil
	}
	return err
}

// Bind the socket to the given address
func (s *Socket) Bind(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_bind(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Unbind the socket from the given address
func (s *Socket) Unbind(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_unbind(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Connect the socket to the given address
func (s *Socket) Connect(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_connect(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Disconnect the socket from the given address
func (s *Socket) Disconnect(address string) error {
	addr := C.CString(address)
	defer C.free(unsafe.Pointer(addr))
	rc, err := C.zmq_disconnect(s.psocket, addr)
	if rc == 0 {
		return nil
	}
	return err
}

// Send data to the socket
func (s *Socket) Send(data []byte, flag SendFlag) error {
	var pdata unsafe.Pointer
	var msg C.zmq_msg_t
	if len(data) == 0 {
		pdata = unsafe.Pointer(&data)
	} else {
		pdata = unsafe.Pointer(&data[0])
	}

	sizeData := C.size_t(len(data))
	// The slice is reused as an unsafe pointer to avoid copy
	// There might be a problem if go gc collect data slice
	// before the message is effectively send.
	// TODO try to use a leaky bucket to mitigate gc collect
	// and improve performances
	C.zmq_msg_init_data(&msg, pdata, sizeData, nil, nil)
	for {
		rc, err := C.zmq_msg_send(&msg, s.psocket, C.int(flag))
		// Retry send on an interrupted system call
		if rc == -1 && C.zmq_errno() == C.int(C.EINTR) {
			continue
		}
		if rc == -1 {
			return err
		}
		break
	}
	C.zmq_msg_close(&msg)
	return nil
}

// SendMultipart sends a message with on or several frames to the socket
func (s *Socket) SendMultipart(data [][]byte, flag SendFlag) error {
	moreFlag := flag | SndMore
	for _, v := range data[:len(data)-1] {
		err := s.Send(v, moreFlag)
		if err != nil {
			return err
		}
	}
	err := s.Send(data[len(data)-1], flag)
	if err != nil {
		return err
	}
	return nil
}

// RecvMultipart receives a multi part message from the socket
func (s *Socket) RecvMultipart(flag SendFlag) (*MessageMultipart, error) {
	msg := &MessageMultipart{}
	msg.parts = make([]*MessagePart, 0, 10)
	i := 0
	for {
		msgPart, err := s.Recv(flag)
		if err != nil {
			return nil, err
		}
		msg.parts = append(msg.parts, msgPart)
		i += 1
		if !msgPart.HasMore() {
			break
		}
	}
	// Make slice iterable
	msg.parts = msg.parts[:i]
	msg.aggregateData()
	return msg, nil
}

var messagePartPool = make(chan *MessagePart, 100)

// Recv receives a message part from the socket
// It is necessary to call CloseMsg on each MessagePart to avoid memory leak
// when the data is not needed anymore
func (s *Socket) Recv(flag SendFlag) (*MessagePart, error) {
	var msg C.zmq_msg_t
	rc, err := C.zmq_msg_init(&msg)
	if rc != 0 {
		return nil, err
	}
	for {
		rc, err = C.zmq_msg_recv(&msg, s.psocket, 0)
		// Retry receive on an interrupted system call
		if rc == -1 && C.zmq_errno() == C.int(C.EINTR) {
			continue
		}
		if rc == -1 {
			C.zmq_msg_close(&msg)
			return nil, err
		}
		break
	}
	data := buildSliceFromMsg(&msg)

	var msgPart *MessagePart
	select {
	case msgPart = <-messagePartPool:
	default:
		msgPart = &MessagePart{}
	}

	msgPart.Data = data
	msgPart.zmqMsg = (*zmqMsg)(&msg)

	return msgPart, nil
}

// SocketOptionInt identifies socket option which returns int value
type SocketOptionInt C.int
// SocketOptionUint64 identifies socket option which returns uint64 value
type SocketOptionUint64 C.int
// SocketOptionInt64 identifies socket option which returns int64 value
type SocketOptionInt64 C.int
// SocketOptionString identifies socket option which returns string value
type SocketOptionString C.int

// Bindings to socket options
const (
	Type                 = SocketOptionInt(C.ZMQ_TYPE)
	Rcvmore              = SocketOptionInt(C.ZMQ_RCVMORE)
	Sndhwm               = SocketOptionInt(C.ZMQ_SNDHWM)
	Rcvhwm               = SocketOptionInt(C.ZMQ_RCVHWM)
	Affinity             = SocketOptionUint64(C.ZMQ_AFFINITY)
	Identity             = SocketOptionString(C.ZMQ_IDENTITY)
	Rate                 = SocketOptionInt(C.ZMQ_RATE)
	RecoveryIvl          = SocketOptionInt(C.ZMQ_RECOVERY_IVL)
	Sndbuf               = SocketOptionInt(C.ZMQ_SNDBUF)
	Rcvbuf               = SocketOptionInt(C.ZMQ_RCVBUF)
	Linger               = SocketOptionInt(C.ZMQ_LINGER)
	ReconnectIvl         = SocketOptionInt(C.ZMQ_RECONNECT_IVL)
	ReconnectIvlMax      = SocketOptionInt(C.ZMQ_RECONNECT_IVL_MAX)
	Backlog              = SocketOptionInt(C.ZMQ_BACKLOG)
	Maxmsgsize           = SocketOptionInt64(C.ZMQ_MAXMSGSIZE)
	MulticastHops        = SocketOptionInt(C.ZMQ_MULTICAST_HOPS)
	Rcvtimeo             = SocketOptionInt(C.ZMQ_RCVTIMEO)
	Sndtimeo             = SocketOptionInt(C.ZMQ_SNDTIMEO)
	Ipv4only             = SocketOptionInt(C.ZMQ_IPV4ONLY)
	DelayAttachOnConnect = SocketOptionInt(C.ZMQ_DELAY_ATTACH_ON_CONNECT)
	Fd                   = SocketOptionInt(C.ZMQ_FD)
	Events               = SocketOptionInt(C.ZMQ_EVENTS)
	LastEndpoint         = SocketOptionString(C.ZMQ_LAST_ENDPOINT)
	TcpKeepalive         = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE)
	TcpKeepaliveIdle     = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_IDLE)
	TcpKeepaliveCnt      = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_CNT)
	TcpKeepaliveIntvl    = SocketOptionInt(C.ZMQ_TCP_KEEPALIVE_INTVL)

	Subscribe       = SocketOptionString(C.ZMQ_SUBSCRIBE)
	Unsubscribe     = SocketOptionString(C.ZMQ_UNSUBSCRIBE)
	RouterMandatory = SocketOptionInt(C.ZMQ_ROUTER_MANDATORY)
	XpubVerbose     = SocketOptionInt(C.ZMQ_XPUB_VERBOSE)
)

func (s *Socket) getOption(option C.int, v interface{}, size *C.size_t) error {
	value := reflect.ValueOf(v)
	pvalue := unsafe.Pointer(value.Pointer())
	rc, err := C.zmq_getsockopt(s.psocket, option, pvalue, size)
	if rc == -1 {
		return err
	}
	return nil
}

// GetOptionInt gets the value of a socket option as an int
func (s *Socket) GetOptionInt(option SocketOptionInt) (int, error) {
	var value int
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionUint64 gets the value of a socket option as an uint64
func (s *Socket) GetOptionUint64(option SocketOptionUint64) (uint64, error) {
	var value uint64
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionInt64 gets the value of a socket option as an int64
func (s *Socket) GetOptionInt64(option SocketOptionUint64) (int64, error) {
	var value int64
	size := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &size)
	return value, err
}

// GetOptionString gets the value of a socket option as a string
func (s *Socket) GetOptionString(option SocketOptionString) (string, error) {
	var value [1024]byte
	sizeString := C.size_t(unsafe.Sizeof(value))
	err := s.getOption(C.int(option), &value, &sizeString)
	if sizeString > 0 {
		// Remove \x00 from zmq string
		return string(value[:sizeString-1]), err
	}
	return "", nil
}

func (s *Socket) setOption(option C.int, v interface{}, size C.size_t) error {
	value := reflect.ValueOf(v)
	pvalue := unsafe.Pointer(value.Pointer())
	rc, err := C.zmq_setsockopt(s.psocket, option, pvalue, size)
	if rc == -1 {
		return err
	}
	return nil
}

// SetOptionInt sets a int socket option to the given value
func (s *Socket) SetOptionInt(option SocketOptionInt, value int) error {
	val := C.int(value)
	size := C.size_t(unsafe.Sizeof(val))
	return s.setOption(C.int(option), &val, size)
}

// SetOptionInt64 sets a int 64 socket option to the given value
func (s *Socket) SetOptionInt64(option SocketOptionInt64, value int64) error {
	size := C.size_t(unsafe.Sizeof(value))
	return s.setOption(C.int(option), &value, size)
}

// SetOptionUint64 sets a uint 64 socket option to the given value
func (s *Socket) SetOptionUint64(option SocketOptionUint64, value uint64) error {
	size := C.size_t(unsafe.Sizeof(value))
	return s.setOption(C.int(option), &value, size)
}

// SetOptionString sets a string socket option to the given value. Can be nil
func (s *Socket) SetOptionString(option SocketOptionString, value *string) error {
	if value == nil {
		return s.setOption(C.int(option), nil, 0)
	}
	size := C.size_t(len(*value))
	cstr := C.CString(*value)
	err := s.setOption(C.int(option), cstr, size)
	C.free(unsafe.Pointer(cstr))
	return err
}

// SocketEvent identifies socket events available
type SocketEvent C.int

// Bindings to socket events
const (
	EventConnected        = SocketEvent(C.ZMQ_EVENT_CONNECTED)
	EventConnectDelayed  = SocketEvent(C.ZMQ_EVENT_CONNECT_DELAYED)
	EventConnectRetried = SocketEvent(C.ZMQ_EVENT_CONNECT_RETRIED)
	EventListening        = SocketEvent(C.ZMQ_EVENT_LISTENING)
	EventBindFailed      = SocketEvent(C.ZMQ_EVENT_BIND_FAILED)
	EventAccepted         = SocketEvent(C.ZMQ_EVENT_ACCEPTED)
	EventAcceptFailed    = SocketEvent(C.ZMQ_EVENT_ACCEPT_FAILED)
	EventClosed           = SocketEvent(C.ZMQ_EVENT_CLOSED)
	EventCloseFailed     = SocketEvent(C.ZMQ_EVENT_CLOSE_FAILED)
	EventDisconnected     = SocketEvent(C.ZMQ_EVENT_DISCONNECTED)
	EventAll              = SocketEvent(C.ZMQ_EVENT_ALL)
)

// Monitor binds event to the socket
func (s *Socket) Monitor(endpoint string, events SocketEvent) error {
	cstr := C.CString(endpoint)
	rc, err := C.zmq_socket_monitor(s.psocket, cstr, C.int(events))
	if rc == -1 {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package test

import (
	"fmt"
	"strings"
	"time"

	"github.com/nfrance-conseil/zeplic/lib"
	"testing"
)

func TestDatasetName(t *testing.T) {
	name := lib.DatasetName("tank/test@SNAP")
	if name != "tank/test" {
		t.Errorf("DatasetName() test failed!")
	}
}

func TestSnapName(t *testing.T) {
	name := lib.SnapName("SNAP")
	year, month, day := time.Now().Date()
	get := fmt.Sprintf("%s_%d-%s-%02d", "SNAP", year, month, day)
	if strings.Contains(name, get) == false {
		t.Errorf("SnapName() test failed!")
	}
}

func TestSnapBackup(t *testing.T) {
	backup := lib.SnapBackup()
	year, month, day := time.Now().Date()
	get := fmt.Sprintf("%s_%d-%s-%02d", "BACKUP", year, month, day)
	if strings.Contains(backup, get) == false {
		t.Errorf("SnapBackup() test failed!")
	}
}

func TestRenamed(t *testing.T) {
	renamed := lib.Renamed("tank/test@SNAP1", "tank/test@SNAP2")
	if renamed == false {
		t.Errorf("Renamed() test failed!")
	}
}
<commit_msg>Update test<commit_after>package test

import (
	"fmt"
	"strings"
	"time"

	"github.com/nfrance-conseil/zeplic/lib"
	"testing"
)

func TestDatasetName(t *testing.T) {
	name := lib.DatasetName("tank/test@SNAP")
	if name != "tank/test" {
		t.Errorf("DatasetName() test failed!")
	}
}

func TestSnapName(t *testing.T) {
	name := lib.SnapName("SNAP")
	year, month, day := time.Now().Date()
	get := fmt.Sprintf("%s_%d-%s-%02d", "SNAP", year, month, day)
	if strings.Contains(name, get) == false {
		t.Errorf("SnapName() test failed!")
	}
}

func TestRenamed(t *testing.T) {
	renamed := lib.Renamed("tank/test@SNAP1", "tank/test@SNAP2")
	if renamed == false {
		t.Errorf("Renamed() test failed!")
	}
}
<|endoftext|>


<commit_before>package sf

type Sprite struct {
	texture *Texture
	rect    Rect
	verts   [4]Vertex
	//T       *Transformable // TODO? This is a workaround since we can't inherit.
}

func NewSprite(t *Texture) *Sprite {
	spr := &Sprite{}
	spr.SetTexture(t)
	spr.SetColor(Color{255, 255, 255, 255})
	//spr.T = NewTransformable()

	return spr
}

func (s *Sprite) Render(t *RenderTarget, states RenderStates) {
	states.Texture = s.texture
	//states.transform.Combine(s.T.Transform())
	t.Render(s.verts[:], Quads, states)
}

func (s *Sprite) SetTexture(t *Texture) {
	s.SetTextureRect(Rect{0, 0, t.Size().X, t.Size().Y})
	s.texture = t
}

func (s *Sprite) SetTextureRect(rect Rect) {
	if rect != s.rect {
		s.rect = rect
		s.updatePositions()
		s.updateTexCoords()
	}
}

func (s *Sprite) SetColor(Color Color) {
	s.verts[0].Color = Color
	s.verts[1].Color = Color
	s.verts[2].Color = Color
	s.verts[3].Color = Color
}

func (s *Sprite) Texture() *Texture {
	return s.texture
}

func (s *Sprite) LocalBounds() Rect {
	return s.rect
}

func (s *Sprite) updatePositions() {
	s.verts[0].Pos = Vector2{}
	s.verts[1].Pos = Vector2{0, s.rect.H}
	s.verts[2].Pos = Vector2{s.rect.W, s.rect.H}
	s.verts[3].Pos = Vector2{s.rect.W, 0}
}

func (s *Sprite) updateTexCoords() {
	left := s.rect.Left
	right := left + s.rect.W
	top := s.rect.Top
	bottom := top + s.rect.H

	s.verts[0].TexCoords = Vector2{left, top}
	s.verts[1].TexCoords = Vector2{left, bottom}
	s.verts[2].TexCoords = Vector2{right, bottom}
	s.verts[3].TexCoords = Vector2{right, top}
}
<commit_msg>Stuff<commit_after>package sf

type Sprite struct {
	texture *Texture
	rect    Rect
	verts   [4]Vertex
	//T       *Transformable // TODO? This is a workaround since we can't inherit.
}

func NewSprite(t *Texture) *Sprite {
	spr := &Sprite{}
	spr.SetTexture(t)
	spr.SetColor(Color{255, 255, 255, 255})
	//spr.T = NewTransformable()

	return spr
}

func (s *Sprite) Render(t *RenderTarget, states RenderStates) {
	states.Texture = s.texture
	//states.Transform.Combine(s.T.Transform())
	t.Render(s.verts[:], Quads, states)
}

func (s *Sprite) SetTexture(t *Texture) {
	s.SetTextureRect(Rect{0, 0, t.Size().X, t.Size().Y})
	s.texture = t
}

func (s *Sprite) SetTextureRect(rect Rect) {
	if rect != s.rect {
		s.rect = rect
		s.updatePositions()
		s.updateTexCoords()
	}
}

func (s *Sprite) SetColor(Color Color) {
	s.verts[0].Color = Color
	s.verts[1].Color = Color
	s.verts[2].Color = Color
	s.verts[3].Color = Color
}

func (s *Sprite) Texture() *Texture {
	return s.texture
}

func (s *Sprite) LocalBounds() Rect {
	return s.rect
}

func (s *Sprite) updatePositions() {
	s.verts[0].Pos = Vector2{}
	s.verts[1].Pos = Vector2{0, s.rect.H}
	s.verts[2].Pos = Vector2{s.rect.W, s.rect.H}
	s.verts[3].Pos = Vector2{s.rect.W, 0}
}

func (s *Sprite) updateTexCoords() {
	left := s.rect.Left
	right := left + s.rect.W
	top := s.rect.Top
	bottom := top + s.rect.H

	s.verts[0].TexCoords = Vector2{left, top}
	s.verts[1].TexCoords = Vector2{left, bottom}
	s.verts[2].TexCoords = Vector2{right, bottom}
	s.verts[3].TexCoords = Vector2{right, top}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package e2e_node

import (
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/kubernetes/pkg/features"
	kubeletconfig "k8s.io/kubernetes/pkg/kubelet/apis/config"
	"k8s.io/kubernetes/test/e2e/common"
	"k8s.io/kubernetes/test/e2e/framework"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

const (
	defaultObservationTimeout = time.Minute * 4
)

var _ = framework.KubeDescribe("StartupProbe [Serial] [Disruptive] [NodeFeature:StartupProbe]", func() {
	f := framework.NewDefaultFramework("critical-pod-test")
	var podClient *framework.PodClient

	/*
		These tests are located here as they require tempSetCurrentKubeletConfig to enable the feature gate for startupProbe.
		Once the feature gate has been removed, these tests should come back to test/e2e/common/container_probe.go.
	*/
	ginkgo.Context("when a container has a startup probe", func() {
		tempSetCurrentKubeletConfig(f, func(initialConfig *kubeletconfig.KubeletConfiguration) {
			if initialConfig.FeatureGates == nil {
				initialConfig.FeatureGates = make(map[string]bool)
			}
			initialConfig.FeatureGates[string(features.StartupProbe)] = true
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, delayed by startup probe
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat the non-existent /tmp/health file. Liveness probe MUST NOT fail until startup probe expires.
		*/
		framework.ConformanceIt("should *not* be restarted with a exec \"cat /tmp/health\" because startup probe delays it [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    60,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 0, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, delayed by startup probe
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat the non-existent /tmp/health file. Liveness probe MUST fail after startup probe expires. The Pod MUST now be killed and restarted incrementing restart count to 1.
		*/
		framework.ConformanceIt("should be restarted with a exec \"cat /tmp/health\" because startup probe does not delay it long enough [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    3,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 1, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, startup finished restart
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat /temp/health file. The Container is started by creating /tmp/startup after 10 seconds, triggering liveness probe to fail. The Pod MUST now be killed and restarted incrementing restart count to 1.
		*/
		framework.ConformanceIt("should be restarted with a exec \"cat /tmp/health\" after startup probe succeeds it [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 10; echo ok >/tmp/startup; sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/startup"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    60,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 1, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod readiness probe, delayed by startup probe
			Description: A Pod is created with startup and readiness probes. The Container is started by creating /tmp/startup after 45 seconds, delaying the ready state by this amount of time. This is similar to the "Pod readiness probe, with initial delay" test.
		*/
		framework.ConformanceIt("should not be ready until startupProbe succeeds [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "echo ok >/tmp/health; sleep 45; echo ok >/tmp/startup; sleep 600"}
			readinessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 0,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/startup"},
					},
				},
				InitialDelaySeconds: 0,
				FailureThreshold:    60,
			}
			p := podClient.Create(startupPodSpec(startupProbe, readinessProbe, nil, cmd))

			p, err := podClient.Get(p.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			f.WaitForPodReady(p.Name)
			isReady, err := testutils.PodRunningReady(p)
			framework.ExpectNoError(err)
			gomega.Expect(isReady).To(gomega.BeTrue(), "pod should be ready")

			// We assume the pod became ready when the container became ready. This
			// is true for a single container pod.
			readyTime, err := common.GetTransitionTimeForReadyCondition(p)
			framework.ExpectNoError(err)
			startedTime, err := common.GetContainerStartedTime(p, "busybox")
			framework.ExpectNoError(err)

			framework.Logf("Container started at %v, pod became ready at %v", startedTime, readyTime)
			if readyTime.Sub(startedTime) < 40*time.Second {
				framework.Failf("Pod became ready before startupProbe succeeded")
			}
		})
	})
})

func startupPodSpec(startupProbe, readinessProbe, livenessProbe *v1.Probe, cmd []string) *v1.Pod {
	return &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:   "startup-" + string(uuid.NewUUID()),
			Labels: map[string]string{"test": "startup"},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:           "busybox",
					Image:          imageutils.GetE2EImage(imageutils.BusyBox),
					Command:        cmd,
					LivenessProbe:  livenessProbe,
					ReadinessProbe: readinessProbe,
					StartupProbe:   startupProbe,
				},
			},
		},
	}
}
<commit_msg>Mark startupProbe test as NodeAlphaFeature and fix podClient instanciation<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package e2e_node

import (
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/kubernetes/pkg/features"
	kubeletconfig "k8s.io/kubernetes/pkg/kubelet/apis/config"
	"k8s.io/kubernetes/test/e2e/common"
	"k8s.io/kubernetes/test/e2e/framework"
	testutils "k8s.io/kubernetes/test/utils"
	imageutils "k8s.io/kubernetes/test/utils/image"

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega"
)

const (
	defaultObservationTimeout = time.Minute * 4
)

var _ = framework.KubeDescribe("StartupProbe [Serial] [Disruptive] [NodeAlphaFeature:StartupProbe]", func() {
	f := framework.NewDefaultFramework("startup-probe-test")
	var podClient *framework.PodClient
	ginkgo.BeforeEach(func() {
		podClient = f.PodClient()
	})

	/*
		These tests are located here as they require tempSetCurrentKubeletConfig to enable the feature gate for startupProbe.
		Once the feature gate has been removed, these tests should come back to test/e2e/common/container_probe.go.
	*/
	ginkgo.Context("when a container has a startup probe", func() {
		tempSetCurrentKubeletConfig(f, func(initialConfig *kubeletconfig.KubeletConfiguration) {
			if initialConfig.FeatureGates == nil {
				initialConfig.FeatureGates = make(map[string]bool)
			}
			initialConfig.FeatureGates[string(features.StartupProbe)] = true
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, delayed by startup probe
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat the non-existent /tmp/health file. Liveness probe MUST NOT fail until startup probe expires.
		*/
		framework.ConformanceIt("should *not* be restarted with a exec \"cat /tmp/health\" because startup probe delays it [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    60,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 0, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, delayed by startup probe
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat the non-existent /tmp/health file. Liveness probe MUST fail after startup probe expires. The Pod MUST now be killed and restarted incrementing restart count to 1.
		*/
		framework.ConformanceIt("should be restarted with a exec \"cat /tmp/health\" because startup probe does not delay it long enough [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    3,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 1, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod liveness probe, using local file, startup finished restart
			Description: A Pod is created with liveness probe that uses ‘exec’ command to cat /temp/health file. The Container is started by creating /tmp/startup after 10 seconds, triggering liveness probe to fail. The Pod MUST now be killed and restarted incrementing restart count to 1.
		*/
		framework.ConformanceIt("should be restarted with a exec \"cat /tmp/health\" after startup probe succeeds it [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "sleep 10; echo ok >/tmp/startup; sleep 600"}
			livenessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    1,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/startup"},
					},
				},
				InitialDelaySeconds: 15,
				FailureThreshold:    60,
			}
			pod := startupPodSpec(startupProbe, nil, livenessProbe, cmd)
			common.RunLivenessTest(f, pod, 1, defaultObservationTimeout)
		})

		/*
			Release : v1.16
			Testname: Pod readiness probe, delayed by startup probe
			Description: A Pod is created with startup and readiness probes. The Container is started by creating /tmp/startup after 45 seconds, delaying the ready state by this amount of time. This is similar to the "Pod readiness probe, with initial delay" test.
		*/
		framework.ConformanceIt("should not be ready until startupProbe succeeds [NodeConformance]", func() {
			cmd := []string{"/bin/sh", "-c", "echo ok >/tmp/health; sleep 45; echo ok >/tmp/startup; sleep 600"}
			readinessProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/health"},
					},
				},
				InitialDelaySeconds: 0,
			}
			startupProbe := &v1.Probe{
				Handler: v1.Handler{
					Exec: &v1.ExecAction{
						Command: []string{"cat", "/tmp/startup"},
					},
				},
				InitialDelaySeconds: 0,
				FailureThreshold:    60,
			}
			p := podClient.Create(startupPodSpec(startupProbe, readinessProbe, nil, cmd))

			p, err := podClient.Get(p.Name, metav1.GetOptions{})
			framework.ExpectNoError(err)

			f.WaitForPodReady(p.Name)
			isReady, err := testutils.PodRunningReady(p)
			framework.ExpectNoError(err)
			gomega.Expect(isReady).To(gomega.BeTrue(), "pod should be ready")

			// We assume the pod became ready when the container became ready. This
			// is true for a single container pod.
			readyTime, err := common.GetTransitionTimeForReadyCondition(p)
			framework.ExpectNoError(err)
			startedTime, err := common.GetContainerStartedTime(p, "busybox")
			framework.ExpectNoError(err)

			framework.Logf("Container started at %v, pod became ready at %v", startedTime, readyTime)
			if readyTime.Sub(startedTime) < 40*time.Second {
				framework.Failf("Pod became ready before startupProbe succeeded")
			}
		})
	})
})

func startupPodSpec(startupProbe, readinessProbe, livenessProbe *v1.Probe, cmd []string) *v1.Pod {
	return &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:   "startup-" + string(uuid.NewUUID()),
			Labels: map[string]string{"test": "startup"},
		},
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				{
					Name:           "busybox",
					Image:          imageutils.GetE2EImage(imageutils.BusyBox),
					Command:        cmd,
					LivenessProbe:  livenessProbe,
					ReadinessProbe: readinessProbe,
					StartupProbe:   startupProbe,
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>package entities

// Copyright 2019 Google LLC
//
// 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
//
// 	https://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.

import (
	"context"
	"testing"

	"github.com/googlecloudplatform/threat-automation/clients/stubs"
	compute "google.golang.org/api/compute/v1"
)

func TestCreateDiskSnapshot(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
		disk      = "test-disk"
		snapshot  = "test-snapshot"
	)
	tests := []struct {
		name             string
		expectedError    error
		expectedResponse *compute.Snapshot
	}{
		{
			name:             "test",
			expectedError:    nil,
			expectedResponse: &compute.Snapshot{},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}
			computeStub.SavedCreateSnapshots = make(map[string]compute.Snapshot)
			ctx := context.Background()
			h := NewHost(computeStub)
			if _, err := h.CreateDiskSnapshot(ctx, projectID, zone, disk, snapshot); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}

func TestStartVm(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
	)
	tests := []struct {
		name          string
		instanceName  string
		expectedError error
	}{
		{
			name:          "test if starts successfully",
			instanceName:  "existentVm",
			expectedError: nil,
		},
		{
			name:          "should notify in case of error",
			instanceName:  "nonexistent",
			expectedError: stubs.ErrNonexistentVM,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}

			ctx := context.Background()
			h := NewHost(computeStub)

			if _, err := h.StartInstance(ctx, projectID, zone, tt.instanceName); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}

func TestDeleteVm(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
	)
	tests := []struct {
		name          string
		instanceName  string
		expectedError error
	}{
		{
			name:          "test if starts successfully",
			instanceName:  "existentVm",
			expectedError: nil,
		},
		{
			name:          "should notify in case of error",
			instanceName:  "nonexistent",
			expectedError: stubs.ErrNonexistentVM,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}

			ctx := context.Background()
			h := NewHost(computeStub)

			if _, err := h.DeleteInstance(ctx, projectID, zone, tt.instanceName); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}
<commit_msg>Fixes comments<commit_after>package entities

// Copyright 2019 Google LLC
//
// 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
//
// 	https://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.

import (
	"context"
	"testing"

	"github.com/googlecloudplatform/threat-automation/clients/stubs"
	compute "google.golang.org/api/compute/v1"
)

func TestCreateDiskSnapshot(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
		disk      = "test-disk"
		snapshot  = "test-snapshot"
	)
	tests := []struct {
		name             string
		expectedError    error
		expectedResponse *compute.Snapshot
	}{
		{
			name:             "test",
			expectedError:    nil,
			expectedResponse: &compute.Snapshot{},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}
			computeStub.SavedCreateSnapshots = make(map[string]compute.Snapshot)
			ctx := context.Background()
			h := NewHost(computeStub)
			if _, err := h.CreateDiskSnapshot(ctx, projectID, zone, disk, snapshot); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}

func TestStartVm(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
	)
	tests := []struct {
		name          string
		instanceName  string
		expectedError error
	}{
		{
			name:          "test if starts successfully",
			instanceName:  "existentVm",
			expectedError: nil,
		},
		{
			name:          "should notify in case of error",
			instanceName:  "nonexistent",
			expectedError: stubs.ErrNonexistentVM,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}

			ctx := context.Background()
			h := NewHost(computeStub)

			if _, err := h.StartInstance(ctx, projectID, zone, tt.instanceName); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}

func TestDeleteVm(t *testing.T) {
	const (
		projectID = "test-project-id"
		zone      = "test-zone"
	)
	tests := []struct {
		name          string
		instanceName  string
		expectedError error
	}{
		{
			name:          "test if deletes successfully",
			instanceName:  "existentVm",
			expectedError: nil,
		},
		{
			name:          "should notify in case of error",
			instanceName:  "nonexistent",
			expectedError: stubs.ErrNonexistentVM,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			computeStub := &stubs.ComputeStub{}

			ctx := context.Background()
			h := NewHost(computeStub)

			if _, err := h.DeleteInstance(ctx, projectID, zone, tt.instanceName); err != tt.expectedError {
				t.Errorf("%v failed exp:%v got: %v", tt.name, tt.expectedError, err)
			}
		})
	}
}
<|endoftext|>


<commit_before>// Package cli is a library to help creating command line tools.
package cli

import (
	"flag"
	"fmt"
	"os"
)

// To add a module, implement this interface. Definition is the command
// definition. Exec is the behaviour that you want to implement as a command
type Command interface {
	Definition() string // usually it's the output for --help
	Exec(args []string) error
}

// Module is the shared structure of commands and sub-commands.
type Module struct {
	children map[string]*Module // Non-nil if sub-command
	command  Command            // Non-nil if command
}

// NewCLI returns a root Module that you can add commands and
// another modules (sub-commands).
func NewCLI() *Module {
	return &Module{
		children: make(map[string]*Module, 0),
	}
}

// AddCommand adds a new command this module.
func (m *Module) AddCommand(name string, command Command) {
	child := &Module{
		command: command,
	}
	m.children[name] = child
}

// AddSubCommand adds a new sub-command this module.
func (m *Module) AddSubCommand(name string) *Module {
	child := &Module{
		children: make(map[string]*Module, 0),
	}
	m.children[name] = child
	return child
}

// Run is the function that is intended to be run from main().
func (m *Module) Run() {
	flag.Parse()
	args := flag.Args()

	command, args, err := m.findCommand(args)
	if err != nil {
		exitErr(err)
	}

	err = command.Exec(args)
	if err != nil {
		exitErr(err)
	}

	os.Exit(0) // just to be explicit
}

func exitErr(err error) {
	fmt.Fprintf(os.Stderr, "%s\n", err.Error())
	os.Exit(1)
}

func (m *Module) findCommand(args []string) (Command, []string, error) {
	newArgs := args // this is the subset of args and will be returned with command

	// Iterate over args and update the module pointer "m"
	for _, arg := range args {
		if m.children == nil {
			// m is a command
			break
		}

		m = m.children[arg]
		newArgs = newArgs[1:]
	}

	if m == nil {
		return nil, nil, fmt.Errorf("Command not found")
	}

	// m is a command or sub-command we don't care because we are
	// returning Command interface
	return m, newArgs, nil
}

////////////////////////////////////////////////////////////////////////
// Methods below implement Command interface for Module (sub-command) //
////////////////////////////////////////////////////////////////////////

func (m *Module) Definition() string {
	if m.command != nil {
		// m is a command
		return m.command.Definition()
	}

	// m is a sub-command
	return fmt.Sprintf("Run to see sub-commands")
}

func (m *Module) Exec(args []string) error {
	if m.command != nil {
		// m is a command
		return m.command.Exec(args)
	}

	// m is a sub-command
	// Print command list
	fmt.Println("Possible commands:")
	for n, module := range m.children {
		fmt.Printf("  %-10s  ", n)

		if module.command != nil {
			fmt.Println(module.command.Definition())
		} else {
			fmt.Println(module.Definition())
		}
	}

	return nil
}
<commit_msg>fix kd cli panic<commit_after>// Package cli is a library to help creating command line tools.
package cli

import (
	"flag"
	"fmt"
	"os"
)

// To add a module, implement this interface. Definition is the command
// definition. Exec is the behaviour that you want to implement as a command
type Command interface {
	Definition() string // usually it's the output for --help
	Exec(args []string) error
}

// Module is the shared structure of commands and sub-commands.
type Module struct {
	children map[string]*Module // Non-nil if sub-command
	command  Command            // Non-nil if command
}

// NewCLI returns a root Module that you can add commands and
// another modules (sub-commands).
func NewCLI() *Module {
	return &Module{
		children: make(map[string]*Module, 0),
	}
}

// AddCommand adds a new command this module.
func (m *Module) AddCommand(name string, command Command) {
	child := &Module{
		command: command,
	}
	m.children[name] = child
}

// AddSubCommand adds a new sub-command this module.
func (m *Module) AddSubCommand(name string) *Module {
	child := &Module{
		children: make(map[string]*Module, 0),
	}
	m.children[name] = child
	return child
}

// Run is the function that is intended to be run from main().
func (m *Module) Run() {
	flag.Parse()
	args := flag.Args()

	command, args, err := m.findCommand(args)
	if err != nil {
		exitErr(err)
	}

	err = command.Exec(args)
	if err != nil {
		exitErr(err)
	}

	os.Exit(0) // just to be explicit
}

func exitErr(err error) {
	fmt.Fprintf(os.Stderr, "%s\n", err.Error())
	os.Exit(1)
}

func (m *Module) findCommand(args []string) (Command, []string, error) {
	newArgs := args // this is the subset of args and will be returned with command

	// Iterate over args and update the module pointer "m"
	for _, arg := range args {
		if m == nil || m.children == nil {
			// m is a command
			break
		}

		m = m.children[arg]
		newArgs = newArgs[1:]
	}

	if m == nil {
		return nil, nil, fmt.Errorf("Command not found")
	}

	// m is a command or sub-command we don't care because we are
	// returning Command interface
	return m, newArgs, nil
}

////////////////////////////////////////////////////////////////////////
// Methods below implement Command interface for Module (sub-command) //
////////////////////////////////////////////////////////////////////////

func (m *Module) Definition() string {
	if m.command != nil {
		// m is a command
		return m.command.Definition()
	}

	// m is a sub-command
	return fmt.Sprintf("Run to see sub-commands")
}

func (m *Module) Exec(args []string) error {
	if m.command != nil {
		// m is a command
		return m.command.Exec(args)
	}

	// m is a sub-command
	// Print command list
	fmt.Println("Possible commands:")
	for n, module := range m.children {
		fmt.Printf("  %-10s  ", n)

		if module.command != nil {
			fmt.Println(module.command.Definition())
		} else {
			fmt.Println(module.Definition())
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(net.IPv4(172, 16, 0, 2)), nil)
		count, _ := strconv.Atoi(os.Args[2])
		for i := 0; i < count; i++ {
			fmt.Println(i)
			vm := virt.VM{
				Id: bson.NewObjectId(),
				IP: utils.IntToIP(<-ipPoolFetch),
			}
			vm.Prepare(nil)
			vm.StartCommand().Run()
		}
	},
}

func main() {
	virt.LoadTemplates("templates")
	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<commit_msg>virt: improved vmtool.<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"koding/tools/utils"
	"koding/virt"
	"labix.org/v2/mgo/bson"
	"net"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

var actions = map[string]func(){
	"shutdown": func() {
		withAll("lxc-shutdown")
	},
	"stop": func() {
		withAll("lxc-stop")
	},
	"unprepare": func() {
		for _, vm := range selectVMs(os.Args[2]) {
			err := vm.Unprepare()
			fmt.Printf("%v: %v\n", vm, err)
		}
	},
	"create-test-vms": func() {
		startIP := net.IPv4(172, 16, 0, 2)
		if len(os.Args) >= 4 {
			startIP = net.ParseIP(os.Args[3])
		}
		ipPoolFetch, _ := utils.NewIntPool(utils.IPToInt(startIP), nil)
		count, _ := strconv.Atoi(os.Args[2])
		for i := 0; i < count; i++ {
			go func(i int) {
				vm := virt.VM{
					Id: bson.NewObjectId(),
					IP: utils.IntToIP(<-ipPoolFetch),
				}
				vm.Prepare(nil)
				vm.StartCommand().Run()
				fmt.Println(i)
			}(i)
		}
	},
}

func main() {
	virt.LoadTemplates("templates")
	action := actions[os.Args[1]]
	action()
}

func selectVMs(selector string) []*virt.VM {
	if selector == "all" {
		dirs, err := ioutil.ReadDir("/var/lib/lxc")
		if err != nil {
			panic(err)
		}
		vms := make([]*virt.VM, 0)
		for _, dir := range dirs {
			if strings.HasPrefix(dir.Name(), "vm-") {
				vms = append(vms, &virt.VM{Id: bson.ObjectIdHex(dir.Name()[3:])})
			}
		}
		return vms
	}
	fmt.Println("Invalid selector: " + selector)
	os.Exit(1)
	return nil
}

func withAll(action string) {
	for _, vm := range selectVMs(os.Args[2]) {
		cmd := exec.Command(action, "-n", vm.String())
		out, err := cmd.CombinedOutput()
		fmt.Println(strings.Join(cmd.Args, " ") + ":")
		fmt.Println(err)
		fmt.Println(string(out))
	}
}
<|endoftext|>


<commit_before>// written by Daniel Oaks <daniel@danieloaks.net>
// released under the ISC license

/*
Package ircfmt handles IRC formatting codes, escaping and unescaping.

This allows for a simpler representation of strings that contain colour codes,
bold codes, and such, without having to write and handle raw bytes when
assembling outgoing messages.

This lets you turn raw IRC messages into our escaped versions, and turn escaped
versions back into raw messages suitable for sending on IRC connections. This
is designed to be used on things like PRIVMSG / NOTICE commands, MOTD blocks,
and such.

The escape character we use in this library is the dollar sign ("$"), along
with the given escape characters:

	----------------------------
	 Name       | Escape | Raw
	----------------------------
	 Dollarsign |   $$   |  $
	 Bold       |   $b   | 0x02
	 Colour     |   $c   | 0x03
	 Italic     |   $i   | 0x1d
	 Underscore |   $u   | 0x1f
	 Reset      |   $r   | 0x0f
	----------------------------

Colours are escaped in a slightly different way, using the actual names of them
rather than just the raw numbers.

In our escaped format, the colours for the fore and background are contained in
square brackets after the colour ("$c") escape. For example:

	Red foreground:
		Escaped:  This is a $c[red]cool message!
		Raw:      This is a 0x034cool message!

	Blue foreground, green background:
		Escaped:  This is a $c[blue,green]rad message!
		Raw:      This is a 0x032,3rad message!

When assembling a raw message, we make sure to use the full colour code
("02" vs just "2") when it could become confused due to numbers just after the
colour escape code. For instance, lines like this will be unescaped correctly:

	No number after colour escape:
		Escaped:  This is a $c[red]cool message!
		Raw:      This is a 0x034cool message!

	Number after colour escape:
		Escaped:  This is $c[blue]20% cooler!
		Raw:      This is 0x030220% cooler

Here are the colour names and codes we recognise:

	--------------------
	 Code | Name
	--------------------
	  00  | white
	  01  | black
	  02  | blue
	  03  | green
	  04  | red
	  05  | brown
	  06  | magenta
	  07  | orange
	  08  | yellow
	  09  | light green
	  10  | cyan
	  11  | light cyan
	  12  | light blue
	  13  | pink
	  14  | grey
	  15  | light grey
	--------------------

This package is in alpha.
*/
package ircfmt
<commit_msg>ircfmt: Beta<commit_after>// written by Daniel Oaks <daniel@danieloaks.net>
// released under the ISC license

/*
Package ircfmt handles IRC formatting codes, escaping and unescaping.

This allows for a simpler representation of strings that contain colour codes,
bold codes, and such, without having to write and handle raw bytes when
assembling outgoing messages.

This lets you turn raw IRC messages into our escaped versions, and turn escaped
versions back into raw messages suitable for sending on IRC connections. This
is designed to be used on things like PRIVMSG / NOTICE commands, MOTD blocks,
and such.

The escape character we use in this library is the dollar sign ("$"), along
with the given escape characters:

	----------------------------
	 Name       | Escape | Raw
	----------------------------
	 Dollarsign |   $$   |  $
	 Bold       |   $b   | 0x02
	 Colour     |   $c   | 0x03
	 Italic     |   $i   | 0x1d
	 Underscore |   $u   | 0x1f
	 Reset      |   $r   | 0x0f
	----------------------------

Colours are escaped in a slightly different way, using the actual names of them
rather than just the raw numbers.

In our escaped format, the colours for the fore and background are contained in
square brackets after the colour ("$c") escape. For example:

	Red foreground:
		Escaped:  This is a $c[red]cool message!
		Raw:      This is a 0x034cool message!

	Blue foreground, green background:
		Escaped:  This is a $c[blue,green]rad message!
		Raw:      This is a 0x032,3rad message!

When assembling a raw message, we make sure to use the full colour code
("02" vs just "2") when it could become confused due to numbers just after the
colour escape code. For instance, lines like this will be unescaped correctly:

	No number after colour escape:
		Escaped:  This is a $c[red]cool message!
		Raw:      This is a 0x034cool message!

	Number after colour escape:
		Escaped:  This is $c[blue]20% cooler!
		Raw:      This is 0x030220% cooler

Here are the colour names and codes we recognise:

	--------------------
	 Code | Name
	--------------------
	  00  | white
	  01  | black
	  02  | blue
	  03  | green
	  04  | red
	  05  | brown
	  06  | magenta
	  07  | orange
	  08  | yellow
	  09  | light green
	  10  | cyan
	  11  | light cyan
	  12  | light blue
	  13  | pink
	  14  | grey
	  15  | light grey
	--------------------

This package is in beta and the API should not change.
*/
package ircfmt
<|endoftext|>


<commit_before>package entities

import (
    "math"
    "encoding/json"
    "fmt"
    "log"
)

type Planet struct {
    coords []int
    Texture int
    Size int
    ShipCount int
    MaxShipCount int
    Owner string
}

func (self Planet) String() string {
    // TODO: Improve this
    return self.Owner
}

func (self Planet) GetKey() string {
    return fmt.Sprintf("planet.%d_%d", self.coords[0], self.coords[1])
}

func (self Planet) GetCoords() []int {
    return self.coords
}

func (self Planet) Serialize() (string, []byte) {
    result, err := json.Marshal(self)
    if err != nil {
        log.Fatal(err)
    }
    return self.GetKey(), result
}

func GeneratePlanets(hash string, sun_position []int) ([]Planet, *Planet) {

    hashElement := func(index int) float64 {
        return float64(hash[index]) - 48
    }

    result := []Planet{}
    ring_offset := float64(80)
    planet_radius := float64(50)

    for ix:=0; ix<9; ix++ {
        planet_in_creation := Planet{[]int{0,0}, 0, 0, 0, 0, ""}
        ring_offset += planet_radius + hashElement(4 * ix)

        planet_in_creation.coords[0] = int(float64(sun_position[0]) + ring_offset * math.Cos(
            hashElement(4 * ix + 1) * 40))
        planet_in_creation.coords[1] = int(float64(sun_position[1]) + ring_offset * math.Sin(
            hashElement(4 * ix + 1) * 40))

        planet_in_creation.Texture = int(hashElement(4 * ix + 2))
        planet_in_creation.Size = 1 + int(hashElement(4 * ix + 3))
        result = append(result, planet_in_creation)
    }
    return result, &result[int(hashElement(37)) - 1]
}

<commit_msg>Write better Planet.String()<commit_after>package entities

import (
    "math"
    "encoding/json"
    "fmt"
    "log"
)

type Planet struct {
    coords []int
    Texture int
    Size int
    ShipCount int
    MaxShipCount int
    Owner string
}

func (self Planet) String() string {
    return fmt.Sprintf("Planet[%s, %s]", self.coords[0], self.coords[1])
}

func (self Planet) GetKey() string {
    return fmt.Sprintf("planet.%d_%d", self.coords[0], self.coords[1])
}

func (self Planet) GetCoords() []int {
    return self.coords
}

func (self Planet) Serialize() (string, []byte) {
    result, err := json.Marshal(self)
    if err != nil {
        log.Fatal(err)
    }
    return self.GetKey(), result
}

func GeneratePlanets(hash string, sun_position []int) ([]Planet, *Planet) {

    hashElement := func(index int) float64 {
        return float64(hash[index]) - 48
    }

    result := []Planet{}
    ring_offset := float64(80)
    planet_radius := float64(50)

    for ix:=0; ix<9; ix++ {
        planet_in_creation := Planet{[]int{0,0}, 0, 0, 0, 0, ""}
        ring_offset += planet_radius + hashElement(4 * ix)

        planet_in_creation.coords[0] = int(float64(sun_position[0]) + ring_offset * math.Cos(
            hashElement(4 * ix + 1) * 40))
        planet_in_creation.coords[1] = int(float64(sun_position[1]) + ring_offset * math.Sin(
            hashElement(4 * ix + 1) * 40))

        planet_in_creation.Texture = int(hashElement(4 * ix + 2))
        planet_in_creation.Size = 1 + int(hashElement(4 * ix + 3))
        result = append(result, planet_in_creation)
    }
    return result, &result[int(hashElement(37)) - 1]
}

<|endoftext|>


<commit_before>package core

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

func Test_GeneralMessageBroadcaster_Create_Events(t *testing.T) {
	model.SetDefaultClock(model.NewFakeClock())

	referentials := NewMemoryReferentials()
	referential := referentials.New("Un Referential Plutot Cool")
	referential.Start()
	defer referential.Stop()

	partner := referential.Partners().New("Un Partner tout autant cool")
	partner.Settings["remote_objectid_kind"] = "internal"
	partner.ConnectorTypes = []string{TEST_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER}
	partner.RefreshConnectors()
	referential.Partners().Save(partner)

	connector, _ := partner.Connector(TEST_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER)

	situation := referential.Model().Situations().New()
	//situation.Save()

	objectid := model.NewObjectID("internal", string(situation.Id()))
	situation.SetObjectID(objectid)

	reference := model.Reference{
		ObjectId: &objectid,
		Id:       string(situation.Id()),
		Type:     "Situation",
	}

	subs := partner.Subscriptions().New()
	subs.Save()
	subs.CreateAddNewResource(reference)
	subs.SetKind(string(subs.Id()))
	subs.Save()
	time.Sleep(10 * time.Millisecond) // Wait for the goRoutine to start ...

	situation.Save()

	time.Sleep(10 * time.Millisecond) // Wait for the Broadcaster and Connector to finish their work
	if len(connector.(*TestGeneralMessageSubscriptionBroadcaster).events) != 1 {
		t.Error("1 event should have been generated got: ", len(connector.(*TestGeneralMessageSubscriptionBroadcaster).events))
	}
}

func Test_GeneralMessageBroadcaster_Receive_Notify(t *testing.T) {
	// Create a test http server

	response := []byte{}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		response, _ = ioutil.ReadAll(r.Body)
		w.Header().Add("Content-Type", "text/xml")
	}))
	defer ts.Close()

	// Create a test http server
	referentials := NewMemoryReferentials()
	referential := referentials.New("Un Referential Plutot Cool")
	referential.Start()
	defer referential.broacasterManager.Stop()

	partner := referential.Partners().New("Un Partner tout autant cool")
	partner.Settings["remote_objectid_kind"] = "internal"
	partner.Settings["remote_credential"] = "external"
	partner.Settings["remote_url"] = ts.URL

	partner.ConnectorTypes = []string{SIRI_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER}
	partner.RefreshConnectors()
	referential.Partners().Save(partner)

	connector, _ := partner.Connector(SIRI_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER)
	connector.(*SIRIGeneralMessageSubscriptionBroadcaster).generalMessageBroadcaster = NewFakeGeneralMessageBroadcaster(connector.(*SIRIGeneralMessageSubscriptionBroadcaster))

	situation := referential.Model().Situations().New()

	objectid := model.NewObjectID("internal", string(situation.Id()))
	situation.SetObjectID(objectid)

	reference := model.Reference{
		ObjectId: &objectid,
		Id:       string(situation.Id()),
		Type:     "Situation",
	}

	subscription, _ := partner.Subscriptions().FindOrCreateByKind("This Kind should normaly be the exterior Subscription Id")
	subscription.CreateAddNewResource(reference)

	time.Sleep(10 * time.Millisecond) // Wait for the goRoutine to start ...
	situation.Save()

	time.Sleep(10 * time.Millisecond) // Wait for the Broadcaster and Connector to finish their work
	connector.(*SIRIGeneralMessageSubscriptionBroadcaster).generalMessageBroadcaster.Start()

	notify, _ := siri.NewXMLNotifyGeneralMessageFromContent(response)
	delivery := notify.GeneralMessagesDeliveries()

	if len(delivery) != 1 {
		t.Errorf("Should have received 1 delivery but got == %v", len(delivery))
	}

	if delivery[0].SubscriberRef() != "external" {
		t.Errorf("SubscriberRef should be external but got == %v", delivery[0].SubscriptionRef())
	}

	sv := delivery[0].XMLGeneralMessages()

	if len(sv) != 1 {
		t.Errorf("Should have received 1 GeneralMessage but got == %v\n%v", len(sv), sv)
	}
}
<commit_msg>Referential stop after test to clean goroutine<commit_after>package core

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/af83/edwig/model"
	"github.com/af83/edwig/siri"
)

func Test_GeneralMessageBroadcaster_Create_Events(t *testing.T) {
	model.SetDefaultClock(model.NewFakeClock())

	referentials := NewMemoryReferentials()
	referential := referentials.New("Un Referential Plutot Cool")
	referential.Start()
	defer referential.Stop()

	partner := referential.Partners().New("Un Partner tout autant cool")
	partner.Settings["remote_objectid_kind"] = "internal"
	partner.ConnectorTypes = []string{TEST_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER}
	partner.RefreshConnectors()
	referential.Partners().Save(partner)

	connector, _ := partner.Connector(TEST_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER)

	situation := referential.Model().Situations().New()
	//situation.Save()

	objectid := model.NewObjectID("internal", string(situation.Id()))
	situation.SetObjectID(objectid)

	reference := model.Reference{
		ObjectId: &objectid,
		Id:       string(situation.Id()),
		Type:     "Situation",
	}

	subs := partner.Subscriptions().New()
	subs.Save()
	subs.CreateAddNewResource(reference)
	subs.SetKind(string(subs.Id()))
	subs.Save()
	time.Sleep(10 * time.Millisecond) // Wait for the goRoutine to start ...

	situation.Save()

	time.Sleep(10 * time.Millisecond) // Wait for the Broadcaster and Connector to finish their work
	if len(connector.(*TestGeneralMessageSubscriptionBroadcaster).events) != 1 {
		t.Error("1 event should have been generated got: ", len(connector.(*TestGeneralMessageSubscriptionBroadcaster).events))
	}
}

func Test_GeneralMessageBroadcaster_Receive_Notify(t *testing.T) {
	// Create a test http server

	response := []byte{}
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		response, _ = ioutil.ReadAll(r.Body)
		w.Header().Add("Content-Type", "text/xml")
	}))
	defer ts.Close()

	// Create a test http server
	referentials := NewMemoryReferentials()
	referential := referentials.New("Un Referential Plutot Cool")
	referential.Start()
	defer referential.Stop()

	partner := referential.Partners().New("Un Partner tout autant cool")
	partner.Settings["remote_objectid_kind"] = "internal"
	partner.Settings["remote_credential"] = "external"
	partner.Settings["remote_url"] = ts.URL

	partner.ConnectorTypes = []string{SIRI_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER}
	partner.RefreshConnectors()
	referential.Partners().Save(partner)

	connector, _ := partner.Connector(SIRI_GENERAL_MESSAGE_SUBSCRIPTION_BROADCASTER)
	connector.(*SIRIGeneralMessageSubscriptionBroadcaster).generalMessageBroadcaster = NewFakeGeneralMessageBroadcaster(connector.(*SIRIGeneralMessageSubscriptionBroadcaster))

	situation := referential.Model().Situations().New()

	objectid := model.NewObjectID("internal", string(situation.Id()))
	situation.SetObjectID(objectid)

	reference := model.Reference{
		ObjectId: &objectid,
		Id:       string(situation.Id()),
		Type:     "Situation",
	}

	subscription, _ := partner.Subscriptions().FindOrCreateByKind("This Kind should normaly be the exterior Subscription Id")
	subscription.CreateAddNewResource(reference)

	time.Sleep(10 * time.Millisecond) // Wait for the goRoutine to start ...
	situation.Save()

	time.Sleep(10 * time.Millisecond) // Wait for the Broadcaster and Connector to finish their work
	connector.(*SIRIGeneralMessageSubscriptionBroadcaster).generalMessageBroadcaster.Start()

	notify, _ := siri.NewXMLNotifyGeneralMessageFromContent(response)
	delivery := notify.GeneralMessagesDeliveries()

	if len(delivery) != 1 {
		t.Errorf("Should have received 1 delivery but got == %v", len(delivery))
	}

	if delivery[0].SubscriberRef() != "external" {
		t.Errorf("SubscriberRef should be external but got == %v", delivery[0].SubscriptionRef())
	}

	sv := delivery[0].XMLGeneralMessages()

	if len(sv) != 1 {
		t.Errorf("Should have received 1 GeneralMessage but got == %v\n%v", len(sv), sv)
	}
}
<|endoftext|>


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

package dhcp

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"runtime"
	"unsafe"

	"golang.org/x/sys/unix"
)

var (
	protoAll     = int(unix.ETH_P_ALL)
	byteOrder    = binary.ByteOrder(binary.BigEndian)
	macBroadcast = []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
	ipBroadcast  = []byte{0xff, 0xff, 0xff, 0xff}
)

func init() {
	// This kernel API is icky through and through. It wants the
	// ethernet protocol number in big-endian form, but receives it as
	// a native-endian integer. Thus, we need to do the moral
	// equivalent of htons().
	i := uint16(1)
	b := *(*byte)(unsafe.Pointer(&i))
	if b == 1 {
		protoAll = protoAll << 8
		byteOrder = binary.LittleEndian
	}
}

// LinuxConn implements Conn using Linux raw sockets.
//
// The advantage compared to PortableConn is that LinuxConn does not
// need to bind to any port, and so can be run alongside other DHCP
// services on the same machine. However, using it requires
// CAP_NET_RAW, whereas PortableConn doesn't.
type LinuxConn struct {
	port       uint16
	ethernetFd int // AF_PACKET SOCK_RAW socket
	ipFd       int // AF_INET SOCK_RAW IPPROTO_RAW socket
}

func closeLinuxConn(c *LinuxConn) {
	if c.ethernetFd != -1 {
		unix.Close(c.ethernetFd)
		c.ethernetFd = -1
	}
	if c.ipFd != -1 {
		unix.Close(c.ipFd)
		c.ipFd = -1
	}
}

// NewLinuxConn creates a LinuxConn that receives DHCP packets on the
// given UDP port (should typically be 67)
func NewLinuxConn(port uint16) (*LinuxConn, error) { // TODO: support for interface binding
	eth, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, protoAll)
	if err != nil {
		return nil, err
	}

	filter := []unix.SockFilter{
		{0x28, 0, 0, 12},           // Load ethernet frame type
		{0x15, 0, 8, 0x0800},       // Is IPv4?
		{0x30, 0, 0, 23},           // Load IP packet type
		{0x15, 0, 6, 17},           // Is UDP?
		{0x28, 0, 0, 20},           // Load fragment offset
		{0x45, 4, 0, 0x1fff},       // Is first/only fragment?
		{0xb1, 0, 0, 14},           // Jump to start of UDP header
		{0x48, 0, 0, 16},           // Load destination port
		{0x15, 0, 1, uint32(port)}, // Is correct port?
		{0x6, 0, 0, 0x40000},       // Yes, receive packet
		{0x6, 0, 0, 0},             // No, ignore packet
	}
	filterProg := &unix.SockFprog{
		Len:    uint16(len(filter)),
		Filter: &filter[0],
	}

	_, _, errno := unix.Syscall6(unix.SYS_SETSOCKOPT, uintptr(eth), uintptr(unix.SOL_SOCKET), uintptr(unix.SO_ATTACH_FILTER), uintptr(unsafe.Pointer(filterProg)), uintptr(unsafe.Sizeof(*filterProg)), 0)
	if errno != 0 {
		unix.Close(eth)
		return nil, errno
	}

	ip, err := unix.Socket(unix.AF_INET, unix.SOCK_RAW, unix.IPPROTO_RAW)
	if err != nil {
		unix.Close(eth)
	}

	ret := &LinuxConn{port, eth, ip}
	runtime.SetFinalizer(ret, closeLinuxConn)
	return ret, nil
}

// Close closes the connection.
func (c *LinuxConn) Close() error {
	closeLinuxConn(c)
	return nil
}

// RecvDHCP implements the Conn RecvDHCP method.
func (c *LinuxConn) RecvDHCP() (*Packet, *net.Interface, error) {
	buf := make([]byte, 1500)
	for {
		n, from, err := unix.Recvfrom(c.ethernetFd, buf, 0)
		if err != nil {
			return nil, nil, err
		}
		bs := buf[:n]
		// Advance past the ethernet, IP and UDP headers, to reach the
		// DHCP packet.
		off := 22 + 4*int(buf[14]&0xf)
		pkt, err := Unmarshal(bs[off:])
		if err != nil {
			// TODO: return temporary error to allow the server to log
			// stuff.
			continue
		}

		if err = validatePacket(bs, pkt); err != nil {
			// TODO: return temporary error to allow the server to log
			// stuff.
			continue
		}

		addr := from.(*unix.SockaddrLinklayer)
		intf, err := net.InterfaceByIndex(addr.Ifindex)
		if err != nil {
			return nil, nil, err
		}

		return pkt, intf, nil
	}
}

// SendDHCP implements the Conn SendDHCP method.
func (c *LinuxConn) SendDHCP(pkt *Packet, intf *net.Interface) error {
	payload, err := pkt.Marshal()
	if err != nil {
		return err
	}
	switch pkt.TxType() {
	case TxBroadcast:
		if intf == nil {
			return errors.New("packet needs to be broadcast, but no interface specified")
		}
		srcIP, err := interfaceIP(intf)
		if err != nil {
			return err
		}
		bs := assemblePacket(intf.HardwareAddr, macBroadcast, srcIP, ipBroadcast, c.port, 68, payload)
		addr := &unix.SockaddrLinklayer{
			Ifindex: intf.Index,
			Halen:   6,
		}
		copy(addr.Addr[:6], intf.HardwareAddr)
		if err = unix.Sendto(c.ethernetFd, bs, 0, addr); err != nil {
			return err
		}
	case TxRelayAddr:
		bs := assemblePacket(nil, nil, nil, pkt.RelayAddr, c.port, 67, payload)
		bs = bs[14:] // Skip the ethernet header
		addr := &unix.SockaddrInet4{}
		if err = unix.Sendto(c.ipFd, bs, 0, addr); err != nil {
			return err
		}
	case TxClientAddr:
		bs := assemblePacket(nil, nil, nil, pkt.ClientAddr, c.port, 68, payload)
		bs = bs[14:] // Skip the ethernet header
		addr := &unix.SockaddrInet4{}
		if err = unix.Sendto(c.ipFd, bs, 0, addr); err != nil {
			return err
		}
	case TxHardwareAddr:
		if intf == nil {
			return errors.New("packet needs to be transmitted to unconfigured client, but no interface specified")
		}
		srcIP, err := interfaceIP(intf)
		if err != nil {
			return err
		}
		bs := assemblePacket(intf.HardwareAddr, pkt.HardwareAddr, srcIP, pkt.YourAddr, c.port, 68, payload)

		addr := &unix.SockaddrLinklayer{
			Ifindex: intf.Index,
			Halen:   6,
		}
		copy(addr.Addr[:6], intf.HardwareAddr)
		if err = unix.Sendto(c.ethernetFd, bs, 0, addr); err != nil {
			return err
		}
	}

	return nil
}

func validatePacket(frame []byte, pkt *Packet) error {
	if pkt.RelayAddr != nil {
		// If the packet is from a relay, no validation is needed.
		return nil
	}

	// ciaddr must match the IP header's source IP (either an actual
	// address, or 0.0.0.0).
	if !bytes.Equal(pkt.ClientAddr, frame[24:28]) {
		return errors.New("ciaddr doesn't match packet source IP")
	}
	// chaddr must match the source MAC address
	if !bytes.Equal(pkt.HardwareAddr, frame[6:12]) {
		return errors.New("chaddr doesn't match packet source MAC")
	}

	return nil
}

func interfaceIP(intf *net.Interface) (net.IP, error) {
	addrs, err := intf.Addrs()
	if err != nil {
		return nil, err
	}

	// Try to find an IPv4 address to use, in the following order:
	// global unicast (includes rfc1918), link-local unicast,
	// loopback.
	fs := [](func(net.IP) bool){
		net.IP.IsGlobalUnicast,
		net.IP.IsLinkLocalUnicast,
		net.IP.IsLoopback,
	}
	for _, f := range fs {
		for _, a := range addrs {
			ipaddr, ok := a.(*net.IPNet)
			if !ok {
				continue
			}
			ip := ipaddr.IP.To4()
			if ip == nil {
				continue
			}
			if f(ip) {
				return ip, nil
			}
		}
	}

	return nil, fmt.Errorf("interface %s has no unicast address usable as a DHCP packet source", intf.Name)
}

func assemblePacket(srcMAC, dstMAC net.HardwareAddr, srcIP, dstIP net.IP, srcPort, dstPort uint16, payload []byte) []byte {
	buf := make([]byte, 42, 42+len(payload))

	// Ethernet header
	copy(buf[:6], dstMAC)
	copy(buf[6:12], srcMAC)
	binary.BigEndian.PutUint16(buf[12:14], 0x0800)

	// IP header
	buf[14] = (4 << 4) + 5                                          // IP version 4, 5-word header (20b)
	buf[15] = 0xc0                                                  // ToS CS6 (Network Control)
	binary.BigEndian.PutUint16(buf[16:18], uint16(28+len(payload))) // IP packet length
	binary.BigEndian.PutUint32(buf[18:22], 0x4000)                  // ID=0, frag_off=0, dont_fragment=1
	buf[22] = 64                                                    // TTL
	buf[23] = 17                                                    // Inner protocol: UDP
	copy(buf[26:30], srcIP)
	copy(buf[30:34], dstIP)

	var cksum uint32
	for i := 14; i < 34; i += 2 {
		cksum += uint32(binary.BigEndian.Uint16(buf[i : i+2]))
	}
	cksum = (cksum >> 16) + (cksum & 0xFFFF)
	binary.BigEndian.PutUint16(buf[24:26], ^uint16(cksum))

	// UDP header
	binary.BigEndian.PutUint16(buf[34:36], srcPort)                // Source port
	binary.BigEndian.PutUint16(buf[36:38], dstPort)                // Destination port
	binary.BigEndian.PutUint16(buf[38:40], uint16(8+len(payload))) // UDP length

	return append(buf, payload...)
}
<commit_msg>Checkpoint some further hacking on AF_PACKET, before bailing.<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package dhcp

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"runtime"
	"unsafe"

	"golang.org/x/sys/unix"
)

var (
	protoAll = int(unix.ETH_P_ALL)
	macbcast = []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
)

func init() {
	// This kernel API is icky through and through. It wants the
	// ethernet protocol number in big-endian form, but receives it as
	// a native-endian integer. Thus, we need to do the moral
	// equivalent of htons().
	i := uint16(1)
	b := *(*byte)(unsafe.Pointer(&i))
	if b == 1 {
		protoAll = protoAll << 8
	}
}

// LinuxConn implements Conn using Linux raw sockets.
//
// The advantage compared to PortableConn is that LinuxConn does not
// need to bind to any port, and so can be run alongside other DHCP
// services on the same machine. However, using it requires
// CAP_NET_RAW, whereas PortableConn doesn't.
type LinuxConn struct {
	port       int
	ethernetFd int // AF_PACKET SOCK_RAW socket
	ipFd       int // AF_INET SOCK_RAW IPPROTO_RAW socket
}

func closeLinuxConn(c *LinuxConn) {
	if c.ethernetFd != -1 {
		unix.Close(c.ethernetFd)
		c.ethernetFd = -1
	}
	if c.ipFd != -1 {
		unix.Close(c.ipFd)
		c.ipFd = -1
	}
}

// NewLinuxConn creates a LinuxConn that receives DHCP packets. TODO better
func NewLinuxConn(addr string) (*LinuxConn, error) {
	if addr == "" {
		addr = ":67"
	}
	udpAddr, err := net.ResolveUDPAddr("udp4", addr)
	if err != nil {
		return nil, err
	}
	udpAddr.IP = udpAddr.IP.To4()

	eth, err := unix.Socket(unix.AF_PACKET, unix.SOCK_RAW, protoAll)
	if err != nil {
		return nil, err
	}

	filter := filterPortOnly(udpAddr.Port)
	if udpAddr.IP != nil {
		filter = filterIPAndPort(udpAddr.IP, udpAddr.Port)
	}
	if err = attachFilter(eth, filter); err != nil {
		unix.Close(eth)
		return nil, err
	}

	ip, err := unix.Socket(unix.AF_INET, unix.SOCK_RAW, unix.IPPROTO_RAW)
	if err != nil {
		unix.Close(eth)
		return nil, err
	}

	// if err = attachFilter(ip, filterNoPackets()); err != nil {
	// 	unix.Close(eth)
	// 	unix.Close(ip)
	// 	return nil, err
	// }

	ret := &LinuxConn{
		port:       udpAddr.Port,
		ethernetFd: eth,
		ipFd:       ip,
	}
	runtime.SetFinalizer(ret, closeLinuxConn)
	return ret, nil
}

// Close closes the connection.
func (c *LinuxConn) Close() error {
	closeLinuxConn(c)
	return nil
}

// RecvDHCP implements the Conn RecvDHCP method.
func (c *LinuxConn) RecvDHCP() (*Packet, *net.Interface, error) {
	buf := make([]byte, 1500)
	for {
		n, from, err := unix.Recvfrom(c.ethernetFd, buf, 0)
		if err != nil {
			return nil, nil, err
		}
		bs := buf[:n]
		// Advance past the ethernet, IP and UDP headers, to reach the
		// DHCP packet.
		off := 22 + 4*int(buf[14]&0xf)
		pkt, err := Unmarshal(bs[off:])
		if err != nil {
			// TODO: return temporary error to allow the server to log
			// stuff.
			continue
		}

		if err = validatePacket(bs, pkt); err != nil {
			// TODO: return temporary error to allow the server to log
			// stuff.
			continue
		}

		addr := from.(*unix.SockaddrLinklayer)
		intf, err := net.InterfaceByIndex(addr.Ifindex)
		if err != nil {
			return nil, nil, err
		}

		return pkt, intf, nil
	}
}

// SendDHCP implements the Conn SendDHCP method.
func (c *LinuxConn) SendDHCP(pkt *Packet, intf *net.Interface) error {
	payload, err := pkt.Marshal()
	if err != nil {
		return err
	}
	switch pkt.TxType() {
	case TxBroadcast:
		if intf == nil {
			return errors.New("packet needs to be broadcast, but no interface specified")
		}
		srcIP, err := interfaceIP(intf)
		if err != nil {
			return err
		}
		bs := assemblePacket(intf.HardwareAddr, macbcast, srcIP, net.IPv4bcast, c.port, 68, payload)
		addr := &unix.SockaddrLinklayer{
			Ifindex: intf.Index,
			Halen:   6,
		}
		copy(addr.Addr[:6], intf.HardwareAddr)
		if err = unix.Sendto(c.ethernetFd, bs, 0, addr); err != nil {
			return err
		}
	case TxRelayAddr:
		bs := assemblePacket(nil, nil, nil, pkt.RelayAddr, c.port, 67, payload)
		bs = bs[14:] // Skip the ethernet header
		addr := &unix.SockaddrInet4{}
		copy(addr.Addr[:], pkt.RelayAddr.To4())
		if err = unix.Sendto(c.ipFd, bs, 0, addr); err != nil {
			return err
		}
	case TxClientAddr:
		bs := assemblePacket(nil, nil, nil, pkt.ClientAddr, c.port, 68, payload)
		bs = bs[14:] // Skip the ethernet header
		addr := &unix.SockaddrInet4{}
		if err = unix.Sendto(c.ipFd, bs, 0, addr); err != nil {
			return err
		}
	case TxHardwareAddr:
		if intf == nil {
			return errors.New("packet needs to be transmitted to unconfigured client, but no interface specified")
		}
		srcIP, err := interfaceIP(intf)
		if err != nil {
			return err
		}
		bs := assemblePacket(intf.HardwareAddr, pkt.HardwareAddr, srcIP, pkt.YourAddr, c.port, 68, payload)

		addr := &unix.SockaddrLinklayer{
			Ifindex: intf.Index,
			Halen:   6,
		}
		copy(addr.Addr[:6], intf.HardwareAddr)
		if err = unix.Sendto(c.ethernetFd, bs, 0, addr); err != nil {
			return err
		}
	}

	return nil
}

func validatePacket(frame []byte, pkt *Packet) error {
	if pkt.RelayAddr != nil {
		// If the packet is from a relay, no validation is needed.
		return nil
	}

	// ciaddr must match the IP header's source IP (either an actual
	// address, or 0.0.0.0).
	if !bytes.Equal(pkt.ClientAddr, frame[24:28]) {
		return errors.New("ciaddr doesn't match packet source IP")
	}
	// chaddr must match the source MAC address
	if !bytes.Equal(pkt.HardwareAddr, frame[6:12]) {
		return errors.New("chaddr doesn't match packet source MAC")
	}

	return nil
}

func interfaceIP(intf *net.Interface) (net.IP, error) {
	addrs, err := intf.Addrs()
	if err != nil {
		return nil, err
	}

	// Try to find an IPv4 address to use, in the following order:
	// global unicast (includes rfc1918), link-local unicast,
	// loopback.
	fs := [](func(net.IP) bool){
		net.IP.IsGlobalUnicast,
		net.IP.IsLinkLocalUnicast,
		net.IP.IsLoopback,
	}
	for _, f := range fs {
		for _, a := range addrs {
			ipaddr, ok := a.(*net.IPNet)
			if !ok {
				continue
			}
			ip := ipaddr.IP.To4()
			if ip == nil {
				continue
			}
			if f(ip) {
				return ip, nil
			}
		}
	}

	return nil, fmt.Errorf("interface %s has no unicast address usable as a DHCP packet source", intf.Name)
}

func assemblePacket(srcMAC, dstMAC net.HardwareAddr, srcIP, dstIP net.IP, srcPort, dstPort int, payload []byte) []byte {
	buf := make([]byte, 42, 42+len(payload))

	// Ethernet header
	copy(buf[:6], dstMAC)
	copy(buf[6:12], srcMAC)
	binary.BigEndian.PutUint16(buf[12:14], 0x0800)

	// IP header
	buf[14] = (4 << 4) + 5                                          // IP version 4, 5-word header (20b)
	buf[15] = 0xc0                                                  // ToS CS6 (Network Control)
	binary.BigEndian.PutUint16(buf[16:18], uint16(28+len(payload))) // IP packet length
	binary.BigEndian.PutUint32(buf[18:22], 0x4000)                  // ID=0, frag_off=0, dont_fragment=1
	buf[22] = 64                                                    // TTL
	buf[23] = 17                                                    // Inner protocol: UDP
	copy(buf[26:30], srcIP.To4())
	copy(buf[30:34], dstIP.To4())
	fmt.Println(buf[30:34])

	var cksum uint32
	for i := 14; i < 34; i += 2 {
		cksum += uint32(binary.BigEndian.Uint16(buf[i : i+2]))
	}
	cksum = (cksum >> 16) + (cksum & 0xFFFF)
	binary.BigEndian.PutUint16(buf[24:26], ^uint16(cksum))

	// UDP header
	binary.BigEndian.PutUint16(buf[34:36], uint16(srcPort))        // Source port
	binary.BigEndian.PutUint16(buf[36:38], uint16(dstPort))        // Destination port
	binary.BigEndian.PutUint16(buf[38:40], uint16(8+len(payload))) // UDP length

	return append(buf, payload...)
}

func attachFilter(fd int, filter *unix.SockFprog) error {
	_, _, errno := unix.Syscall6(unix.SYS_SETSOCKOPT, uintptr(fd), uintptr(unix.SOL_SOCKET), uintptr(unix.SO_ATTACH_FILTER), uintptr(unsafe.Pointer(filter)), uintptr(unsafe.Sizeof(*filter)), 0)
	if errno != 0 {
		return errno
	}
	return nil
}

func filterPortOnly(port int) *unix.SockFprog {
	// This filter comes from:
	// tcpdump -dd 'ip and udp dst port 68'
	filter := []unix.SockFilter{
		{0x28, 0, 0, 12},           // Load ethernet frame type
		{0x15, 0, 8, 0x0800},       // Is IPv4?
		{0x30, 0, 0, 23},           // Load IP packet type
		{0x15, 0, 6, 17},           // Is UDP?
		{0x28, 0, 0, 20},           // Load fragment offset
		{0x45, 4, 0, 0x1fff},       // Is first/only fragment?
		{0xb1, 0, 0, 14},           // Jump to start of UDP header
		{0x48, 0, 0, 16},           // Load destination port
		{0x15, 0, 1, uint32(port)}, // Is correct port?
		{0x6, 0, 0, 0x40000},       // Yes, receive packet
		{0x6, 0, 0, 0},             // No, ignore packet
	}
	return &unix.SockFprog{
		Len:    uint16(len(filter)),
		Filter: &filter[0],
	}
}

func filterIPAndPort(dstIP net.IP, port int) *unix.SockFprog {
	d := binary.BigEndian.Uint32([]byte(dstIP.To4()))
	// This filter comes from:
	// tcpdump -dd 'ip and udp and (dst 192.168.2.2 or dst 255.255.255.255) and dst port 68'
	filter := []unix.SockFilter{
		{0x28, 0, 0, 12},           // Load ethernet frame type
		{0x15, 0, 11, 0x0800},      // Is IPv4?
		{0x30, 0, 0, 23},           // Load IP packet type
		{0x15, 0, 9, 17},           // Is UDP?
		{0x20, 0, 0, 30},           // Load destination IP
		{0x15, 1, 0, d},            // Is target IP?
		{0x15, 0, 6, 0xffffffff},   // Is Broadcast?
		{0x28, 0, 0, 20},           // Load fragment offset
		{0x45, 4, 0, 0x1fff},       // Is first/only fragment?
		{0xb1, 0, 0, 14},           // Jump to start of UDP header
		{0x48, 0, 0, 16},           // Load destination port
		{0x15, 0, 1, uint32(port)}, // Is correct port?
		{0x6, 0, 0, 0x40000},       // Yes, receive packet
		{0x6, 0, 0, 0},             // No, ignore packet
	}
	return &unix.SockFprog{
		Len:    uint16(len(filter)),
		Filter: &filter[0],
	}
}

func filterNoPackets() *unix.SockFprog {
	filter := []unix.SockFilter{
		{0x6, 0, 0, 0}, // ignore packet
	}
	return &unix.SockFprog{
		Len:    uint16(len(filter)),
		Filter: &filter[0],
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/k0kubun/pp"
	"github.com/pkg/errors"
)

var cnt int

func subErr(err error) error {
	err = errors.Wrapf(err, "Suberror #%d", cnt)
	cnt++
	return err
}

func getError() error {
	err := errors.New("my error")
	err = subErr(err)
	return errors.Wrap(err, "exit")
}

func main() {
	fmt.Println("Hello world!")
	err := getError()
	err = errors.Wrap(err, "open failed")
	err = subErr(err)
	err = errors.Wrap(err, "read config failed")

	pp.Println("Cause: ", errors.Cause(err))
	err = errors.Wrap(err, "New message")
	pp.Println("Error: ", err)
	fmt.Printf("[%+v]\n", err)
	fmt.Printf("{%+v}\n", errors.Cause(err))
}
<commit_msg>use wrapper instead internal error<commit_after>package main

import (
	"fmt"

	"github.com/k0kubun/pp"
	"github.com/pkg/errors"
)

var cnt int

func subErr(err error) error {
	err = errors.Wrapf(err, "Suberror #%d", cnt)
	cnt++
	return err
}

func getError() error {
	err := fmt.Errorf("my error")
	err = subErr(err)
	return errors.Wrap(err, "exit")
}

func main() {
	fmt.Println("Hello world!")
	err := getError()
	err = errors.Wrap(err, "open failed")
	err = subErr(err)
	err = errors.Wrap(err, "read config failed")

	pp.Println("Cause: ", errors.Cause(err))
	err = errors.Wrap(err, "New message")
	pp.Println("Error: ", err)
	fmt.Printf("[%+v]\n", err)
	fmt.Printf("{%+v}\n", errors.Cause(err))
	fmt.Printf("[%v]\n", err)
}
<|endoftext|>