<commit_before>/*
 * Copyright 2013 Brett Vickers. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *    1. Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY COPYRIGHT HOLDER ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL COPYRIGHT HOLDER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package etree

import (
	"errors"
	"strconv"
	"strings"
)

var (
	errPath = errors.New("etree: invalid path")
)

// A segment is a portion of a path between "/" characters.
// If contains one selector and zero or more [filters].
type segment struct {
	sel     selector
	filters []filter
}

// A selector selects XML elements for consideration by the
// path traversal.
type selector interface {
	apply(e *Element, p *pather)
}

// A filter pares down a list of candidate XML elements based
// on a path filter in [brackets].
type filter interface {
	apply(p *pather)
}

func (seg *segment) apply(e *Element, p *pather) {
	seg.sel.apply(e, p)
	for _, f := range seg.filters {
		f.apply(p)
	}
}

// A pather is used to traverse an element tree, collecting results
// that match a series of path selectors and filters.
type pather struct {
	stack      []node
	results    []*Element
	inResults  map[*Element]bool
	candidates []*Element
}

// A node represents an element and the remaining path segments that
// should be applied against it by the pather.
type node struct {
	e        *Element
	segments []segment
}

func newPather() *pather {
	return &pather{
		stack:      make([]node, 0),
		results:    make([]*Element, 0),
		inResults:  make(map[*Element]bool),
		candidates: make([]*Element, 0),
	}
}

func (p *pather) push(n node) {
	p.stack = append(p.stack, n)
}

func (p *pather) pop() node {
	n := p.stack[len(p.stack)-1]
	p.stack = p.stack[:len(p.stack)-1]
	return n
}

func (p *pather) empty() bool {
	return len(p.stack) == 0
}

// traverse follows the path from the element e, collecting
// and then returning all elements that match the path's selectors
// and filters.
func (p *pather) traverse(e *Element, pathstr string) []*Element {
	segments := parsePath(pathstr)
	for p.push(node{e, segments}); !p.empty(); {
		p.eval(p.pop())
	}
	return p.results
}

// eval evalutes the current path node by applying the remaining
// path's selector rules against the node's element.
func (p *pather) eval(n node) {
	p.candidates = p.candidates[0:0]
	seg, remain := n.segments[0], n.segments[1:]
	seg.apply(n.e, p)

	if len(remain) == 0 {
		for _, c := range p.candidates {
			if in := p.inResults[c]; !in {
				p.inResults[c] = true
				p.results = append(p.results, c)
			}
		}
	} else {
		for _, c := range p.candidates {
			p.push(node{c, remain})
		}
	}
}

// parsePath parses an XPath-like string describing a path
// through an element tree and returns a slice of segment
// descriptors.
func parsePath(path string) []segment {
	// Path can start with //, but not /
	if strings.HasPrefix(path, "//") {
		path = path[1:]
	} else if strings.HasPrefix(path, "/") {
		panic(errPath)
	}

	segments := make([]segment, 0)
	for _, s := range strings.Split(path, "/") {
		segments = append(segments, parseSegment(s))
	}
	return segments
}

// parseSegment parses a path segment between / characters.
func parseSegment(path string) segment {
	pieces := strings.Split(path, "[")
	seg := segment{
		sel:     parseSelector(pieces[0]),
		filters: make([]filter, 0),
	}
	for i := 1; i < len(pieces); i++ {
		fpath := pieces[i]
		if fpath[len(fpath)-1] != ']' {
			panic(errPath)
		}
		seg.filters = append(seg.filters, parseFilter(fpath[:len(fpath)-1]))
	}
	return seg
}

// parseSelector parses a selector at the start of a path segment.
func parseSelector(path string) selector {
	switch path {
	case ".":
		return new(selectSelf)
	case "..":
		return new(selectParent)
	case "*":
		return new(selectChildren)
	case "":
		return new(selectDescendants)
	default:
		return &selectChildrenTag{path}
	}
}

// parseFilter parses a path filter contained within [brackets].
func parseFilter(path string) filter {
	if len(path) == 0 {
		panic(errPath)
	}
	eqindex := strings.Index(path, "='")
	if eqindex == -1 {
		switch {
		case path[0] == '@':
			return &filterAttr{path[1:]}
		case isInteger(path):
			pos, _ := strconv.Atoi(path)
			if pos == 0 {
				pos = 1
			}
			return &filterPos{pos - 1}
		default:
			return &filterChild{path}
		}
	} else {
		rindex := nextIndex(path, "'", eqindex+2)
		if rindex != len(path)-1 {
			panic(errPath)
		}
		switch {
		case path[0] == '@':
			return &filterAttrVal{path[1:eqindex], path[eqindex+2 : rindex]}
		default:
			return &filterChildText{path[:eqindex], path[eqindex+2 : rindex]}
		}
	}
	return nil
}

// selectSelf selects the current element into the candidate list.
type selectSelf struct{}

func (s *selectSelf) apply(e *Element, p *pather) {
	p.candidates = append(p.candidates, e)
}

// selectParent selects the element's parent into the candidate list.
type selectParent struct{}

func (s *selectParent) apply(e *Element, p *pather) {
	if e.Parent != nil {
		p.candidates = append(p.candidates, e.Parent)
	}
}

// selectChildren selects the element's child elements into the
// candidate list.
type selectChildren struct{}

func (s *selectChildren) apply(e *Element, p *pather) {
	for _, c := range e.Child {
		if c, ok := c.(*Element); ok {
			p.candidates = append(p.candidates, c)
		}
	}
}

// selectDescendants selects all descendant child elements
// of the element into the candidate list.
type selectDescendants struct{}

func (s *selectDescendants) apply(e *Element, p *pather) {
	stack := elementStack{e}
	for !stack.empty() {
		e := stack.pop()
		p.candidates = append(p.candidates, e)
		for _, c := range e.Child {
			if c, ok := c.(*Element); ok {
				p.candidates = append(p.candidates, c)
			}
		}
	}
}

// selectChildrenTag selects into the candidate list all child
// elements of the element having the specified tag.
type selectChildrenTag struct {
	tag string
}

func (s *selectChildrenTag) apply(e *Element, p *pather) {
	for _, c := range e.Child {
		if c, ok := c.(*Element); ok && c.Tag == s.tag {
			p.candidates = append(p.candidates, c)
		}
	}
}

// filterPos filters the candidate list, keeping only the
// candidate at the specified index.
type filterPos struct {
	index int
}

func (f *filterPos) apply(p *pather) {
	result := make([]*Element, 0)
	if f.index < len(p.candidates) {
		result = append(result, p.candidates[f.index])
	}
	p.candidates = result
}

// filterAttr filters the candidate list for elements having
// the specified attribute.
type filterAttr struct {
	attr string
}

func (f *filterAttr) apply(p *pather) {
	result := make([]*Element, 0)
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil {
			result = append(result, c)
		}
	}
	p.candidates = result
}

// filterAttrVal filters the candidate list for elements having
// the specified attribute with the specified value.
type filterAttrVal struct {
	attr, val string
}

func (f *filterAttrVal) apply(p *pather) {
	result := make([]*Element, 0)
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil && a.Value == f.val {
			result = append(result, c)
		}
	}
	p.candidates = result
}

// filterChild filters the candidate list for elements having
// a child element with the specified tag.
type filterChild struct {
	tag string
}

func (f *filterChild) apply(p *pather) {
	result := make([]*Element, 0)
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag {
				result = append(result, c)
			}
		}
	}
	p.candidates = result
}

// filterChildText filters the candidate list for elements having
// a child element with the specified tag and text.
type filterChildText struct {
	tag, text string
}

func (f *filterChildText) apply(p *pather) {
	result := make([]*Element, 0)
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag && cc.Text() == f.text {
				result = append(result, c)
			}
		}
	}
	p.candidates = result
}
<commit_msg>minor memory optimization to path.<commit_after>/*
 * Copyright 2013 Brett Vickers. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *    1. Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY COPYRIGHT HOLDER ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL COPYRIGHT HOLDER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package etree

import (
	"errors"
	"strconv"
	"strings"
)

var (
	errPath = errors.New("etree: invalid path")
)

// A segment is a portion of a path between "/" characters.
// If contains one selector and zero or more [filters].
type segment struct {
	sel     selector
	filters []filter
}

// A selector selects XML elements for consideration by the
// path traversal.
type selector interface {
	apply(e *Element, p *pather)
}

// A filter pares down a list of candidate XML elements based
// on a path filter in [brackets].
type filter interface {
	apply(p *pather)
}

func (seg *segment) apply(e *Element, p *pather) {
	seg.sel.apply(e, p)
	for _, f := range seg.filters {
		f.apply(p)
	}
}

// A pather is used to traverse an element tree, collecting results
// that match a series of path selectors and filters.
type pather struct {
	stack         []node
	results       []*Element
	inResults     map[*Element]bool
	candidates    []*Element
	candidatesAlt []*Element
}

// A node represents an element and the remaining path segments that
// should be applied against it by the pather.
type node struct {
	e        *Element
	segments []segment
}

func newPather() *pather {
	return &pather{
		stack:         make([]node, 0),
		results:       make([]*Element, 0),
		inResults:     make(map[*Element]bool),
		candidates:    make([]*Element, 0),
		candidatesAlt: make([]*Element, 0),
	}
}

func (p *pather) push(n node) {
	p.stack = append(p.stack, n)
}

func (p *pather) pop() node {
	n := p.stack[len(p.stack)-1]
	p.stack = p.stack[:len(p.stack)-1]
	return n
}

func (p *pather) empty() bool {
	return len(p.stack) == 0
}

// traverse follows the path from the element e, collecting
// and then returning all elements that match the path's selectors
// and filters.
func (p *pather) traverse(e *Element, pathstr string) []*Element {
	segments := parsePath(pathstr)
	for p.push(node{e, segments}); !p.empty(); {
		p.eval(p.pop())
	}
	return p.results
}

// eval evalutes the current path node by applying the remaining
// path's selector rules against the node's element.
func (p *pather) eval(n node) {
	p.candidates = p.candidates[0:0]
	seg, remain := n.segments[0], n.segments[1:]
	seg.apply(n.e, p)

	if len(remain) == 0 {
		for _, c := range p.candidates {
			if in := p.inResults[c]; !in {
				p.inResults[c] = true
				p.results = append(p.results, c)
			}
		}
	} else {
		for _, c := range p.candidates {
			p.push(node{c, remain})
		}
	}
}

// parsePath parses an XPath-like string describing a path
// through an element tree and returns a slice of segment
// descriptors.
func parsePath(path string) []segment {
	// Path can start with //, but not /
	if strings.HasPrefix(path, "//") {
		path = path[1:]
	} else if strings.HasPrefix(path, "/") {
		panic(errPath)
	}

	segments := make([]segment, 0)
	for _, s := range strings.Split(path, "/") {
		segments = append(segments, parseSegment(s))
	}
	return segments
}

// parseSegment parses a path segment between / characters.
func parseSegment(path string) segment {
	pieces := strings.Split(path, "[")
	seg := segment{
		sel:     parseSelector(pieces[0]),
		filters: make([]filter, 0),
	}
	for i := 1; i < len(pieces); i++ {
		fpath := pieces[i]
		if fpath[len(fpath)-1] != ']' {
			panic(errPath)
		}
		seg.filters = append(seg.filters, parseFilter(fpath[:len(fpath)-1]))
	}
	return seg
}

// parseSelector parses a selector at the start of a path segment.
func parseSelector(path string) selector {
	switch path {
	case ".":
		return new(selectSelf)
	case "..":
		return new(selectParent)
	case "*":
		return new(selectChildren)
	case "":
		return new(selectDescendants)
	default:
		return &selectChildrenTag{path}
	}
}

// parseFilter parses a path filter contained within [brackets].
func parseFilter(path string) filter {
	if len(path) == 0 {
		panic(errPath)
	}
	eqindex := strings.Index(path, "='")
	if eqindex == -1 {
		switch {
		case path[0] == '@':
			return &filterAttr{path[1:]}
		case isInteger(path):
			pos, _ := strconv.Atoi(path)
			if pos == 0 {
				pos = 1
			}
			return &filterPos{pos - 1}
		default:
			return &filterChild{path}
		}
	} else {
		rindex := nextIndex(path, "'", eqindex+2)
		if rindex != len(path)-1 {
			panic(errPath)
		}
		switch {
		case path[0] == '@':
			return &filterAttrVal{path[1:eqindex], path[eqindex+2 : rindex]}
		default:
			return &filterChildText{path[:eqindex], path[eqindex+2 : rindex]}
		}
	}
	return nil
}

// selectSelf selects the current element into the candidate list.
type selectSelf struct{}

func (s *selectSelf) apply(e *Element, p *pather) {
	p.candidates = append(p.candidates, e)
}

// selectParent selects the element's parent into the candidate list.
type selectParent struct{}

func (s *selectParent) apply(e *Element, p *pather) {
	if e.Parent != nil {
		p.candidates = append(p.candidates, e.Parent)
	}
}

// selectChildren selects the element's child elements into the
// candidate list.
type selectChildren struct{}

func (s *selectChildren) apply(e *Element, p *pather) {
	for _, c := range e.Child {
		if c, ok := c.(*Element); ok {
			p.candidates = append(p.candidates, c)
		}
	}
}

// selectDescendants selects all descendant child elements
// of the element into the candidate list.
type selectDescendants struct{}

func (s *selectDescendants) apply(e *Element, p *pather) {
	stack := elementStack{e}
	for !stack.empty() {
		e := stack.pop()
		p.candidates = append(p.candidates, e)
		for _, c := range e.Child {
			if c, ok := c.(*Element); ok {
				p.candidates = append(p.candidates, c)
			}
		}
	}
}

// selectChildrenTag selects into the candidate list all child
// elements of the element having the specified tag.
type selectChildrenTag struct {
	tag string
}

func (s *selectChildrenTag) apply(e *Element, p *pather) {
	for _, c := range e.Child {
		if c, ok := c.(*Element); ok && c.Tag == s.tag {
			p.candidates = append(p.candidates, c)
		}
	}
}

// filterPos filters the candidate list, keeping only the
// candidate at the specified index.
type filterPos struct {
	index int
}

func (f *filterPos) apply(p *pather) {
	p.candidatesAlt = p.candidatesAlt[:0]
	if f.index < len(p.candidates) {
		p.candidatesAlt = append(p.candidatesAlt, p.candidates[f.index])
	}
	p.candidates = p.candidatesAlt
}

// filterAttr filters the candidate list for elements having
// the specified attribute.
type filterAttr struct {
	attr string
}

func (f *filterAttr) apply(p *pather) {
	p.candidatesAlt = p.candidatesAlt[:0]
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil {
			p.candidatesAlt = append(p.candidatesAlt, c)
		}
	}
	p.candidates = p.candidatesAlt
}

// filterAttrVal filters the candidate list for elements having
// the specified attribute with the specified value.
type filterAttrVal struct {
	attr, val string
}

func (f *filterAttrVal) apply(p *pather) {
	p.candidatesAlt = p.candidatesAlt[:0]
	for _, c := range p.candidates {
		if a := c.SelectAttr(f.attr); a != nil && a.Value == f.val {
			p.candidatesAlt = append(p.candidatesAlt, c)
		}
	}
	p.candidates = p.candidatesAlt
}

// filterChild filters the candidate list for elements having
// a child element with the specified tag.
type filterChild struct {
	tag string
}

func (f *filterChild) apply(p *pather) {
	p.candidatesAlt = p.candidatesAlt[:0]
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag {
				p.candidatesAlt = append(p.candidatesAlt, c)
			}
		}
	}
	p.candidates = p.candidatesAlt
}

// filterChildText filters the candidate list for elements having
// a child element with the specified tag and text.
type filterChildText struct {
	tag, text string
}

func (f *filterChildText) apply(p *pather) {
	p.candidatesAlt = p.candidatesAlt[:0]
	for _, c := range p.candidates {
		for _, cc := range c.Child {
			if cc, ok := cc.(*Element); ok && cc.Tag == f.tag && cc.Text() == f.text {
				p.candidatesAlt = append(p.candidatesAlt, c)
			}
		}
	}
	p.candidates = p.candidatesAlt
}
<|endoftext|>


<commit_before>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t * p, struct pcap_pkthdr **pkt_header,
		   u_char ** pkt_data)
{
	return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

// Pcap wraps a pcap_t struct.
type Pcap struct {
	cptr     *C.pcap_t
	sampling int // 1:N; a bit of a hack
}

type pcapError struct{ string }

// Stat records statistics about packets received and dropped.
type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

// Interface describes a single network interface on a host.
type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

// IFAddress is the interface address.
type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

// Version returns the current pcap library version.
func Version() string { return C.GoString(C.pcap_lib_version()) }

// Datalink TODO
func (p *Pcap) Datalink() int { return int(C.pcap_datalink(p.cptr)) }

func (e *pcapError) Error() string { return e.string }

// Geterror converts the last pcap error to a Go error.
func (p *Pcap) Geterror() error { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }

// Next wraps libpcap NextEx
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

// Create TODO
func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// SetBufferSize sets buffer size (in bytes) on the activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetPromisc sets promiscuous mode on the handle. It should be called before
// activation.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetSnapLen TODO
func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetReadTimeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeoutMS int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeoutMS), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// Close calls pcap_close on the underlying pcap_t.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

// SetSampling sets the sample rate of the handle. We perform sampling in the
// Go wrapper library, as efficiently as we can at that layer.
func (p *Pcap) SetSampling(rate float64) {
	// Take 1 packet every N (approximate).
	//  rate=0.50 sampling=1.0/0.50=2
	//  rate=0.33 sampling=1.0/0.33=3
	//  rate=0.15 sampling=1.0/0.15=6.66=6 (alas)
	p.sampling = int(1.0 / rate)
}

// NextEx gets the next packet on the handle.
func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	var result int32
	for i := 0; i < p.sampling; i++ {
		// "The struct pcap_pkthdr and the packet data are not to be freed by
		// the caller, and are not guaranteed to be valid after the next call
		// ... if the code needs them to remain valid, it must make a copy of
		// them." --pcap_next_ex(3)
		//
		// I believe this is not a terrible way to do sampling. We still incur
		// the cost of pcap_next_ex, which does copy from kernel to userspace,
		// but we skip any further allocations and a lot of analysis in our
		// client code.
		result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	}

	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   make([]byte, pkthdr.caplen),
	}
	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return pkt, result
}

// Getstats TODO
func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

// SetFilter TODO
func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

// SetDatalink TODO
func (p *Pcap) SetDatalink(dlt int) error {
	if C.pcap_set_datalink(p.cptr, C.int(dlt)) == -1 {
		return p.Geterror()
	}
	return nil
}

// DatalinkValueToName string
func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

// DatalinkValueToDescription TODO
func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

// FindAllDevs TODO
func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<commit_msg>Remove unneeded wrapper<commit_after>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

// Pcap wraps a pcap_t struct.
type Pcap struct {
	cptr     *C.pcap_t
	sampling int // 1:N; a bit of a hack
}

type pcapError struct{ string }

// Stat records statistics about packets received and dropped.
type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

// Interface describes a single network interface on a host.
type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

// IFAddress is the interface address.
type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

// Version returns the current pcap library version.
func Version() string { return C.GoString(C.pcap_lib_version()) }

// Datalink TODO
func (p *Pcap) Datalink() int { return int(C.pcap_datalink(p.cptr)) }

func (e *pcapError) Error() string { return e.string }

// Geterror converts the last pcap error to a Go error.
func (p *Pcap) Geterror() error { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }

// Next wraps libpcap NextEx
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

// Create TODO
func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// SetBufferSize sets buffer size (in bytes) on the activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetPromisc sets promiscuous mode on the handle. It should be called before
// activation.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetSnapLen TODO
func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetReadTimeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeoutMS int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeoutMS), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// Close calls pcap_close on the underlying pcap_t.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

// SetSampling sets the sample rate of the handle. We perform sampling in the
// Go wrapper library, as efficiently as we can at that layer.
func (p *Pcap) SetSampling(rate float64) {
	// Take 1 packet every N (approximate).
	//  rate=0.50 sampling=1.0/0.50=2
	//  rate=0.33 sampling=1.0/0.33=3
	//  rate=0.15 sampling=1.0/0.15=6.66=6 (alas)
	p.sampling = int(1.0 / rate)
}

// NextEx gets the next packet on the handle.
func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	var result int32
	for i := 0; i < p.sampling; i++ {
		// "The struct pcap_pkthdr and the packet data are not to be freed by
		// the caller, and are not guaranteed to be valid after the next call
		// ... if the code needs them to remain valid, it must make a copy of
		// them." --pcap_next_ex(3)
		//
		// I believe this is not a terrible way to do sampling. We still incur
		// the cost of pcap_next_ex, which does copy from kernel to userspace,
		// but we skip any further allocations and a lot of analysis in our
		// client code.
		result = int32(C.pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	}

	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   make([]byte, pkthdr.caplen),
	}
	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return pkt, result
}

// Getstats TODO
func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

// SetFilter TODO
func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

// SetDatalink TODO
func (p *Pcap) SetDatalink(dlt int) error {
	if C.pcap_set_datalink(p.cptr, C.int(dlt)) == -1 {
		return p.Geterror()
	}
	return nil
}

// DatalinkValueToName string
func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

// DatalinkValueToDescription TODO
func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

// FindAllDevs TODO
func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<|endoftext|>


<commit_before>// Interface to both live and offline pcap parsing.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t *p, struct pcap_pkthdr **pkt_header,
                      u_char **pkt_data) {
    return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"
import (
	"errors"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func (p *Pcap) Next() (pkt *Packet) {
	rv, _ := p.NextEx()
	return rv
}

// Openlive opens a device and returns a *Pcap handler
func Openlive(device string, snaplen int32, promisc bool, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func Openoffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr *C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	if nil == buf {
		return
	}

	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000)
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = C.GoBytes(buf, C.int(pkthdr.caplen))
	return
}

func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) Geterror() error {
	return errors.New(C.GoString(C.pcap_geterr(p.cptr)))
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

func (p *Pcap) Setfilter(expr string) (err error) {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if -1 == C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) {
		return p.Geterror()
	}

	if -1 == C.pcap_setfilter(p.cptr, &bpf) {
		C.pcap_freecode(&bpf)
		return p.Geterror()
	}
	C.pcap_freecode(&bpf)
	return nil
}

func Version() string {
	return C.GoString(C.pcap_lib_version())
}

func (p *Pcap) Datalink() int {
	return int(C.pcap_datalink(p.cptr))
}

func (p *Pcap) Setdatalink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func Findalldevs() (ifs []Interface, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findalladdresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findalladdresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		var a IFAddress
		var err error
		if a.IP, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddr_to_IP(rsa *syscall.RawSockaddr) (IP []byte, err error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		IP = make([]byte, 4)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		IP = make([]byte, 16)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	}
	err = errors.New("Unsupported address type")
	return
}

func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if -1 == C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) {
		err = p.Geterror()
	}
	C.free(unsafe.Pointer(buf))
	return
}
<commit_msg>Added support for Windows.<commit_after>// Interface to both live and offline pcap parsing.
package pcap

/*
#cgo linux LDFLAGS: -lpcap
#cgo freebsd LDFLAGS: -lpcap
#cgo darwin LDFLAGS: -lpcap
#cgo windows CFLAGS: -I C:/WpdPack/Include
#cgo windows,386 LDFLAGS: -L C:/WpdPack/Lib -lwpcap
#cgo windows,amd64 LDFLAGS: -L C:/WpdPack/Lib/x64 -lwpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t *p, struct pcap_pkthdr **pkt_header,
                      u_char **pkt_data) {
    return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"
import (
	"errors"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func (p *Pcap) Next() (pkt *Packet) {
	rv, _ := p.NextEx()
	return rv
}

// Openlive opens a device and returns a *Pcap handler
func Openlive(device string, snaplen int32, promisc bool, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func Openoffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = errors.New(C.GoString(buf))
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr *C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	if nil == buf {
		return
	}

	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)*1000)
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = C.GoBytes(buf, C.int(pkthdr.caplen))
	return
}

func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) Geterror() error {
	return errors.New(C.GoString(C.pcap_geterr(p.cptr)))
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

func (p *Pcap) Setfilter(expr string) (err error) {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if -1 == C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) {
		return p.Geterror()
	}

	if -1 == C.pcap_setfilter(p.cptr, &bpf) {
		C.pcap_freecode(&bpf)
		return p.Geterror()
	}
	C.pcap_freecode(&bpf)
	return nil
}

func Version() string {
	return C.GoString(C.pcap_lib_version())
}

func (p *Pcap) Datalink() int {
	return int(C.pcap_datalink(p.cptr))
}

func (p *Pcap) Setdatalink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func Findalldevs() (ifs []Interface, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findalladdresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findalladdresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		var a IFAddress
		var err error
		if a.IP, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddr_to_IP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddr_to_IP(rsa *syscall.RawSockaddr) (IP []byte, err error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		IP = make([]byte, 4)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		IP = make([]byte, 16)
		for i := 0; i < len(IP); i++ {
			IP[i] = pp.Addr[i]
		}
		return
	}
	err = errors.New("Unsupported address type")
	return
}

func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if -1 == C.pcap_sendpacket(p.cptr, (*C.u_char)(unsafe.Pointer(buf)), (C.int)(len(data))) {
		err = p.Geterror()
	}
	C.free(unsafe.Pointer(buf))
	return
}
<|endoftext|>


<commit_before>package eth

import (
	"github.com/ethereum/ethchain-go"
	"github.com/ethereum/ethutil-go"
	"github.com/ethereum/ethwire-go"
	"log"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

const (
	// The size of the output buffer for writing messages
	outputBufferSize = 50
)

type Peer struct {
	// Ethereum interface
	ethereum *Ethereum
	// Net connection
	conn net.Conn
	// Output queue which is used to communicate and handle messages
	outputQueue chan *ethwire.Msg
	// Quit channel
	quit chan bool
	// Determines whether it's an inbound or outbound peer
	inbound bool
	// Flag for checking the peer's connectivity state
	connected  int32
	disconnect int32
	// Last known message send
	lastSend time.Time
	// Indicated whether a verack has been send or not
	// This flag is used by writeMessage to check if messages are allowed
	// to be send or not. If no version is known all messages are ignored.
	versionKnown bool

	// Last received pong message
	lastPong int64
	// Indicates whether a MsgGetPeersTy was requested of the peer
	// this to prevent receiving false peers.
	requestedPeerList bool

	// Determines whether this is a seed peer
	seed bool
}

func NewPeer(conn net.Conn, ethereum *Ethereum, inbound bool) *Peer {
	return &Peer{
		outputQueue: make(chan *ethwire.Msg, outputBufferSize),
		quit:        make(chan bool),
		ethereum:    ethereum,
		conn:        conn,
		inbound:     inbound,
		disconnect:  0,
		connected:   1,
	}
}

func NewOutboundPeer(addr string, ethereum *Ethereum) *Peer {
	p := &Peer{
		outputQueue: make(chan *ethwire.Msg, outputBufferSize),
		quit:        make(chan bool),
		ethereum:    ethereum,
		inbound:     false,
		connected:   0,
		disconnect:  0,
	}

	// Set up the connection in another goroutine so we don't block the main thread
	go func() {
		conn, err := net.DialTimeout("tcp", addr, 30*time.Second)

		if err != nil {
			log.Println("Connection to peer failed", err)
			p.Stop()
			return
		}
		p.conn = conn

		// Atomically set the connection state
		atomic.StoreInt32(&p.connected, 1)
		atomic.StoreInt32(&p.disconnect, 0)

		p.Start(false)
	}()

	return p
}

// Outputs any RLP encoded data to the peer
func (p *Peer) QueueMessage(msg *ethwire.Msg) {
	p.outputQueue <- msg
}

func (p *Peer) writeMessage(msg *ethwire.Msg) {
	// Ignore the write if we're not connected
	if atomic.LoadInt32(&p.connected) != 1 {
		return
	}

	if !p.versionKnown {
		switch msg.Type {
		case ethwire.MsgHandshakeTy: // Ok
		default: // Anything but ack is allowed
			return
		}
	}

	err := ethwire.WriteMessage(p.conn, msg)
	if err != nil {
		log.Println("Can't send message:", err)
		// Stop the client if there was an error writing to it
		p.Stop()
		return
	}

	// XXX TMP CODE FOR TESTNET
	switch msg.Type {
	case ethwire.MsgPeersTy:
		if p.seed {
			p.Stop()
		}
	}
}

// Outbound message handler. Outbound messages are handled here
func (p *Peer) HandleOutbound() {
	// The ping timer. Makes sure that every 2 minutes a ping is send to the peer
	tickleTimer := time.NewTicker(2 * time.Minute)
out:
	for {
		select {
		// Main message queue. All outbound messages are processed through here
		case msg := <-p.outputQueue:
			p.writeMessage(msg)

			p.lastSend = time.Now()

		case <-tickleTimer.C:
			p.writeMessage(ethwire.NewMessage(ethwire.MsgPingTy, ""))

			// Break out of the for loop if a quit message is posted
		case <-p.quit:
			break out
		}
	}

clean:
	// This loop is for draining the output queue and anybody waiting for us
	for {
		select {
		case <-p.outputQueue:
			// TODO
		default:
			break clean
		}
	}
}

// Inbound handler. Inbound messages are received here and passed to the appropriate methods
func (p *Peer) HandleInbound() {

out:
	for atomic.LoadInt32(&p.disconnect) == 0 {
		// Wait for a message from the peer
		msgs, err := ethwire.ReadMessages(p.conn)
		for _, msg := range msgs {
			if err != nil {
				log.Println(err)

				break out
			}

			switch msg.Type {
			case ethwire.MsgHandshakeTy:
				// Version message
				p.handleHandshake(msg)

				p.QueueMessage(ethwire.NewMessage(ethwire.MsgGetPeersTy, ""))
			case ethwire.MsgDiscTy:
				p.Stop()
			case ethwire.MsgPingTy:
				// Respond back with pong
				p.QueueMessage(ethwire.NewMessage(ethwire.MsgPongTy, ""))
			case ethwire.MsgPongTy:
				// If we received a pong back from a peer we set the
				// last pong so the peer handler knows this peer is still
				// active.
				p.lastPong = time.Now().Unix()
			case ethwire.MsgBlockTy:
				// Get all blocks and process them
				msg.Data = msg.Data
				for i := msg.Data.Length() - 1; i >= 0; i-- {
					block := ethchain.NewBlockFromRlpValue(msg.Data.Get(i))
					err := p.ethereum.BlockManager.ProcessBlock(block)

					if err != nil {
						log.Println(err)
					}
				}
			case ethwire.MsgTxTy:
				// If the message was a transaction queue the transaction
				// in the TxPool where it will undergo validation and
				// processing when a new block is found
				for i := 0; i < msg.Data.Length(); i++ {
					p.ethereum.TxPool.QueueTransaction(ethchain.NewTransactionFromRlpValue(msg.Data.Get(i)))
				}
			case ethwire.MsgGetPeersTy:
				// Flag this peer as a 'requested of new peers' this to
				// prevent malicious peers being forced.
				p.requestedPeerList = true
				// Peer asked for list of connected peers
				p.pushPeers()
			case ethwire.MsgPeersTy:
				// Received a list of peers (probably because MsgGetPeersTy was send)
				// Only act on message if we actually requested for a peers list
				if p.requestedPeerList {
					data := msg.Data
					// Create new list of possible peers for the ethereum to process
					peers := make([]string, data.Length())
					// Parse each possible peer
					for i := 0; i < data.Length(); i++ {
						peers[i] = data.Get(i).Get(0).AsString() + ":" + strconv.Itoa(int(data.Get(i).Get(1).AsUint()))
					}

					// Connect to the list of peers
					p.ethereum.ProcessPeerList(peers)
					// Mark unrequested again
					p.requestedPeerList = false

				}
			case ethwire.MsgGetChainTy:
				var parent *ethchain.Block
				// Length minus one since the very last element in the array is a count
				l := msg.Data.Length() - 1
				// Ignore empty get chains
				if l <= 1 {
					break
				}

				// Amount of parents in the canonical chain
				amountOfBlocks := msg.Data.Get(l).AsUint()
				// Check each SHA block hash from the message and determine whether
				// the SHA is in the database
				for i := 0; i < l; i++ {
					if data := msg.Data.Get(i).AsBytes(); p.ethereum.BlockManager.BlockChain().HasBlock(data) {
						parent = p.ethereum.BlockManager.BlockChain().GetBlock(data)
						break
					}
				}

				// If a parent is found send back a reply
				if parent != nil {
					chain := p.ethereum.BlockManager.BlockChain().GetChainFromHash(parent.Hash(), amountOfBlocks)
					p.QueueMessage(ethwire.NewMessage(ethwire.MsgBlockTy, append(chain, amountOfBlocks)))
				} else {
					// If no blocks are found we send back a reply with msg not in chain
					// and the last hash from get chain
					lastHash := msg.Data.Get(l - 1)
					log.Printf("Sending not in chain with hash %x\n", lastHash.AsRaw())
					p.QueueMessage(ethwire.NewMessage(ethwire.MsgNotInChainTy, []interface{}{lastHash.AsRaw()}))
				}
			case ethwire.MsgNotInChainTy:
				log.Printf("Not in chain %x\n", msg.Data)
				// TODO

				// Unofficial but fun nonetheless
			case ethwire.MsgTalkTy:
				log.Printf("%v says: %s\n", p.conn.RemoteAddr(), msg.Data.AsString())
			}
		}
	}

	p.Stop()
}

func (p *Peer) Start(seed bool) {
	p.seed = seed

	peerHost, _, _ := net.SplitHostPort(p.conn.LocalAddr().String())
	servHost, _, _ := net.SplitHostPort(p.conn.RemoteAddr().String())
	if peerHost == servHost {
		log.Println("Connected to self")

		p.Stop()

		return
	}

	//if !p.inbound {
	err := p.pushHandshake()
	if err != nil {
		log.Printf("Peer can't send outbound version ack", err)

		p.Stop()

		return
	}
	//}

	// Run the outbound handler in a new goroutine
	go p.HandleOutbound()
	// Run the inbound handler in a new goroutine
	go p.HandleInbound()

}

func (p *Peer) Stop() {
	if atomic.AddInt32(&p.disconnect, 1) != 1 {
		return
	}

	close(p.quit)
	if atomic.LoadInt32(&p.connected) != 0 {
		p.writeMessage(ethwire.NewMessage(ethwire.MsgDiscTy, ""))
		p.conn.Close()
	}

	log.Println("Peer shutdown")
}

func (p *Peer) pushHandshake() error {
	msg := ethwire.NewMessage(ethwire.MsgHandshakeTy, []interface{}{
		uint32(0), uint32(0), "/Ethereum(G) v0.0.1/",
	})

	p.QueueMessage(msg)

	return nil
}

// Pushes the list of outbound peers to the client when requested
func (p *Peer) pushPeers() {
	outPeers := make([]interface{}, len(p.ethereum.InOutPeers()))
	// Serialise each peer
	for i, peer := range p.ethereum.InOutPeers() {
		outPeers[i] = peer.RlpData()
	}

	// Send message to the peer with the known list of connected clients
	msg := ethwire.NewMessage(ethwire.MsgPeersTy, outPeers)

	p.QueueMessage(msg)
}

func (p *Peer) handleHandshake(msg *ethwire.Msg) {
	c := msg.Data
	// [PROTOCOL_VERSION, NETWORK_ID, CLIENT_ID]
	p.versionKnown = true

	var istr string
	// If this is an inbound connection send an ack back
	if p.inbound {
		/*
			err := p.pushHandshake()
			if err != nil {
				log.Println("Peer can't send ack back")

				p.Stop()
			}
		*/
		istr = "inbound"
	} else {
		//msg := ethwire.NewMessage(ethwire.MsgGetChainTy, []interface{}{p.ethereum.BlockManager.BlockChain().CurrentBlock.Hash(), uint64(100)})
		//p.QueueMessage(msg)

		istr = "outbound"
	}

	log.Printf("peer connect (%s) %v %s\n", istr, p.conn.RemoteAddr(), c.Get(2).AsString())
}

func (p *Peer) RlpData() []interface{} {
	host, prt, err := net.SplitHostPort(p.conn.RemoteAddr().String())
	if err != nil {
		return nil
	}

	port, err := strconv.Atoi(prt)
	if err != nil {
		return nil
	}

	//port := ethutil.NumberToBytes(uint16(i), 16)

	return []interface{}{host, port}
}

func (p *Peer) RlpEncode() []byte {
	host, prt, err := net.SplitHostPort(p.conn.RemoteAddr().String())
	if err != nil {
		return nil
	}

	i, err := strconv.Atoi(prt)
	if err != nil {
		return nil
	}

	port := ethutil.NumberToBytes(uint16(i), 16)

	return ethutil.Encode([]interface{}{host, port})
}
<commit_msg>For the testnet always 30303 for now to make it easy<commit_after>package eth

import (
	"github.com/ethereum/ethchain-go"
	"github.com/ethereum/ethutil-go"
	"github.com/ethereum/ethwire-go"
	"log"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

const (
	// The size of the output buffer for writing messages
	outputBufferSize = 50
)

type Peer struct {
	// Ethereum interface
	ethereum *Ethereum
	// Net connection
	conn net.Conn
	// Output queue which is used to communicate and handle messages
	outputQueue chan *ethwire.Msg
	// Quit channel
	quit chan bool
	// Determines whether it's an inbound or outbound peer
	inbound bool
	// Flag for checking the peer's connectivity state
	connected  int32
	disconnect int32
	// Last known message send
	lastSend time.Time
	// Indicated whether a verack has been send or not
	// This flag is used by writeMessage to check if messages are allowed
	// to be send or not. If no version is known all messages are ignored.
	versionKnown bool

	// Last received pong message
	lastPong int64
	// Indicates whether a MsgGetPeersTy was requested of the peer
	// this to prevent receiving false peers.
	requestedPeerList bool

	// Determines whether this is a seed peer
	seed bool
}

func NewPeer(conn net.Conn, ethereum *Ethereum, inbound bool) *Peer {
	return &Peer{
		outputQueue: make(chan *ethwire.Msg, outputBufferSize),
		quit:        make(chan bool),
		ethereum:    ethereum,
		conn:        conn,
		inbound:     inbound,
		disconnect:  0,
		connected:   1,
	}
}

func NewOutboundPeer(addr string, ethereum *Ethereum) *Peer {
	p := &Peer{
		outputQueue: make(chan *ethwire.Msg, outputBufferSize),
		quit:        make(chan bool),
		ethereum:    ethereum,
		inbound:     false,
		connected:   0,
		disconnect:  0,
	}

	// Set up the connection in another goroutine so we don't block the main thread
	go func() {
		conn, err := net.DialTimeout("tcp", addr, 30*time.Second)

		if err != nil {
			log.Println("Connection to peer failed", err)
			p.Stop()
			return
		}
		p.conn = conn

		// Atomically set the connection state
		atomic.StoreInt32(&p.connected, 1)
		atomic.StoreInt32(&p.disconnect, 0)

		p.Start(false)
	}()

	return p
}

// Outputs any RLP encoded data to the peer
func (p *Peer) QueueMessage(msg *ethwire.Msg) {
	p.outputQueue <- msg
}

func (p *Peer) writeMessage(msg *ethwire.Msg) {
	// Ignore the write if we're not connected
	if atomic.LoadInt32(&p.connected) != 1 {
		return
	}

	if !p.versionKnown {
		switch msg.Type {
		case ethwire.MsgHandshakeTy: // Ok
		default: // Anything but ack is allowed
			return
		}
	}

	err := ethwire.WriteMessage(p.conn, msg)
	if err != nil {
		log.Println("Can't send message:", err)
		// Stop the client if there was an error writing to it
		p.Stop()
		return
	}

	// XXX TMP CODE FOR TESTNET
	switch msg.Type {
	case ethwire.MsgPeersTy:
		if p.seed {
			p.Stop()
		}
	}
}

// Outbound message handler. Outbound messages are handled here
func (p *Peer) HandleOutbound() {
	// The ping timer. Makes sure that every 2 minutes a ping is send to the peer
	tickleTimer := time.NewTicker(2 * time.Minute)
out:
	for {
		select {
		// Main message queue. All outbound messages are processed through here
		case msg := <-p.outputQueue:
			p.writeMessage(msg)

			p.lastSend = time.Now()

		case <-tickleTimer.C:
			p.writeMessage(ethwire.NewMessage(ethwire.MsgPingTy, ""))

			// Break out of the for loop if a quit message is posted
		case <-p.quit:
			break out
		}
	}

clean:
	// This loop is for draining the output queue and anybody waiting for us
	for {
		select {
		case <-p.outputQueue:
			// TODO
		default:
			break clean
		}
	}
}

// Inbound handler. Inbound messages are received here and passed to the appropriate methods
func (p *Peer) HandleInbound() {

out:
	for atomic.LoadInt32(&p.disconnect) == 0 {
		// Wait for a message from the peer
		msgs, err := ethwire.ReadMessages(p.conn)
		for _, msg := range msgs {
			if err != nil {
				log.Println(err)

				break out
			}

			switch msg.Type {
			case ethwire.MsgHandshakeTy:
				// Version message
				p.handleHandshake(msg)

				p.QueueMessage(ethwire.NewMessage(ethwire.MsgGetPeersTy, ""))
			case ethwire.MsgDiscTy:
				p.Stop()
			case ethwire.MsgPingTy:
				// Respond back with pong
				p.QueueMessage(ethwire.NewMessage(ethwire.MsgPongTy, ""))
			case ethwire.MsgPongTy:
				// If we received a pong back from a peer we set the
				// last pong so the peer handler knows this peer is still
				// active.
				p.lastPong = time.Now().Unix()
			case ethwire.MsgBlockTy:
				// Get all blocks and process them
				msg.Data = msg.Data
				for i := msg.Data.Length() - 1; i >= 0; i-- {
					block := ethchain.NewBlockFromRlpValue(msg.Data.Get(i))
					err := p.ethereum.BlockManager.ProcessBlock(block)

					if err != nil {
						log.Println(err)
					}
				}
			case ethwire.MsgTxTy:
				// If the message was a transaction queue the transaction
				// in the TxPool where it will undergo validation and
				// processing when a new block is found
				for i := 0; i < msg.Data.Length(); i++ {
					p.ethereum.TxPool.QueueTransaction(ethchain.NewTransactionFromRlpValue(msg.Data.Get(i)))
				}
			case ethwire.MsgGetPeersTy:
				// Flag this peer as a 'requested of new peers' this to
				// prevent malicious peers being forced.
				p.requestedPeerList = true
				// Peer asked for list of connected peers
				p.pushPeers()
			case ethwire.MsgPeersTy:
				// Received a list of peers (probably because MsgGetPeersTy was send)
				// Only act on message if we actually requested for a peers list
				if p.requestedPeerList {
					data := msg.Data
					// Create new list of possible peers for the ethereum to process
					peers := make([]string, data.Length())
					// Parse each possible peer
					for i := 0; i < data.Length(); i++ {
						peers[i] = data.Get(i).Get(0).AsString() + ":" + strconv.Itoa(int(data.Get(i).Get(1).AsUint()))
					}

					// Connect to the list of peers
					p.ethereum.ProcessPeerList(peers)
					// Mark unrequested again
					p.requestedPeerList = false

				}
			case ethwire.MsgGetChainTy:
				var parent *ethchain.Block
				// Length minus one since the very last element in the array is a count
				l := msg.Data.Length() - 1
				// Ignore empty get chains
				if l <= 1 {
					break
				}

				// Amount of parents in the canonical chain
				amountOfBlocks := msg.Data.Get(l).AsUint()
				// Check each SHA block hash from the message and determine whether
				// the SHA is in the database
				for i := 0; i < l; i++ {
					if data := msg.Data.Get(i).AsBytes(); p.ethereum.BlockManager.BlockChain().HasBlock(data) {
						parent = p.ethereum.BlockManager.BlockChain().GetBlock(data)
						break
					}
				}

				// If a parent is found send back a reply
				if parent != nil {
					chain := p.ethereum.BlockManager.BlockChain().GetChainFromHash(parent.Hash(), amountOfBlocks)
					p.QueueMessage(ethwire.NewMessage(ethwire.MsgBlockTy, append(chain, amountOfBlocks)))
				} else {
					// If no blocks are found we send back a reply with msg not in chain
					// and the last hash from get chain
					lastHash := msg.Data.Get(l - 1)
					log.Printf("Sending not in chain with hash %x\n", lastHash.AsRaw())
					p.QueueMessage(ethwire.NewMessage(ethwire.MsgNotInChainTy, []interface{}{lastHash.AsRaw()}))
				}
			case ethwire.MsgNotInChainTy:
				log.Printf("Not in chain %x\n", msg.Data)
				// TODO

				// Unofficial but fun nonetheless
			case ethwire.MsgTalkTy:
				log.Printf("%v says: %s\n", p.conn.RemoteAddr(), msg.Data.AsString())
			}
		}
	}

	p.Stop()
}

func (p *Peer) Start(seed bool) {
	p.seed = seed

	peerHost, _, _ := net.SplitHostPort(p.conn.LocalAddr().String())
	servHost, _, _ := net.SplitHostPort(p.conn.RemoteAddr().String())
	if peerHost == servHost {
		log.Println("Connected to self")

		p.Stop()

		return
	}

	//if !p.inbound {
	err := p.pushHandshake()
	if err != nil {
		log.Printf("Peer can't send outbound version ack", err)

		p.Stop()

		return
	}
	//}

	// Run the outbound handler in a new goroutine
	go p.HandleOutbound()
	// Run the inbound handler in a new goroutine
	go p.HandleInbound()

}

func (p *Peer) Stop() {
	if atomic.AddInt32(&p.disconnect, 1) != 1 {
		return
	}

	close(p.quit)
	if atomic.LoadInt32(&p.connected) != 0 {
		p.writeMessage(ethwire.NewMessage(ethwire.MsgDiscTy, ""))
		p.conn.Close()
	}

	log.Println("Peer shutdown")
}

func (p *Peer) pushHandshake() error {
	msg := ethwire.NewMessage(ethwire.MsgHandshakeTy, []interface{}{
		uint32(0), uint32(0), "/Ethereum(G) v0.0.1/",
	})

	p.QueueMessage(msg)

	return nil
}

// Pushes the list of outbound peers to the client when requested
func (p *Peer) pushPeers() {
	outPeers := make([]interface{}, len(p.ethereum.InOutPeers()))
	// Serialise each peer
	for i, peer := range p.ethereum.InOutPeers() {
		outPeers[i] = peer.RlpData()
	}

	// Send message to the peer with the known list of connected clients
	msg := ethwire.NewMessage(ethwire.MsgPeersTy, outPeers)

	p.QueueMessage(msg)
}

func (p *Peer) handleHandshake(msg *ethwire.Msg) {
	c := msg.Data
	// [PROTOCOL_VERSION, NETWORK_ID, CLIENT_ID]
	p.versionKnown = true

	var istr string
	// If this is an inbound connection send an ack back
	if p.inbound {
		/*
			err := p.pushHandshake()
			if err != nil {
				log.Println("Peer can't send ack back")

				p.Stop()
			}
		*/
		istr = "inbound"
	} else {
		//msg := ethwire.NewMessage(ethwire.MsgGetChainTy, []interface{}{p.ethereum.BlockManager.BlockChain().CurrentBlock.Hash(), uint64(100)})
		//p.QueueMessage(msg)

		istr = "outbound"
	}

	log.Printf("peer connect (%s) %v %s\n", istr, p.conn.RemoteAddr(), c.Get(2).AsString())
}

func (p *Peer) RlpData() []interface{} {
	host, _, err := net.SplitHostPort(p.conn.RemoteAddr().String())
	if err != nil {
		return nil
	}

	/* FIXME
	port, err := strconv.Atoi(prt)
	if err != nil {
		return nil
	}
	*/

	//port := ethutil.NumberToBytes(uint16(i), 16)

	return []interface{}{host, uint16(30303) /*port*/}
}

func (p *Peer) RlpEncode() []byte {
	host, prt, err := net.SplitHostPort(p.conn.RemoteAddr().String())
	if err != nil {
		return nil
	}

	i, err := strconv.Atoi(prt)
	if err != nil {
		return nil
	}

	port := ethutil.NumberToBytes(uint16(i), 16)

	return ethutil.Encode([]interface{}{host, port})
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"net"
	"os"
	"regexp"
	"strings"
	"time"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	fping "github.com/tatsushid/go-fastping"
)

type PingPlugin struct {
	Hosts    []string
	Tempfile string
}

func (pp PingPlugin) FetchMetrics() (map[string]interface{}, error) {
	stat := make(map[string]interface{})

	pinger := fping.NewPinger()
	pinger.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		rttMicroSec := float64(rtt.Nanoseconds()) / 1000.0 / 1000.0
		stat[escapeHostName(addr.String())] = rttMicroSec
	}

	for _, host := range pp.Hosts {
		ra, err := net.ResolveIPAddr("ip4:icmp", host)
		if err != nil {
			return nil, err
		}

		pinger.AddIPAddr(ra)
	}

	err := pinger.Run()
	if err != nil {
		return nil, err
	}

	return stat, nil
}

func (pp PingPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
	for _, host := range pp.Hosts {
		metrics = append(metrics, mp.Metrics{
			Name:    escapeHostName(host),
			Label:   host,
			Diff:    false,
			Stacked: false,
		})
	}

	return map[string](mp.Graphs){
		"ping.rtt": mp.Graphs{
			Label:   "Ping Round Trip Times",
			Unit:    "float",
			Metrics: metrics,
		},
	}
}

func escapeHostName(host string) string {
	return strings.Replace(host, ".", "_", -1)
}

func escapeHostNames(hosts string) string {
	return strings.Replace(escapeHostName(hosts), ",", "-", -1)
}

func validate(ipAddr string) bool {
	r := regexp.MustCompile("^\\d+\\.\\d+\\.\\d+\\.\\d+$")
	return r.MatchString(ipAddr)
}

func parseHostsString(optHost string) ([]string, error) {
	hosts := strings.Split(optHost, ",")
	for _, host := range hosts {
		if !validate(host) {
			msg := fmt.Sprintf("error: %v must be ipv4 address format\n", host)
			return nil, errors.New(msg)
		}
	}

	return hosts, nil
}

func main() {
	optHost := flag.String("host", "127.0.0.1", "IP Address(es)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	hosts, err := parseHostsString(*optHost)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	var pp PingPlugin
	pp.Hosts = hosts

	helper := mp.NewMackerelPlugin(pp)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-ping-%s", escapeHostNames(*optHost))
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>Fix bug<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"net"
	"os"
	"regexp"
	"strings"
	"time"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	fping "github.com/tatsushid/go-fastping"
)

type PingPlugin struct {
	Hosts    []string
	Tempfile string
}

func (pp PingPlugin) FetchMetrics() (map[string]interface{}, error) {
	stat := make(map[string]interface{})

	var prevRtt time.Duration

	pinger := fping.NewPinger()
	pinger.OnRecv = func(addr *net.IPAddr, rtt time.Duration) {
		rttMicroSec := float64(rtt.Nanoseconds()-prevRtt.Nanoseconds()) / 1000.0 / 1000.0
		fmt.Println(rttMicroSec)
		stat[escapeHostName(addr.String())] = rttMicroSec

		prevRtt = rtt
	}

	for _, host := range pp.Hosts {
		ra, err := net.ResolveIPAddr("ip4:icmp", host)
		if err != nil {
			return nil, err
		}

		pinger.AddIPAddr(ra)
	}

	err := pinger.Run()
	if err != nil {
		return nil, err
	}

	return stat, nil
}

func (pp PingPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := []mp.Metrics{}
	for _, host := range pp.Hosts {
		metrics = append(metrics, mp.Metrics{
			Name:    escapeHostName(host),
			Label:   host,
			Diff:    false,
			Stacked: false,
		})
	}

	return map[string](mp.Graphs){
		"ping.rtt": mp.Graphs{
			Label:   "Ping Round Trip Times",
			Unit:    "float",
			Metrics: metrics,
		},
	}
}

func escapeHostName(host string) string {
	return strings.Replace(host, ".", "_", -1)
}

func escapeHostNames(hosts string) string {
	return strings.Replace(escapeHostName(hosts), ",", "-", -1)
}

func validate(ipAddr string) bool {
	r := regexp.MustCompile("^\\d+\\.\\d+\\.\\d+\\.\\d+$")
	return r.MatchString(ipAddr)
}

func parseHostsString(optHost string) ([]string, error) {
	hosts := strings.Split(optHost, ",")
	for _, host := range hosts {
		if !validate(host) {
			msg := fmt.Sprintf("error: %v must be ipv4 address format\n", host)
			return nil, errors.New(msg)
		}
	}

	return hosts, nil
}

func main() {
	optHost := flag.String("host", "127.0.0.1", "IP Address(es)")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	hosts, err := parseHostsString(*optHost)
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	var pp PingPlugin
	pp.Hosts = hosts

	helper := mp.NewMackerelPlugin(pp)

	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-ping-%s", escapeHostNames(*optHost))
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013 IBM Corp.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *    Seth Hoenig
 *    Allan Stockdill-Mander
 *    Mike Robertson
 */

package mqtt

import (
	"errors"
	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git/packets"
	"sync"
	"time"
)

type lastcontact struct {
	sync.Mutex
	lasttime time.Time
}

func (l *lastcontact) update() {
	l.Lock()
	defer l.Unlock()
	l.lasttime = time.Now()

}

func (l *lastcontact) get() time.Time {
	l.Lock()
	defer l.Unlock()
	return l.lasttime
}

func keepalive(c *Client) {
	DEBUG.Println(PNG, "keepalive starting")

	for {
		select {
		case <-c.stop:
			DEBUG.Println(PNG, "keepalive stopped")
			c.workers.Done()
			return
		default:
			last := uint(time.Since(c.lastContact.get()).Seconds())
			//DEBUG.Printf("%s last contact: %d (timeout: %d)", PNG, last, uint(c.options.KeepAlive.Seconds()))
			if last > uint(c.options.KeepAlive.Seconds()) {
				if !c.pingOutstanding {
					DEBUG.Println(PNG, "keepalive sending ping")
					ping := packets.NewControlPacket(packets.Pingreq).(*packets.PingreqPacket)
					//We don't want to wait behind large messages being sent, the Write call
					//will block until it it able to send the packet.
					ping.Write(c.conn)
					c.pingOutstanding = true
				} else {
					CRITICAL.Println(PNG, "pingresp not received, disconnecting")
					c.workers.Done()
					c.internalConnLost(errors.New("pingresp not received, disconnecting"))
					return
				}
			}
			time.Sleep(1 * time.Second)
		}
	}
}
<commit_msg>Set pingoutstanding to be false when starting the keepalive, required when reconnecting in some situations.<commit_after>/*
 * Copyright (c) 2013 IBM Corp.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *    Seth Hoenig
 *    Allan Stockdill-Mander
 *    Mike Robertson
 */

package mqtt

import (
	"errors"
	"git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git/packets"
	"sync"
	"time"
)

type lastcontact struct {
	sync.Mutex
	lasttime time.Time
}

func (l *lastcontact) update() {
	l.Lock()
	defer l.Unlock()
	l.lasttime = time.Now()

}

func (l *lastcontact) get() time.Time {
	l.Lock()
	defer l.Unlock()
	return l.lasttime
}

func keepalive(c *Client) {
	DEBUG.Println(PNG, "keepalive starting")
	c.pingOutstanding = false

	for {
		select {
		case <-c.stop:
			DEBUG.Println(PNG, "keepalive stopped")
			c.workers.Done()
			return
		default:
			last := uint(time.Since(c.lastContact.get()).Seconds())
			//DEBUG.Printf("%s last contact: %d (timeout: %d)", PNG, last, uint(c.options.KeepAlive.Seconds()))
			if last > uint(c.options.KeepAlive.Seconds()) {
				if !c.pingOutstanding {
					DEBUG.Println(PNG, "keepalive sending ping")
					ping := packets.NewControlPacket(packets.Pingreq).(*packets.PingreqPacket)
					//We don't want to wait behind large messages being sent, the Write call
					//will block until it it able to send the packet.
					ping.Write(c.conn)
					c.pingOutstanding = true
				} else {
					CRITICAL.Println(PNG, "pingresp not received, disconnecting")
					c.workers.Done()
					c.internalConnLost(errors.New("pingresp not received, disconnecting"))
					return
				}
			}
			time.Sleep(1 * time.Second)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	config "github.com/spf13/viper"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

//Runtime
func init() {
	config.SetConfigName("config")
	config.AddConfigPath(".\\")
	config.ReadInConfig()

	config.SetDefault("Binding", ":8080")
	config.SetDefault("ScriptFolder", ".\\scripts\\")
}

func main() {
	mx := mux.NewRouter()

	mx.HandleFunc("/", IndexHandler)
	mx.HandleFunc("/exit", ExitHandler)
	//mx.HandleFunc("/command/{name:\\S+}", RunShell)
	mx.HandleFunc("/script/{name:\\S+}", RunScript)

	log.Info("Listening at " + config.GetString("Binding"))
	http.ListenAndServe(config.GetString("Binding"), mx)
}

//Functions
func ParseArgs(r *http.Request) string {
	var argsBuffer bytes.Buffer

	pUrl, err := url.Parse(r.RequestURI)
	if err != nil {
		log.Error("Error Parsing URL")
	}

	for key, value := range pUrl.Query() {
		argsBuffer.WriteString(" " + key)
		argsBuffer.WriteString(" " + value[0])
	}

	return argsBuffer.String()
}

func exec_script(sc string) string {
	sc = string(strings.TrimSpace(sc))

	cmd := exec.Command("powershell.exe", "-NoLogo", "-NonInteractive", "-Command", "&{", sc, "}")
	out, err := cmd.Output()

	if err != nil {
		log.WithFields(log.Fields{
			"command": string(sc),
		}).Error(err.Error())
	}

	return string(out)
}

//Handlers
func IndexHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Service Running"))
}

func ExitHandler(w http.ResponseWriter, r *http.Request) {
	defer os.Exit(0)

	log.Info("Shutting Down")
	w.Write([]byte(fmt.Sprintf("Shutting Down")))
	time.Sleep(3000 * time.Millisecond)
}

func RunShell(w http.ResponseWriter, r *http.Request) {
	var commbuffer bytes.Buffer
	commbuffer.WriteString(mux.Vars(r)["name"])
	commbuffer.WriteString(ParseArgs(r))
	commbuffer.WriteString(" | ConvertTo-Json")

	w.Write([]byte(fmt.Sprintf(exec_script(commbuffer.String()))))
}

func RunScript(w http.ResponseWriter, r *http.Request) {
	var commbuffer bytes.Buffer

	commbuffer.WriteString("&\"")
	commbuffer.WriteString(filepath.Join(config.GetString("ScriptFolder"), mux.Vars(r)["name"]))
	commbuffer.WriteString("\"")
	commbuffer.WriteString(ParseArgs(r))

	w.Write([]byte(fmt.Sprintf(exec_script(commbuffer.String()))))
}
<commit_msg>Remove explicit file extension<commit_after>package main

import (
	"bytes"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	config "github.com/spf13/viper"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

//Runtime
func init() {
	config.SetConfigName("config")
	config.AddConfigPath(".\\")
	config.ReadInConfig()

	config.SetDefault("Binding", ":8080")
	config.SetDefault("ScriptFolder", ".\\scripts\\")
}

func main() {
	mx := mux.NewRouter()

	mx.HandleFunc("/", IndexHandler)
	mx.HandleFunc("/exit", ExitHandler)
	mx.HandleFunc("/command/{name:\\S+}", RunShell)
	mx.HandleFunc("/script/{name:\\S+}", RunScript)

	log.Info("Listening at " + config.GetString("Binding"))
	http.ListenAndServe(config.GetString("Binding"), mx)
}

//Functions
func ParseArgs(r *http.Request) string {
	var argsBuffer bytes.Buffer

	pUrl, err := url.Parse(r.RequestURI)
	if err != nil {
		log.Error("Error Parsing URL")
	}

	for key, value := range pUrl.Query() {
		argsBuffer.WriteString(" " + key)
		argsBuffer.WriteString(" " + value[0])
	}

	return argsBuffer.String()
}

func exec_script(sc string) string {
	sc = string(strings.TrimSpace(sc))

	cmd := exec.Command("powershell.exe", "-NoLogo", "-NonInteractive", "-Command", "&{", sc, "}")
	out, err := cmd.Output()

	if err != nil {
		log.WithFields(log.Fields{
			"command": string(sc),
		}).Error(err.Error())
	}

	return string(out)
}

//Handlers
func IndexHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("Service Running"))
}

func ExitHandler(w http.ResponseWriter, r *http.Request) {
	defer os.Exit(0)

	log.Info("Shutting Down")
	w.Write([]byte(fmt.Sprintf("Shutting Down")))
	time.Sleep(3000 * time.Millisecond)
}

func RunShell(w http.ResponseWriter, r *http.Request) {
	var commbuffer bytes.Buffer
	commbuffer.WriteString(mux.Vars(r)["name"])
	commbuffer.WriteString(ParseArgs(r))
	commbuffer.WriteString(" | ConvertTo-Json")

	w.Write([]byte(fmt.Sprintf(exec_script(commbuffer.String()))))
}

func RunScript(w http.ResponseWriter, r *http.Request) {
	var commbuffer bytes.Buffer

	commbuffer.WriteString("&\"")
	commbuffer.WriteString(filepath.Join(config.GetString("ScriptFolder"), mux.Vars(r)["name"]))
	commbuffer.WriteString("\"")
	commbuffer.WriteString(ParseArgs(r))

	w.Write([]byte(fmt.Sprintf(exec_script(commbuffer.String()))))
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package pipeline

import (
	"encoding/json"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/Jeffail/benthos/lib/log"
	"github.com/Jeffail/benthos/lib/message"
	"github.com/Jeffail/benthos/lib/metrics"
	"github.com/Jeffail/benthos/lib/processor"
	"github.com/Jeffail/benthos/lib/response"
	"github.com/Jeffail/benthos/lib/types"
)

func TestSanitise(t *testing.T) {
	var actObj interface{}
	var act []byte
	var err error

	exp := `{` +
		`"processors":[],` +
		`"threads":10` +
		`}`

	conf := NewConfig()
	conf.Threads = 10
	conf.Processors = nil

	if actObj, err = SanitiseConfig(conf); err != nil {
		t.Fatal(err)
	}
	if act, err = json.Marshal(actObj); err != nil {
		t.Fatal(err)
	}
	if string(act) != exp {
		t.Errorf("Wrong sanitised output: %s != %v", act, exp)
	}

	exp = `{` +
		`"processors":[` +
		`{` +
		`"type":"log",` +
		`"log":{` +
		`"level":"INFO",` +
		`"message":""` +
		`}` +
		`},` +
		`{` +
		`"type":"archive",` +
		`"archive":{` +
		`"format":"binary",` +
		`"path":"nope"` +
		`}` +
		`}` +
		`],` +
		`"threads":10` +
		`}`

	proc := processor.NewConfig()
	proc.Type = "log"
	conf.Processors = append(conf.Processors, proc)

	proc = processor.NewConfig()
	proc.Type = "archive"
	proc.Archive.Path = "nope"
	conf.Processors = append(conf.Processors, proc)

	if actObj, err = SanitiseConfig(conf); err != nil {
		t.Fatal(err)
	}
	if act, err = json.Marshal(actObj); err != nil {
		t.Fatal(err)
	}
	if string(act) != exp {
		t.Errorf("Wrong sanitised output: %s != %v", act, exp)
	}
}

func TestProcCtor(t *testing.T) {
	firstProc := processor.NewConfig()
	firstProc.Type = "bounds_check"
	firstProc.BoundsCheck.MinPartSize = 5

	secondProc := processor.NewConfig()
	secondProc.Type = "insert_part"
	secondProc.InsertPart.Content = "1"
	secondProc.InsertPart.Index = 0

	conf := NewConfig()
	conf.Processors = append(conf.Processors, firstProc)

	pipe, err := New(
		conf, nil,
		log.New(os.Stdout, log.Config{LogLevel: "NONE"}),
		metrics.DudType{},
		func() (types.Processor, error) {
			return processor.New(
				secondProc, nil,
				log.New(os.Stdout, log.Config{LogLevel: "NONE"}),
				metrics.DudType{},
			)
		},
	)
	if err != nil {
		t.Fatal(err)
	}

	tChan := make(chan types.Transaction)
	resChan := make(chan types.Response)

	if err = pipe.Consume(tChan); err != nil {
		t.Fatal(err)
	}

	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case tChan <- types.NewTransaction(
		message.New([][]byte{[]byte("foo bar baz")}), resChan,
	):
	}

	var tran types.Transaction
	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case tran = <-pipe.TransactionChan():
	}

	exp := [][]byte{
		[]byte("1"),
		[]byte("foo bar baz"),
	}
	if act := message.GetAllBytes(tran.Payload); !reflect.DeepEqual(exp, act) {
		t.Errorf("Wrong contents: %s != %s", act, exp)
	}

	go func() {
		select {
		case <-time.After(time.Second):
			t.Fatal("timed out")
		case tran.ResponseChan <- response.NewAck():
		}
	}()

	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case res := <-resChan:
		if res.Error() != nil {
			t.Error(res.Error())
		}
	}

	pipe.CloseAsync()
	if err = pipe.WaitForClose(time.Second); err != nil {
		t.Error(err)
	}
}
<commit_msg>Fix test<commit_after>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package pipeline

import (
	"encoding/json"
	"os"
	"reflect"
	"testing"
	"time"

	"github.com/Jeffail/benthos/lib/log"
	"github.com/Jeffail/benthos/lib/message"
	"github.com/Jeffail/benthos/lib/metrics"
	"github.com/Jeffail/benthos/lib/processor"
	"github.com/Jeffail/benthos/lib/response"
	"github.com/Jeffail/benthos/lib/types"
)

func TestSanitise(t *testing.T) {
	var actObj interface{}
	var act []byte
	var err error

	exp := `{` +
		`"processors":[],` +
		`"threads":10` +
		`}`

	conf := NewConfig()
	conf.Threads = 10
	conf.Processors = nil

	if actObj, err = SanitiseConfig(conf); err != nil {
		t.Fatal(err)
	}
	if act, err = json.Marshal(actObj); err != nil {
		t.Fatal(err)
	}
	if string(act) != exp {
		t.Errorf("Wrong sanitised output: %s != %v", act, exp)
	}

	exp = `{` +
		`"processors":[` +
		`{` +
		`"type":"log",` +
		`"log":{` +
		`"fields":{},` +
		`"level":"INFO",` +
		`"message":""` +
		`}` +
		`},` +
		`{` +
		`"type":"archive",` +
		`"archive":{` +
		`"format":"binary",` +
		`"path":"nope"` +
		`}` +
		`}` +
		`],` +
		`"threads":10` +
		`}`

	proc := processor.NewConfig()
	proc.Type = "log"
	conf.Processors = append(conf.Processors, proc)

	proc = processor.NewConfig()
	proc.Type = "archive"
	proc.Archive.Path = "nope"
	conf.Processors = append(conf.Processors, proc)

	if actObj, err = SanitiseConfig(conf); err != nil {
		t.Fatal(err)
	}
	if act, err = json.Marshal(actObj); err != nil {
		t.Fatal(err)
	}
	if string(act) != exp {
		t.Errorf("Wrong sanitised output: %s != %v", act, exp)
	}
}

func TestProcCtor(t *testing.T) {
	firstProc := processor.NewConfig()
	firstProc.Type = "bounds_check"
	firstProc.BoundsCheck.MinPartSize = 5

	secondProc := processor.NewConfig()
	secondProc.Type = "insert_part"
	secondProc.InsertPart.Content = "1"
	secondProc.InsertPart.Index = 0

	conf := NewConfig()
	conf.Processors = append(conf.Processors, firstProc)

	pipe, err := New(
		conf, nil,
		log.New(os.Stdout, log.Config{LogLevel: "NONE"}),
		metrics.DudType{},
		func() (types.Processor, error) {
			return processor.New(
				secondProc, nil,
				log.New(os.Stdout, log.Config{LogLevel: "NONE"}),
				metrics.DudType{},
			)
		},
	)
	if err != nil {
		t.Fatal(err)
	}

	tChan := make(chan types.Transaction)
	resChan := make(chan types.Response)

	if err = pipe.Consume(tChan); err != nil {
		t.Fatal(err)
	}

	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case tChan <- types.NewTransaction(
		message.New([][]byte{[]byte("foo bar baz")}), resChan,
	):
	}

	var tran types.Transaction
	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case tran = <-pipe.TransactionChan():
	}

	exp := [][]byte{
		[]byte("1"),
		[]byte("foo bar baz"),
	}
	if act := message.GetAllBytes(tran.Payload); !reflect.DeepEqual(exp, act) {
		t.Errorf("Wrong contents: %s != %s", act, exp)
	}

	go func() {
		select {
		case <-time.After(time.Second):
			t.Fatal("timed out")
		case tran.ResponseChan <- response.NewAck():
		}
	}()

	select {
	case <-time.After(time.Second):
		t.Fatal("timed out")
	case res := <-resChan:
		if res.Error() != nil {
			t.Error(res.Error())
		}
	}

	pipe.CloseAsync()
	if err = pipe.WaitForClose(time.Second); err != nil {
		t.Error(err)
	}
}
<|endoftext|>